import os
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import streamlit as st
import torch
import torch.nn as nn
import torch.nn.functional as F
from sentence_transformers import SentenceTransformer
from sklearn.feature_extraction.text import CountVectorizer
import pandas as pd
import numpy as np
import time


# ==========================================
# 1. PAGE CONFIG & UI SETUP
# ==========================================
st.set_page_config(page_title="Neural VAE Architecture", layout="wide")

# (Make sure these two lines only appear ONCE in your entire file!)
st.title("🧠 Variational Autoencoder (VAE): Inside the Black Box")
st.markdown("Watch the mathematics of a Neural Network transform text into topics in real-time.")

# --- DISPLAY ARCHITECTURE DIAGRAM ---
try:
    st.image("pipeline.png", use_container_width=True)
except Exception:
    st.info("💡 Tip: Upload a 'pipeline.png' image to display the architecture diagram here.")

st.divider()

# --- CACHE THE HEAVY BERT MODEL ---
@st.cache_resource
def load_bert():
    return SentenceTransformer("all-MiniLM-L6-v2")


# --- SIDEBAR CONTROLS ---
st.sidebar.header("1. Input & Text Prep")
n_topics = st.sidebar.slider("Number of Topics (Latent Dimensions)", 2, 10, 3)
words_per_theme = st.sidebar.slider("Top Words per Topic", 3, 10, 5)

st.sidebar.header("2. Neural Network Settings")
activation_choice = st.sidebar.selectbox("Encoder Activation", ["ReLU", "Softplus", "GELU"])
dropout_rate = st.sidebar.slider("Dropout Rate (Noise)", 0.0, 0.5, 0.2)

st.sidebar.header("3. Optimization (Loss)")
kl_beta = st.sidebar.slider("Max KL Beta (β)", 0.0, 1.0, 0.1)
kl_warmup = st.sidebar.slider("KL Annealing (Epochs)", 0, 500, 200)
learning_rate = st.sidebar.selectbox("Learning Rate", [0.001, 0.005, 0.01, 0.05], index=2)
epochs = st.sidebar.slider("Training Epochs", 100, 1000, 500)

sample_abstract = """
Students who registered for the Mapping with Google massive open online course (MOOC) 
were asked several questions during the registration process to identify prior 
experience with eleven skills as well as their goals for registering for the course. 
At the end of the course, we compared students' self reports of goal achievement 
with behavioral click-stream analysis. In addition, we assessed how well prior 
skill in a subject predicts a student's course completion and found no correlation.
"""
raw_data = st.sidebar.text_area("Training Data:", value=sample_abstract, height=150)

# ==========================================
# 2. THE PYTORCH VAE ARCHITECTURE
# ==========================================
class ContextualVAE(nn.Module):
    def __init__(self, input_dim, hidden_dim, latent_dim, vocab_size, activation, dropout):
        super().__init__()
        act_fn = nn.ReLU() if activation == "ReLU" else nn.Softplus() if activation == "Softplus" else nn.GELU()

        # 1. ENCODER
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim), 
            act_fn, 
            nn.Dropout(dropout), 
            nn.BatchNorm1d(hidden_dim)
        )
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
        
        # 2. DECODER
        self.decoder = nn.Linear(latent_dim, vocab_size, bias=False)

    def encode(self, x):
        h = self.encoder(x)
        return self.fc_mu(h), self.fc_logvar(h)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        p_topic = F.softmax(z, dim=-1)
        recon_log_probs = F.log_softmax(self.decoder(p_topic), dim=-1)
        return recon_log_probs, mu, logvar, z

# ==========================================
# 3. THE ARCHITECTURE DASHBOARD (UI LAYOUT)
# ==========================================
st.divider()

col1, col2, col3 = st.columns(3)

with col1:
    st.markdown("### 📥 1. The Encoder")
    st.latex(r"\mu, \log(\sigma^2) = \text{NeuralNet}(X_{bert})")
    st.caption("Compresses 384D BERT vectors into Mean and Variance coordinates.")
    ui_mu = st.empty()
    ui_var = st.empty()
    with st.expander("View Encoder History"):
        ui_encoder_table = st.empty()

with col2:
    st.markdown("### 🕳️ 2. The Latent Space")
    st.latex(r"Z = \mu + \sigma \odot \epsilon")
    st.caption("The Reparameterization Trick. Samples a topic vector adding random noise.")
    ui_z = st.empty()
    ui_softmax = st.empty()
    with st.expander("View Latent History"):
        ui_latent_table = st.empty()

with col3:
    st.markdown("### 📤 3. The Decoder")
    st.latex(r"\hat{Y} = \text{Softmax}(Z) \cdot W_{decoder}")
    st.caption("Maps the tiny topic vector back to the full Vocabulary size.")
    ui_decoder = st.empty()
    with st.expander("View Decoder History"):
        ui_decoder_table = st.empty()

st.divider()
st.markdown("### 📉 4. Optimization (Backpropagation)")
st.latex(r"\mathcal{L} = D_{KL}(Y || \hat{Y}) + \beta \cdot D_{KL}(\mathcal{N}(\mu, \sigma) || \mathcal{N}(0, I))")

col_loss_text, col_loss_graph = st.columns([1, 2])
with col_loss_text:
    ui_beta = st.empty()
    ui_recon_loss = st.empty()
    ui_kl_loss = st.empty()
with col_loss_graph:
    loss_chart_placeholder = st.empty()

# ==========================================
# 4. TRAINING EXECUTION
# ==========================================
if st.sidebar.button("🚀 Start Backpropagation", type="primary", use_container_width=True):
    
    sentences = [s.strip() for s in raw_data.split('.') if len(s.strip()) > 10]
    
    # Data Prep
    stop_w = list(CountVectorizer(stop_words='english').get_stop_words()) + ['students', 'course']
    vectorizer = CountVectorizer(stop_words=stop_w)
    bow_matrix = vectorizer.fit_transform(sentences).toarray()
    vocab = vectorizer.get_feature_names_out()
    
    target_dist = torch.tensor(bow_matrix).float()
    target_dist = target_dist / (target_dist.sum(dim=1, keepdim=True) + 1e-10)
    
    with st.spinner("Loading BERT (Cached)..."):
        embedder = load_bert() 
        bert_inputs = torch.tensor(embedder.encode(sentences))
    
    model = ContextualVAE(384, 128, n_topics, len(vocab), activation_choice, dropout_rate)
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    
    recon_hist, kl_hist = [], []
    
    # Trackers for the expanding history tables
    history_encoder = []
    history_latent = []
    history_decoder = []
    
# --- PYTORCH TRAINING LOOP ---
    for epoch in range(1, epochs + 1):
        model.train()
        optimizer.zero_grad()
        
        recon, mu, logvar, z = model(bert_inputs)
        
        recon_loss = F.kl_div(recon, target_dist, reduction='batchmean')
        kl_loss = -0.5 * torch.mean(torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1))
        
        current_beta = kl_beta * min(1.0, epoch / max(1, kl_warmup))
        total_loss = recon_loss + (current_beta * kl_loss)
        
        total_loss.backward()
        optimizer.step()
        
        # ==========================================
        # 1. RECORD DATA EVERY SINGLE EPOCH
        # ==========================================
        # Grab math for Sentence 1
        s1_mu = np.round(mu[0].detach().numpy(), 3).tolist()
        s1_var = np.round(torch.exp(0.5 * logvar[0]).detach().numpy(), 3).tolist()
        s1_z = np.round(z[0].detach().numpy(), 3).tolist()
        s1_soft = np.round(F.softmax(z[0], dim=-1).detach().numpy(), 3).tolist()
        
        # Save to Encoder & Latent History
        history_encoder.append({"Epoch": epoch, "Mean (μ)": str(s1_mu), "Variance (σ)": str(s1_var)})
        history_latent.append({"Epoch": epoch, "Sampled Vector (Z)": str(s1_z), "Topic Probs": str(s1_soft)})
        
        # Dynamically create separate columns for EACH topic in the Decoder history!
        weights = model.decoder.weight.T
        dec_dict = {"Epoch": epoch}
        dec_text = ""
        
        for i in range(n_topics):
            top_idx = torch.argsort(weights[i], descending=True)[:words_per_theme]
            theme_words = [vocab[j] for j in top_idx]
            
            # Format for the UI Box
            dec_text += f"**T{i+1}:** {', '.join(theme_words)}\n\n"
            # Format for the DataFrame (Separate column for Topic 1, Topic 2, etc.)
            dec_dict[f"Topic {i+1}"] = ", ".join(theme_words)
            
        history_decoder.append(dec_dict)
        
        recon_hist.append(recon_loss.item())
        kl_hist.append((current_beta * kl_loss).item())
        
        # ==========================================
        # 2. SAFE UI RENDER (Draws the accumulated data every 20 epochs)
        # ==========================================
        if epoch % 20 == 0 or epoch == epochs:
            
            # Update Live Values
            ui_mu.info(f"**Mean (μ) for Sentence 1:**\n`{s1_mu}`")
            ui_var.warning(f"**Variance (σ) for Sentence 1:**\n`{s1_var}`")
            ui_z.success(f"**Sampled 'Z' for Sentence 1:**\n`{s1_z}`")
            ui_softmax.error(f"**Topic Probabilities:**\n`{s1_soft}`")
            ui_decoder.info(dec_text)
            
            # Update Tables (These now contain ALL epochs up to this point!)
            ui_encoder_table.dataframe(pd.DataFrame(history_encoder), hide_index=True, use_container_width=True)
            ui_latent_table.dataframe(pd.DataFrame(history_latent), hide_index=True, use_container_width=True)
            ui_decoder_table.dataframe(pd.DataFrame(history_decoder), hide_index=True, use_container_width=True)
            
            # Update Optimization Metrics
            ui_beta.metric("Current Beta (Annealing)", f"{current_beta:.4f}")
            ui_recon_loss.metric("Reconstruction Loss", f"{recon_loss.item():.4f}")
            ui_kl_loss.metric("KL Penalty", f"{(current_beta * kl_loss).item():.4f}")
            
            # Update Graph
            chart_df = pd.DataFrame({
                "Recon Loss": recon_hist, 
                "KL Penalty": kl_hist
            })
            loss_chart_placeholder.line_chart(chart_df)
            
            time.sleep(0.05)

    st.success("Neural Network Training Complete!")