app.py

import os
import time
import streamlit as st
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import plotly.express as px

from models.supervised_ae import Autoencoder
from models.variational_ae import ConvGMVAE, gmvae_loss_function # Updated to GM-VAE
from utils import get_device, get_2d_projections
#
#  --- Page Config ---
st.set_page_config(page_title="Autoencoder Lab", layout="wide")
st.title("🧠 Convolutional Autoencoder Lab")
st.write("Experiment with the information bottleneck and see how the model reconstructs digits.")

# --- Sidebar: Architecture & Model Selection ---
with st.sidebar:
    architecture = st.selectbox(
        "Choose Architecture",
        ["Supervised Autoencoder", "Variational Autoencoder"],
    )

    if architecture == "Supervised Autoencoder":
        st.header("Model Settings")
        embedding_dim = st.slider("Embedding Dimension (The Bottleneck)", min_value=2, max_value=256, value=64)
        capacity = st.selectbox("Conv Capacity", options=[8, 16, 32], index=1)

        st.header("Training Settings")
        epochs = st.number_input("Epochs", min_value=1, max_value=20, value=5)
        lr = st.number_input("Learning Rate", value=1e-3, format="%.4f")
        batch_size = st.number_input("Batch Size", value=64)
        class_weight = st.slider("Classification Weight (0 = Pure Autoencoder)", min_value=0.0, max_value=1.0, value=0.1, step=0.05)

        train_button = st.button("🚀 Train New Model")

        st.header("📂 Model Browser")
        # Scan the outputs directory for .pth files
        os.makedirs("outputs", exist_ok=True)
        available_files = [f for f in os.listdir("outputs") if f.endswith(".pth")]

        if available_files:
            selected_file = st.selectbox("Load an existing checkpoint:", sorted(available_files))

            if st.button("🔄 Load Selected Model"):
                # Parse params from filename (handles both old and new formats)
                try:
                    parts = selected_file.replace(".pth", "").split("_")
                    loaded_cap, loaded_dim, loaded_cw = 16, 64, 0.0
                    for part in parts:
                        if part.startswith("cap"):
                            loaded_cap = int(part.replace("cap", ""))
                        elif part.startswith("dim"):
                            loaded_dim = int(part.replace("dim", ""))
                        elif part.startswith("cw"):
                            loaded_cw = float(part.replace("cw", ""))

                    # Re-initialize model architecture to match the file
                    device = get_device()
                    model = Autoencoder(capacity=loaded_cap, embedding_dim=loaded_dim).to(device)
                    model.load_state_dict(torch.load(os.path.join("outputs", selected_file), map_location=device))

                    # Store in session state
                    st.session_state['trained_model'] = model
                    st.session_state['device'] = device
                    st.session_state['current_model_name'] = selected_file
                    st.session_state['loaded_cap'] = loaded_cap
                    st.session_state['loaded_dim'] = loaded_dim
                    st.session_state['loaded_cw'] = loaded_cw
                    st.success(f"Loaded {selected_file}")
                except Exception as e:
                    st.error(f"Error loading model: {e}")
        else:
            available_files = []
            st.info("No saved models found in /outputs. Train one to start!")
    elif architecture == "Variational Autoencoder":
        st.header("VAE Settings")
        # Unique keys added to sliders/inputs to prevent Streamlit widget clashes
        num_components = st.slider("Mixture Components (K)", min_value=1, max_value=20, value=10, key="vae_k")
        latent_dim = st.slider("Latent Dimension (The Bottleneck)", min_value=2, max_value=256, value=64, key="vae_dim")
        capacity = st.selectbox("Conv Capacity", options=[8, 16, 32], index=2, key="vae_cap")

        st.header("Training Settings")
        epochs = st.number_input("Epochs", min_value=1, max_value=20, value=5, key="vae_ep")
        lr = st.number_input("Learning Rate", value=1e-3, format="%.4f", key="vae_lr")
        batch_size = st.number_input("Batch Size", value=128, key="vae_bs")
        
        # KLD Weight slider replaces the Classification Weight slider
        kld_weight = st.slider("KL Divergence Weight (Beta)", min_value=0.0, max_value=5.0, value=1.0, step=0.1, key="vae_kld")

        train_button = st.button("🚀 Train GM-VAE Model")

        st.header("📂 Model Browser")
        os.makedirs("outputs", exist_ok=True)
        # Scan for VAE specific checkpoints
        available_files = [f for f in os.listdir("outputs") if f.endswith(".pth") and f.startswith("vae_")]

        if available_files:
            selected_file = st.selectbox("Load an existing VAE checkpoint:", sorted(available_files))

            if st.button("🔄 Load Selected VAE"):
                try:
                    # Parse params from filename (e.g., vae_cap16_dim64_kld1.0_epoch5.pth)
                    # Special case: vae_model.pth from train_vae.py uses capacity=32, dim=64
                    if selected_file == "vae_model.pth":
                        loaded_cap, loaded_dim, loaded_kld, loaded_k = 32, 64, 1.0, 10
                    else:
                        parts = selected_file.replace(".pth", "").split("_")
                        loaded_cap, loaded_dim, loaded_kld, loaded_k = 16, 64, 1.0, 10
                        for part in parts:
                            if part.startswith("cap"):
                                loaded_cap = int(part.replace("cap", ""))
                            elif part.startswith("dim"):
                                loaded_dim = int(part.replace("dim", ""))
                            elif part.startswith("kld"):
                                loaded_kld = float(part.replace("kld", ""))
                            elif part.startswith("k") and not part.startswith("kld"):
                                loaded_k = int(part.replace("k", ""))

                    # Re-initialize model architecture
                    device = get_device()
                    model = ConvGMVAE(capacity=loaded_cap, latent_dim=loaded_dim, num_components=loaded_k).to(device)
                    model.load_state_dict(torch.load(os.path.join("outputs", selected_file), map_location=device))

                    # Store in session state
                    st.session_state['trained_model'] = model
                    st.session_state['device'] = device
                    st.session_state['current_model_name'] = selected_file
                    st.session_state['loaded_cap'] = loaded_cap
                    st.session_state['loaded_dim'] = loaded_dim
                    st.session_state['loaded_kld'] = loaded_kld
                    st.session_state['loaded_k'] = loaded_k
                    st.success(f"Loaded {selected_file}")
                except Exception as e:
                    st.error(f"Error loading model: {e}")
        else:
            st.info("No VAE models found in /outputs. Train one to start!")
# --- Initialize Data ---
@st.cache_resource
def get_data():
    transform = transforms.Compose([transforms.ToTensor()])
    train_ds = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    return train_ds

train_dataset = get_data()
# --- Training Logic (Variational Autoencoder) ---
if architecture == "Variational Autoencoder" and train_button:
    device = get_device()
    model = ConvGMVAE(capacity=capacity, latent_dim=latent_dim, num_components=num_components).to(device)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    
    st.subheader(f"Training VAE (Bottleneck: {latent_dim})")
    progress_bar = st.progress(0)
    loss_chart = st.empty()
    status_text = st.empty()
    
    # Track all three metrics for a multi-line chart
    loss_history = {"Total Loss": [], "Recon Loss": [], "KLD Loss": []}
    
    os.makedirs("outputs", exist_ok=True)
    model.train()
    
    for epoch in range(epochs):
        total_loss_epoch = 0.0
        total_recon_epoch = 0.0
        total_kld_epoch = 0.0
        
        for img, _ in train_loader:
            img = img.to(device)
            
            # Forward pass
            recon, mu, logvar, q_c, prior_mu, prior_logvar = model(img)
            loss, recon_loss, kld_loss = gmvae_loss_function(
                recon, img, mu, logvar, q_c, prior_mu, prior_logvar, kld_weight
            )
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            # Note: We divide by batch size here to get the average loss per item 
            # so the chart numbers are easier to read
            current_batch_size = img.size(0)
            total_loss_epoch += loss.item() / current_batch_size
            total_recon_epoch += recon_loss.item() / current_batch_size
            total_kld_epoch += kld_loss.item() / current_batch_size
        
        avg_loss = total_loss_epoch / len(train_loader)
        avg_recon = total_recon_epoch / len(train_loader)
        avg_kld = total_kld_epoch / len(train_loader)
        
        loss_history["Total Loss"].append(avg_loss)
        loss_history["Recon Loss"].append(avg_recon)
        loss_history["KLD Loss"].append(avg_kld)
        
        status_text.write(f"Epoch [{epoch+1}/{epochs}] | Total: {avg_loss:.4f} | Recon: {avg_recon:.4f} | KLD: {avg_kld:.4f}")
        loss_chart.line_chart(loss_history)
        progress_bar.progress((epoch + 1) / epochs)        
        
        # Save checkpoint
        checkpoint_name = f"vae_cap{capacity}_dim{latent_dim}_k{num_components}_kld{kld_weight}_epoch{epoch+1}.pth"
        save_path = os.path.join("outputs", checkpoint_name)
        torch.save(model.state_dict(), save_path)

    st.session_state['trained_model'] = model
    st.session_state['device'] = device
    st.session_state['current_model_name'] = f"vae_cap{capacity}_dim{latent_dim}_k{num_components}_epoch{epochs}.pth"
    st.session_state['loaded_cap'] = capacity
    st.session_state['loaded_dim'] = latent_dim
    st.session_state['loaded_k'] = num_components
    st.session_state['loaded_kld'] = kld_weight
    st.success(f"Training complete! Checkpoints saved to outputs/ (vae_cap{capacity}_dim{latent_dim}_k{num_components}_epoch*.pth)")
# --- Training Logic (Supervised Autoencoder only) ---
if architecture == "Supervised Autoencoder" and train_button:
    device = get_device()
    model = Autoencoder(capacity=capacity, embedding_dim=embedding_dim).to(device)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.MSELoss()
    criterion_ce = nn.CrossEntropyLoss()
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    
    st.subheader(f"Training Progress (Bottleneck: {embedding_dim})")
    progress_bar = st.progress(0)
    loss_chart = st.empty()
    status_text = st.empty()
    loss_history = []
    
    os.makedirs("outputs", exist_ok=True)
    loss_file_path = os.path.join("outputs", f"loss_cap{capacity}_dim{embedding_dim}.txt")
    with open(loss_file_path, "w") as f:
        pass  # truncate for new run
    model.train()
    for epoch in range(epochs):
        total_loss_epoch = 0.0
        total_recon_epoch = 0.0
        total_class_epoch = 0.0
        for img, label in train_loader:
            img = img.to(device)
            label = label.to(device)
            output, logits = model(img)
            loss_recon = criterion(output, img)
            loss_class = criterion_ce(logits, label)
            loss = loss_recon + (class_weight * loss_class)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss_epoch += loss.item()
            total_recon_epoch += loss_recon.item()
            total_class_epoch += loss_class.item()
        
        avg_loss = total_loss_epoch / len(train_loader)
        avg_recon = total_recon_epoch / len(train_loader)
        avg_class = total_class_epoch / len(train_loader)
        loss_history.append(avg_loss)
        status_text.write(f"Epoch [{epoch+1}/{epochs}] | Total: {avg_loss:.4f} | Recon: {avg_recon:.4f} | Class: {avg_class:.4f}")
        loss_chart.line_chart(loss_history)
        progress_bar.progress((epoch + 1) / epochs)        
        
        # Save a specific checkpoint for this run
        checkpoint_name = f"cap{capacity}_dim{embedding_dim}_cw{class_weight}_epoch{epoch+1}.pth"
        save_path = os.path.join("outputs", checkpoint_name)
        torch.save(model.state_dict(), save_path)
        
        # Track loss history in a text file for this run
        with open(loss_file_path, "a") as f:
            f.write(f"{epoch+1},{avg_loss}\n")    


    st.session_state['trained_model'] = model
    st.session_state['device'] = device
    st.session_state['current_model_name'] = f"cap{capacity}_dim{embedding_dim}_epoch{epochs}.pth"
    st.session_state['loaded_cap'] = capacity
    st.session_state['loaded_dim'] = embedding_dim
    st.session_state['loaded_cw'] = class_weight
    st.success(f"Training complete! Checkpoints saved to outputs/ (cap{capacity}_dim{embedding_dim}_epoch*.pth)")

# --- Visualizing Results (Tabs) ---
if architecture == "Supervised Autoencoder" and 'trained_model' in st.session_state:
    model = st.session_state['trained_model']
    device = st.session_state['device']
    model.eval()

    loaded_cap = st.session_state.get('loaded_cap', 16)
    loaded_dim = st.session_state.get('loaded_dim', 64)
    prefix = f"cap{loaded_cap}_dim{loaded_dim}"
    related_epochs = sorted([f for f in available_files if f.startswith(prefix)])

    tab_recon, tab_morph, tab_evolution, tab_map, tab_xray, tab_weights, tab_dreams, tab_arch, tab_fingerprints = st.tabs([
        "🖼️ Reconstructions",
        "🧪 Morphing",
        "⏳ Evolution",
        "🗺️ Latent Map",
        "🔍 Layer X-Ray",
        "🧠 Weight Inspector",
        "🔮 Filter Dreams & Mining",
        "🗺️ Architecture Map",
        "🧬 Embedding Fingerprints",
    ])

    # --- Tab: Reconstructions ---
    with tab_recon:
        st.subheader("Test Reconstructions")
        n_images = st.slider("Number of images to preview", 1, 10, 5)

        # Pin controls: use pinned indices if set, else random
        if "pinned_recon_indices" in st.session_state and len(st.session_state["pinned_recon_indices"]) == n_images:
            indices = st.session_state["pinned_recon_indices"]
        else:
            indices = list(np.random.choice(len(train_dataset), n_images, replace=False))

        pin_col, random_col, _ = st.columns([1, 1, 4])
        with pin_col:
            if st.button("📌 Pin these images"):
                st.session_state["pinned_recon_indices"] = list(indices)
                st.rerun()
        with random_col:
            if st.button("🎲 Random"):
                if "pinned_recon_indices" in st.session_state:
                    del st.session_state["pinned_recon_indices"]
                st.rerun()

        if "pinned_recon_indices" in st.session_state:
            st.caption("Images pinned. Switch models to compare reconstructions on the same originals.")

        cols = st.columns(n_images)
        for i, idx in enumerate(indices):
            img, _ = train_dataset[idx]
            with torch.no_grad():
                recon, _ = model(img.unsqueeze(0).to(device))
                recon = recon.cpu().squeeze().numpy()

            with cols[i]:
                fig, ax = plt.subplots(2, 1)
                ax[0].imshow(img.squeeze().numpy(), cmap='gray')
                ax[0].set_title(f"Original (idx {idx})")
                ax[0].axis('off')
                ax[1].imshow(recon, cmap='gray')
                ax[1].set_title("Reconstructed")
                ax[1].axis('off')
                st.pyplot(fig)

        st.divider()
        st.subheader("Explore the Latent Vector")
        st.write("This is the embedding vector (or whatever size you chose) for a random image.")
        sample_img, _ = train_dataset[indices[0]]
        with torch.no_grad():
            vector = model.encode(sample_img.unsqueeze(0).to(device)).cpu().numpy()
        st.bar_chart(vector.flatten())
        st.caption("Each bar above represents one 'feature' the model has learned to compress the image into.")

    # --- Tab: Morphing ---
    with tab_morph:
        st.subheader("🧪 Latent Space Morphing")
        st.write("Pick two images and watch the 'in-between' states as we blend their embedding summaries.")

        col_a, col_b = st.columns(2)
        default_a = st.session_state.get("pinned_morph_idx_a", 0)
        default_b = st.session_state.get("pinned_morph_idx_b", 1)
        with col_a:
            idx_a_input = st.number_input("Index of Image A", min_value=0, max_value=len(train_dataset)-1, value=default_a)
        with col_b:
            idx_b_input = st.number_input("Index of Image B", min_value=0, max_value=len(train_dataset)-1, value=default_b)

        pin_morph_col, unpin_morph_col, _ = st.columns([1, 1, 4])
        with pin_morph_col:
            if st.button("📌 Pin morph images", key="pin_morph"):
                st.session_state["pinned_morph_idx_a"] = idx_a_input
                st.session_state["pinned_morph_idx_b"] = idx_b_input
                st.rerun()
        with unpin_morph_col:
            if st.button("🎲 Unpin", key="unpin_morph"):
                if "pinned_morph_idx_a" in st.session_state:
                    del st.session_state["pinned_morph_idx_a"]
                if "pinned_morph_idx_b" in st.session_state:
                    del st.session_state["pinned_morph_idx_b"]
                st.rerun()

        if "pinned_morph_idx_a" in st.session_state:
            st.caption("Images pinned. Switch models to compare interpolations on the same originals.")

        idx_a = st.session_state.get("pinned_morph_idx_a", idx_a_input)
        idx_b = st.session_state.get("pinned_morph_idx_b", idx_b_input)

        img_a, _ = train_dataset[idx_a]
        img_b, _ = train_dataset[idx_b]

        with torch.no_grad():
            vec_a = model.encode(img_a.unsqueeze(0).to(device))
            vec_b = model.encode(img_b.unsqueeze(0).to(device))

        alpha = st.slider("Morphing Strength", 0.0, 1.0, 0.5)
        vec_morphed = (1 - alpha) * vec_a + alpha * vec_b

        with torch.no_grad():
            recon_morphed = model.decoder(vec_morphed).cpu().squeeze().numpy()

        m_col1, m_col2, m_col3 = st.columns(3)
        with m_col1:
            st.image(img_a.numpy().squeeze(), caption="Image A", width=200)
        with m_col2:
            st.image(recon_morphed, caption=f"Interpolated (Alpha={alpha})", width=200)
        with m_col3:
            st.image(img_b.numpy().squeeze(), caption="Image B", width=200)

    # --- Tab: Evolution ---
    with tab_evolution:
        st.subheader("⏳ Training Evolution")
        st.write("Compare the same image across different saved epochs of this model.")

        if related_epochs:
            col_left, col_right = st.columns([1, 1])
            with col_left:
                playback_speed = st.slider("Playback Speed (seconds)", min_value=0.1, max_value=1.0, value=0.3, step=0.1)
                play_button = st.button("▶️ Play Training History")
            with col_right:
                selected_epoch_file = st.select_slider("Select Epoch Checkpoint", options=related_epochs)

            image_placeholder = st.empty()

            comp_model = Autoencoder(capacity=loaded_cap, embedding_dim=loaded_dim).to(device)
            comp_model.load_state_dict(torch.load(os.path.join("outputs", selected_epoch_file), map_location=device))
            comp_model.eval()

            test_img, _ = train_dataset[0]
            with torch.no_grad():
                comp_recon, _ = comp_model(test_img.unsqueeze(0).to(device))
                comp_recon = comp_recon.cpu().squeeze().numpy()
            image_placeholder.image(comp_recon, caption=f"Reconstruction at {selected_epoch_file}", width=200)

            if play_button:
                for epoch_file in related_epochs:
                    comp_model = Autoencoder(capacity=loaded_cap, embedding_dim=loaded_dim).to(device)
                    comp_model.load_state_dict(torch.load(os.path.join("outputs", epoch_file), map_location=device))
                    comp_model.eval()

                    with torch.no_grad():
                        comp_recon, _ = comp_model(test_img.unsqueeze(0).to(device))
                        comp_recon = comp_recon.cpu().squeeze().numpy()
                    image_placeholder.image(comp_recon, caption=f"Reconstruction at {epoch_file}", width=200)
                    time.sleep(playback_speed)
        else:
            st.info("No epoch checkpoints found for this model. Train with this config to generate them.")

    # --- Tab: Latent Map ---
    with tab_map:
        st.subheader("🗺️ High-Dimensional Latent Map")
        st.write("Compare how different algorithms unroll the latent bottleneck.")

        reduction_method = st.radio(
            "Choose Algorithm:",
            ["PCA (Linear, Fast)", "t-SNE (Non-linear, Local Clusters)", "UMAP (Non-linear, Global Structure)"],
            horizontal=True,
            key="ae_reduce_method"
        )

        if st.button("Generate Latent Map", key="ae_map_btn"):
            model.eval()
            n_samples = 1000
            sample_indices = np.random.choice(len(train_dataset), n_samples, replace=False)

            embeddings = []
            labels = []
            with torch.no_grad():
                for idx in sample_indices:
                    img, label = train_dataset[idx]
                    # Standard AE just encodes directly to the vector
                    emb = model.encode(img.unsqueeze(0).to(device)).cpu().numpy().flatten()
                    embeddings.append(emb)
                    labels.append(str(int(label)))

            embeddings = np.array(embeddings)

            try:
                with st.spinner(f"Running {reduction_method.split(' ')[0]}..."):
                    # Offload to utils.py (Standard AE has no cluster centers to pass)
                    coords_2d, _ = get_2d_projections(embeddings, method=reduction_method)

                df = pd.DataFrame({
                    "x": coords_2d[:, 0],
                    "y": coords_2d[:, 1],
                    "digit": labels,
                })

                fig = px.scatter(
                    df, x="x", y="y", color="digit", 
                    opacity=0.7, 
                    title=f"Supervised AE Latent Space ({reduction_method.split(' ')[0]})"
                )
                st.plotly_chart(fig, width='stretch')
                
            except Exception as e:
                st.error(f"Error generating map: {e}")
    # --- Tab: Layer X-Ray ---
    with tab_xray:
        st.subheader("🔍 Layer X-Ray")
        st.write("Visualize the intermediate convolutional feature maps.")

        model.eval()

        img_idx = st.number_input(
            "Image index",
            min_value=0,
            max_value=len(train_dataset) - 1,
            value=0,
            key="xray_img_idx",
        )

        img, _ = train_dataset[img_idx]
        img_batch = img.unsqueeze(0).to(device)

        with torch.no_grad():
            activations = model.get_activations(img_batch)

        # Conv1 (14×14 per filter)
        st.caption("Conv1 output: 14×14 per filter")
        conv1_out = activations["conv1"]
        conv1_np = conv1_out.cpu().squeeze(0).numpy()
        n_filters1 = conv1_np.shape[0]
        n_cols = 4
        n_rows1 = (n_filters1 + n_cols - 1) // n_cols
        fig1, axes1 = plt.subplots(n_rows1, n_cols, figsize=(3 * n_cols, 3 * n_rows1))
        axes1 = axes1.flatten()
        for i in range(n_filters1):
            axes1[i].imshow(conv1_np[i], cmap="gray")
            axes1[i].set_title(f"Filter {i}")
            axes1[i].axis("off")
        for i in range(n_filters1, len(axes1)):
            axes1[i].axis("off")
        xray_col1, xray_col2 = st.columns([3, 1])
        with xray_col1:
            st.pyplot(fig1)
        with xray_col2:
            st.image(img.numpy().squeeze(), caption="Original input", width=200)
        plt.close(fig1)

        # Conv2 (7×7 per filter)
        st.divider()
        st.caption("Conv2 output: 7×7 per filter")
        conv2_out = activations["conv2"]
        conv2_np = conv2_out.cpu().squeeze(0).numpy()
        n_filters2 = conv2_np.shape[0]
        n_rows2 = (n_filters2 + n_cols - 1) // n_cols
        fig2, axes2 = plt.subplots(n_rows2, n_cols, figsize=(3 * n_cols, 3 * n_rows2))
        axes2 = axes2.flatten()
        for i in range(n_filters2):
            axes2[i].imshow(conv2_np[i], cmap="gray")
            axes2[i].set_title(f"Filter {i}")
            axes2[i].axis("off")
        for i in range(n_filters2, len(axes2)):
            axes2[i].axis("off")
        st.pyplot(fig2)
        plt.close(fig2)

    # --- Tab: Weight Inspector ---
    with tab_weights:
        model = st.session_state["trained_model"]

        st.subheader("The Compression Matrix (Encoder FC Layer)")
        fc_weights = model.encoder.fc.weight.data.cpu().numpy()
        fig_fc = px.imshow(
            fc_weights,
            color_continuous_scale="RdBu_r",
            aspect="auto",
            title="Linear Layer Weights (Embedding Dim x Flattened Features)",
        )
        st.plotly_chart(fig_fc, width='stretch')

        st.subheader("Conv1 Kernels (3x3)")
        conv1_weights = model.encoder.conv1.weight.data.cpu().squeeze().numpy()
        v1_min, v1_max = float(conv1_weights.min()), float(conv1_weights.max())
        for row_start in range(0, len(conv1_weights), 4):
            cols = st.columns(4)
            for j in range(4):
                i = row_start + j
                if i < len(conv1_weights):
                    with cols[j]:
                        fig_k, ax = plt.subplots(figsize=(2, 2))
                        im = ax.imshow(conv1_weights[i], cmap="RdBu_r", vmin=v1_min, vmax=v1_max)
                        ax.set_xticks([])
                        ax.set_yticks([])
                        ax.set_title(f"Filter {i}")
                        for (ri, ci), v in np.ndenumerate(conv1_weights[i]):
                            ax.text(ci, ri, f"{v:.2f}", ha="center", va="center", fontsize=8)
                        st.pyplot(fig_k)
                        plt.close(fig_k)

        st.subheader("Conv2 Kernels (3x3, averaged over input channels)")
        conv2_weights = model.encoder.conv2.weight.data.cpu().numpy()
        conv2_avg = conv2_weights.mean(axis=1)
        v2_min, v2_max = float(conv2_avg.min()), float(conv2_avg.max())
        for row_start in range(0, len(conv2_avg), 4):
            cols = st.columns(4)
            for j in range(4):
                i = row_start + j
                if i < len(conv2_avg):
                    with cols[j]:
                        fig_k, ax = plt.subplots(figsize=(2, 2))
                        im = ax.imshow(conv2_avg[i], cmap="RdBu_r", vmin=v2_min, vmax=v2_max)
                        ax.set_xticks([])
                        ax.set_yticks([])
                        ax.set_title(f"Filter {i}")
                        for (ri, ci), v in np.ndenumerate(conv2_avg[i]):
                            ax.text(ci, ri, f"{v:.2f}", ha="center", va="center", fontsize=8)
                        st.pyplot(fig_k)
                        plt.close(fig_k)

    # --- Tab: Filter Dreams & Mining ---
    with tab_dreams:
        model = st.session_state["trained_model"]
        model.eval()

        if st.button("🚀 Generate Full Filter Dashboard", key="full_dashboard_btn"):
            loaded_cap = st.session_state.get("loaded_cap", 16)

            # Vectorized dataset mining
            sample_indices = np.random.choice(len(train_dataset), 1000, replace=False)
            batch_imgs = []
            batch_labels = []
            for idx in sample_indices:
                img, label = train_dataset[idx]
                batch_imgs.append(img)
                batch_labels.append(label)
            batch_imgs = torch.stack(batch_imgs).to(device)
            batch_labels = torch.tensor(batch_labels)

            with torch.no_grad():
                batch_acts_conv1 = model.encoder.relu(model.encoder.conv1(batch_imgs))
                batch_acts_conv2 = model.encoder.relu(model.encoder.conv2(batch_acts_conv1))
            mean_acts_conv1 = batch_acts_conv1.mean(dim=(2, 3))
            mean_acts_conv2 = batch_acts_conv2.mean(dim=(2, 3))

            st.subheader("Conv1 Filters")
            for f in range(loaded_cap):
                # Gradient ascent (the dream)
                noise_img = torch.randn(1, 1, 28, 28, requires_grad=True, device=device)
                optimizer = torch.optim.Adam([noise_img], lr=0.1)

                for _ in range(50):
                    optimizer.zero_grad()
                    out = model.encoder.relu(model.encoder.conv1(noise_img))
                    loss = -out[0, f].mean()
                    loss.backward()
                    optimizer.step()

                dream_np = noise_img.detach().cpu().squeeze().numpy()
                d_min, d_max = dream_np.min(), dream_np.max()
                dream_np = (dream_np - d_min) / (d_max - d_min + 1e-8)

                # Top 2 mined images
                top_indices = mean_acts_conv1[:, f].argsort(descending=True)[:2]

                # Favorite digits: average activation per digit 0-9
                filter_acts = mean_acts_conv1[:, f].cpu()
                class_means = []
                for digit in range(10):
                    mask = (batch_labels == digit)
                    if mask.sum() > 0:
                        class_means.append(filter_acts[mask].mean().item())
                    else:
                        class_means.append(0.0)

                # Activation maps
                act_map1 = batch_acts_conv1[top_indices[0], f].cpu().numpy()
                act_map2 = batch_acts_conv1[top_indices[1], f].cpu().numpy()
                a1_min, a1_max = act_map1.min(), act_map1.max()
                a2_min, a2_max = act_map2.min(), act_map2.max()
                act_map1 = (act_map1 - a1_min) / (a1_max - a1_min + 1e-8)
                act_map2 = (act_map2 - a2_min) / (a2_max - a2_min + 1e-8)

                # Dashboard UI
                st.divider()
                st.subheader(f"Filter {f}")

                cols = st.columns([1.5, 1, 1, 1, 1, 2])
                with cols[0]:
                    st.image(dream_np, caption="Dream", width='stretch')
                with cols[1]:
                    img1 = batch_imgs[top_indices[0]].cpu().squeeze().numpy()
                    st.image(img1, caption="Top Match 1", width='stretch')
                with cols[2]:
                    st.image(act_map1, caption="What it sees in Match 1", width='stretch')
                with cols[3]:
                    img2 = batch_imgs[top_indices[1]].cpu().squeeze().numpy()
                    st.image(img2, caption="Top Match 2", width='stretch')
                with cols[4]:
                    st.image(act_map2, caption="What it sees in Match 2", width='stretch')
                with cols[5]:
                    st.caption("Average Activation per Digit")
                    chart_data = pd.DataFrame({"Digit": [str(i) for i in range(10)], "Activation": class_means})
                    st.bar_chart(chart_data, x="Digit", y="Activation", height=150)

            # Conv2 filters
            st.divider()
            st.subheader("Conv2 Filters")
            for f in range(loaded_cap * 2):
                # Gradient ascent (the dream) - through conv1 and conv2
                noise_img = torch.randn(1, 1, 28, 28, requires_grad=True, device=device)
                optimizer = torch.optim.Adam([noise_img], lr=0.1)

                for _ in range(50):
                    optimizer.zero_grad()
                    out1 = model.encoder.relu(model.encoder.conv1(noise_img))
                    out2 = model.encoder.relu(model.encoder.conv2(out1))
                    loss = -out2[0, f].mean()
                    loss.backward()
                    optimizer.step()

                dream_np = noise_img.detach().cpu().squeeze().numpy()
                d_min, d_max = dream_np.min(), dream_np.max()
                dream_np = (dream_np - d_min) / (d_max - d_min + 1e-8)

                top_indices = mean_acts_conv2[:, f].argsort(descending=True)[:2]

                filter_acts = mean_acts_conv2[:, f].cpu()
                class_means = []
                for digit in range(10):
                    mask = (batch_labels == digit)
                    if mask.sum() > 0:
                        class_means.append(filter_acts[mask].mean().item())
                    else:
                        class_means.append(0.0)

                act_map1 = batch_acts_conv2[top_indices[0], f].cpu().numpy()
                act_map2 = batch_acts_conv2[top_indices[1], f].cpu().numpy()
                a1_min, a1_max = act_map1.min(), act_map1.max()
                a2_min, a2_max = act_map2.min(), act_map2.max()
                act_map1 = (act_map1 - a1_min) / (a1_max - a1_min + 1e-8)
                act_map2 = (act_map2 - a2_min) / (a2_max - a2_min + 1e-8)

                st.divider()
                st.subheader(f"Filter {f}")

                cols = st.columns([1.5, 1, 1, 1, 1, 2])
                with cols[0]:
                    st.image(dream_np, caption="Dream", width='stretch')
                with cols[1]:
                    img1 = batch_imgs[top_indices[0]].cpu().squeeze().numpy()
                    st.image(img1, caption="Top Match 1", width='stretch')
                with cols[2]:
                    st.image(act_map1, caption="What it sees in Match 1", width='stretch')
                with cols[3]:
                    img2 = batch_imgs[top_indices[1]].cpu().squeeze().numpy()
                    st.image(img2, caption="Top Match 2", width='stretch')
                with cols[4]:
                    st.image(act_map2, caption="What it sees in Match 2", width='stretch')
                with cols[5]:
                    st.caption("Average Activation per Digit")
                    chart_data = pd.DataFrame({"Digit": [str(i) for i in range(10)], "Activation": class_means})
                    st.bar_chart(chart_data, x="Digit", y="Activation", height=150)

    # --- Tab: Architecture Map ---
    with tab_arch:
        loaded_cap = st.session_state.get("loaded_cap", 16)
        loaded_dim = st.session_state.get("loaded_dim", 64)
        flat_size = loaded_cap * 2 * 7 * 7

        st.subheader("Autoencoder Pipeline")

        dot_graph = f"""
digraph NN {{
    rankdir=TB;
    node [shape=box, style=filled, fontname="Arial"];
    
    Input [label="Input Image\\n1 x 28 x 28", fillcolor=lightgrey];
    Conv1 [label="Encoder Conv1\\n3x3 Kernel, Stride 2\\nLeakyReLU", fillcolor=lightblue];
    Conv2 [label="Encoder Conv2\\n3x3 Kernel, Stride 2\\nLeakyReLU", fillcolor=lightblue];
    Flatten [label="Flatten", fillcolor=lightgrey];
    Bottleneck [label="Bottleneck (Linear)\\n{loaded_dim} Features", fillcolor=salmon];
    DecFC [label="Decoder Linear\\n{flat_size} Features", fillcolor=lightgreen];
    Reshape [label="Reshape\\n{loaded_cap*2} x 7 x 7", fillcolor=lightgrey];
    ConvT1 [label="Decoder ConvTranspose1\\n3x3 Kernel, Stride 2\\nLeakyReLU", fillcolor=lightgreen];
    ConvT2 [label="Decoder ConvTranspose2\\n3x3 Kernel, Stride 2\\nSigmoid", fillcolor=lightgreen];
    Output [label="Reconstructed Image\\n1 x 28 x 28", fillcolor=lightgrey];
    Classifier [label="Classifier (Linear)\\n10 Classes", fillcolor=gold];
    Logits [label="Digit Logits\\n10-dim", fillcolor=lightgrey];
Input -> Conv1 [label=" 1 x 28 x 28"];
Conv1 -> Conv2 [label=" {loaded_cap} x 14 x 14"];
Conv2 -> Flatten [label=" {loaded_cap*2} x 7 x 7"];
Flatten -> Bottleneck [label=" {flat_size}"];
Bottleneck -> DecFC [label=" {loaded_dim}"];
Bottleneck -> Classifier [label=" {loaded_dim}"];
Classifier -> Logits [label=" 10"];
DecFC -> Reshape [label=" {flat_size}"];
Reshape -> ConvT1 [label=" {loaded_cap*2} x 7 x 7"];
ConvT1 -> ConvT2 [label=" {loaded_cap} x 14 x 14"];
ConvT2 -> Output [label=" 1 x 28 x 28"];
}}
"""
        st.graphviz_chart(dot_graph, width='stretch')

        st.subheader("Layer Parameter Breakdown")
        model = st.session_state["trained_model"]

        param_rows = []
        for name, param in model.named_parameters():
            n = param.numel()
            param_rows.append((name, n))

        total = sum(n for _, n in param_rows)
        header = "| Layer | Parameters |\n|-------|------------|"
        body = "\n".join(f"| {name} | {n:,} |" for name, n in param_rows)
        st.markdown(f"{header}\n{body}\n\n**Total:** {total:,} parameters")

    # --- Tab: Embedding Fingerprints ---
    with tab_fingerprints:
        model = st.session_state["trained_model"]
        model.eval()

        if st.button("🔍 Analyze Digit Fingerprints", key="fingerprints_btn"):
            indices = np.random.choice(len(train_dataset), 2000, replace=False)

            batch_imgs = []
            batch_labels = []
            for idx in indices:
                img, label = train_dataset[idx]
                batch_imgs.append(img)
                batch_labels.append(label)
            batch_imgs = torch.stack(batch_imgs).to(device)
            batch_labels = torch.tensor(batch_labels)

            with torch.no_grad():
                embeddings = model.encode(batch_imgs).cpu()

            centroids = []
            for digit in range(10):
                mask = (batch_labels == digit)
                if mask.sum() > 0:
                    centroids.append(embeddings[mask].mean(dim=0))
                else:
                    centroids.append(torch.zeros(embeddings.shape[1]))
            centroids_tensor = torch.stack(centroids)

            st.subheader("The 10 Digit Fingerprints")
            loaded_dim = st.session_state.get("loaded_dim", 64)
            fig_fp = px.imshow(
                centroids_tensor.numpy(),
                color_continuous_scale="RdBu_r",
                aspect="auto",
                labels=dict(x=f"Latent Dimension (0-{loaded_dim-1})", y="Digit Class", color="Activation"),
                y=[str(i) for i in range(10)],
                title="Average Latent Vector per Digit",
            )
            st.plotly_chart(fig_fp, width='stretch')

            st.subheader("The 'Perfect' Digits")
            st.write("These are generated by passing the average embedding for each digit back through the decoder.")

            with torch.no_grad():
                ideal_recons = model.decoder(centroids_tensor.to(device)).cpu().squeeze().numpy()

            fp_cols = st.columns(10)
            for i in range(10):
                with fp_cols[i]:
                    st.image(ideal_recons[i], caption=str(i), width='stretch')

elif architecture == "Variational Autoencoder" and 'trained_model' in st.session_state:
    # --- Visualizing Results (Variational Autoencoder) ---
    model = st.session_state['trained_model']
    device = st.session_state['device']
    model.eval()

    loaded_cap = st.session_state.get('loaded_cap', 32)
    loaded_dim = st.session_state.get('loaded_dim', 64)
    
    # Helper function: Allows us to pass an arbitrary latent vector (z) straight into the decoder
    def decode_z(z_vector):
        with torch.no_grad():
            z_proj = model.decoder_input(z_vector)
            # Reshape dynamically based on the model's capacity
            z_reshaped = z_proj.view(-1, loaded_cap * 2, 7, 7)
            recon = model.decoder(z_reshaped)
        return recon

    tab_recon, tab_morph, tab_dream, tab_map = st.tabs([
        "🖼️ Reconstructions",
        "🧪 Smooth Morphing",
        "✨ Dreaming (Generation)",
        "🗺️ Latent Map",
    ])

    # --- Tab: Reconstructions ---
    with tab_recon:
        st.subheader("Test Reconstructions")
        st.write("Notice how the VAE might be slightly blurrier than the Supervised AE. This is the trade-off for a continuous latent space!")
        n_images = st.slider("Number of images to preview", 1, 10, 5, key="vae_n_imgs")

        # Reuse pinned logic to allow direct comparison if you switch back and forth between AE and VAE
        if "pinned_recon_indices" in st.session_state and len(st.session_state["pinned_recon_indices"]) == n_images:
            indices = st.session_state["pinned_recon_indices"]
        else:
            indices = list(np.random.choice(len(train_dataset), n_images, replace=False))

        pin_col, random_col, _ = st.columns([1, 1, 4])
        with pin_col:
            if st.button("📌 Pin these images", key="vae_pin"):
                st.session_state["pinned_recon_indices"] = list(indices)
                st.rerun()
        with random_col:
            if st.button("🎲 Random", key="vae_rand"):
                if "pinned_recon_indices" in st.session_state:
                    del st.session_state["pinned_recon_indices"]
                st.rerun()

        cols = st.columns(n_images)
        for i, idx in enumerate(indices):
            img, _ = train_dataset[idx]
            with torch.no_grad():
                # The VAE returns (reconstruction, mu, logvar)
                outputs = model(img.unsqueeze(0).to(device))
                recon = outputs[0]
                mu = outputs[1]
                logvar = outputs[2]
                recon = recon.cpu().squeeze().numpy()
                mu = mu.cpu().squeeze().numpy()
                logvar = logvar.cpu().squeeze().numpy()

            with cols[i]:
                fig, ax = plt.subplots(2, 1)
                ax[0].imshow(img.squeeze().numpy(), cmap='gray')
                ax[0].set_title(f"Original (idx {idx})")
                ax[0].axis('off')
                ax[1].imshow(recon, cmap='gray')
                ax[1].set_title("Reconstructed")
                ax[1].axis('off')
                st.pyplot(fig)

    # --- Tab: Morphing ---
    with tab_morph:
        st.subheader("🧪 Smooth Latent Space Morphing")
        st.write("Because the GM-VAE forces the latent space to be continuous, we can transition cleanly between digits.")

        col_a, col_b = st.columns(2)
        default_a = st.session_state.get("pinned_morph_idx_a", 0)
        default_b = st.session_state.get("pinned_morph_idx_b", 1)
        
        with col_a:
            idx_a_input = st.number_input("Index of Image A", min_value=0, max_value=len(train_dataset)-1, value=default_a, key="vae_idx_a")
        with col_b:
            idx_b_input = st.number_input("Index of Image B", min_value=0, max_value=len(train_dataset)-1, value=default_b, key="vae_idx_b")

        idx_a = st.session_state.get("pinned_morph_idx_a", idx_a_input)
        idx_b = st.session_state.get("pinned_morph_idx_b", idx_b_input)

        img_a, _ = train_dataset[idx_a]
        img_b, _ = train_dataset[idx_b]

        with torch.no_grad():
            # Safe indexing: outputs[1] is the mean (mu) of the encoded image
            out_a = model(img_a.unsqueeze(0).to(device))
            mu_a = out_a[1]
            
            out_b = model(img_b.unsqueeze(0).to(device))
            mu_b = out_b[1]

        alpha = st.slider("Morphing Strength", 0.0, 1.0, 0.5, key="vae_alpha")
        
        # Draw a straight line between the two cluster centers in 64D space
        z_morphed = (1 - alpha) * mu_a + alpha * mu_b

        with torch.no_grad():
            recon_morphed = decode_z(z_morphed).cpu().squeeze().numpy()

        m_col1, m_col2, m_col3 = st.columns(3)
        with m_col1:
            st.image(img_a.numpy().squeeze(), caption="Image A", width=200)
        with m_col2:
            st.image(recon_morphed, caption=f"Interpolated (Alpha={alpha})", width=200)
        with m_col3:
            st.image(img_b.numpy().squeeze(), caption="Image B", width=200)

    # --- Tab: Dreaming ---
    with tab_dream:
        st.subheader("✨ Directed Dreaming (GM-VAE Generation)")
        st.write("Unlike a standard VAE, a GM-VAE learns $K$ distinct distributions. Pick a specific cluster and sample random images directly from it!")
        
        with torch.no_grad():
            # Pass a dummy tensor just to trigger the forward pass and extract the learned prior
            dummy_img = torch.zeros(1, 1, 28, 28).to(device)
            dummy_out = model(dummy_img)
            
            # Safe indexing: outputs[4] is prior_mu, outputs[5] is prior_logvar
            prior_mu = dummy_out[4]       
            prior_logvar = dummy_out[5]   
            
            num_clusters = prior_mu.size(0)

        col_settings, col_visuals = st.columns([1, 2])
        
        with col_settings:
            # Dynamically generate the dropdown based on the K value of the loaded model
            selected_cluster = st.selectbox("Select a Learned Cluster to sample from:", range(num_clusters), key="vae_cluster_select")
            n_dreams = st.slider("Images to generate", 1, 10, 5, key="vae_n_dreams")
            dream_btn = st.button("🎲 Dream digits", key="vae_dream_btn")

        if dream_btn:
            with col_visuals:
                st.write(f"Sampling from Cluster {selected_cluster}...")
                
                # Get the mean and variance for the specific cluster the user chose
                cluster_mu = prior_mu[selected_cluster]
                cluster_std = torch.exp(0.5 * prior_logvar[selected_cluster])
                
                # Sample random noise and shift/scale it to match the cluster's exact shape
                eps = torch.randn(n_dreams, loaded_dim).to(device)
                random_z = cluster_mu + eps * cluster_std
                
                with torch.no_grad():
                    dreamed_imgs = decode_z(random_z).cpu().squeeze().numpy()
                
                cols = st.columns(n_dreams)
                for i in range(n_dreams):
                    with cols[i]:
                        img_to_show = dreamed_imgs if n_dreams == 1 else dreamed_imgs[i]
                        st.image(img_to_show, caption=f"Sample {i+1}", width=150)

    # --- Tab: Latent Map ---
    with tab_map:
        st.subheader("🗺️ High-Dimensional Latent Map")
        st.write("Notice where the GM-VAE placed its learned cluster centers (marked with X) compared to the raw data!")
        
        reduction_method = st.radio(
            "Choose Algorithm:", 
            ["PCA (Linear, Fast)", "t-SNE (Non-linear, Local Clusters)", "UMAP (Non-linear, Global Structure)"], 
            horizontal=True, 
            key="vae_reduce_method"
        )

        if st.button("Generate Map", key="vae_map_btn"):
            n_samples = 1000
            sample_indices = np.random.choice(len(train_dataset), n_samples, replace=False)

            embeddings = []
            labels = []
            
            with torch.no_grad():
                for idx in sample_indices:
                    img, label = train_dataset[idx]
                    
                    # GM-VAE signature: recon, mu, logvar, q_c, prior_mu, prior_logvar
                    outputs = model(img.unsqueeze(0).to(device))
                    mu = outputs[1]
                    prior_mu = outputs[4]
                    
                    embeddings.append(mu.cpu().numpy().flatten())
                    labels.append(str(int(label)))
                
                # The K cluster centers the GM-VAE learned
                cluster_centers = prior_mu.cpu().numpy()

            embeddings = np.array(embeddings)

            try:
                with st.spinner(f"Running {reduction_method.split(' ')[0]}..."):
                    # Pass the centers in so they get reduced alongside the data
                    coords_2d_data, coords_2d_centers = get_2d_projections(
                        embeddings, 
                        cluster_centers=cluster_centers, 
                        method=reduction_method
                    )
                
                df = pd.DataFrame({
                    "x": coords_2d_data[:, 0],
                    "y": coords_2d_data[:, 1],
                    "digit": labels,
                })
                
                fig = px.scatter(
                    df, x="x", y="y", color="digit", 
                    opacity=0.6, 
                    title=f"GM-VAE Latent Space ({reduction_method.split(' ')[0]})"
                )
                
                # Overlay the Learned GM-VAE Cluster Centers
                fig.add_scatter(
                    x=coords_2d_centers[:, 0], 
                    y=coords_2d_centers[:, 1], 
                    mode='markers', 
                    marker=dict(symbol='x', size=15, color='black', line=dict(width=2, color='white')),
                    name="Learned Cluster Centers"
                )
                
                st.plotly_chart(fig, width='stretch')
                
            except Exception as e:
                st.error(f"Error generating map: {e}")