import streamlit as st
import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from pipeline import AdvancedFashionCNN, ActivationTelemetryEngine, transform_custom_image

# --- DASHBOARD LAYOUT MANAGEMENT ---
st.set_page_config(
    page_title="AI Portfolio | Computer Vision Explainer Engine",
    page_icon="🧠",
    layout="wide"
)

# Load global definitions
CLASS_NAMES = ["T-shirt/top", "Trouser", "Pullover", "Dress", "Coat", "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"]

@st.cache_resource
def bootstrap_model_engine():
    """Compiles framework layers and acts as a safe memory cache layer."""
    instance = AdvancedFashionCNN()
    # Weights initialize gracefully for telemetry presentations
    return instance, ActivationTelemetryEngine(instance)

model, tele_engine = bootstrap_model_engine()

# --- SIDEBAR INTERFACE STRUCTURE ---
with st.sidebar:
    st.title("PKRAMAN")
    st.caption("🚀 Deep Learning & Computer Vision Research Engineer")
    st.markdown("🌐 [GitHub](https://github.com) | 💼 [LinkedIn](https://linkedin.com)")
    st.markdown("---")
    
    st.subheader("🛠️ Core Architecture Features")
    st.markdown("""
    - **Squeeze-and-Excitation (SE):** Multi-channel attention mapping.
    - **Grad-CAM Tracking:** Latent-space backpropagation.
    - **Residual Paths:** Multi-layer identity mappings.
    """)
    st.markdown("---")
    
    input_method = st.radio("Asset Ingestion Strategy", ["Database Testing Samples", "Custom Image Upload"])

# --- MAIN SCREEN PORTFOLIO HERO SEGMENT ---
st.title("🧠 Advanced Explainer Engine: Interpretability Matrix for Fashion CNNs")
st.caption("Interactive Neural Layer Visualization & Explainer Engine for Production Models")

# Metric KPIs
m1, m2, m3 = st.columns(3)
m1.metric("Feature Extractor Layer", "Residual Conv + SE-Attention", "Dynamic Hooks Active")
m2.metric("Gradient Localization Engine", "Grad-CAM Backend", "Target: Layer 3 Conv")
m3.metric("Dataset Domain Space", "FashionMNIST", "10-Class Optimization Loop")

st.markdown("---")

# Setup Image Core Inputs
image_tensor, reference_label = None, None

if input_method == "Database Testing Samples":
    # Download baseline test sets for quick selection verification
    test_dataset = torchvision.datasets.FashionMNIST(
        root='./data', train=False, download=True,
        transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
    )
    sample_index = st.slider("Select Sample Index from Database Validation Set", 0, len(test_dataset)-1, 16)
    image_tensor, reference_label = test_dataset[sample_index]
else:
    uploaded_file = st.file_uploader("Upload an Image File (Auto-Grayscale & Resized to 28x28)", type=["jpg", "jpeg", "png"])
    if uploaded_file is not None:
        try:
            image_tensor = transform_custom_image(uploaded_file)
        except Exception as e:
            st.error(f"Ingestion compilation format error: {e}")

# --- METRIC GENERATION DISPLAY SEQUENCE ---
if image_tensor is not None:
    # Forward pass predictions
    with torch.no_grad():
        logits = model(image_tensor.unsqueeze(0).to(tele_engine.device))
        probs = torch.softmax(logits, dim=1).squeeze().cpu().numpy()
    
    predicted_class = np.argmax(probs)
    confidence = probs[predicted_class]

    # Run interpretability computations
    saliency_map = tele_engine.compute_saliency(image_tensor)
    gradcam_map = tele_engine.compute_gradcam(image_tensor)

    # UI presentation space
    col_img, col_metrics = st.columns([1.2, 0.8])
    
    with col_img:
        st.subheader("🎯 Visual Explainer Core Diagnostics")
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        
        # Raw Display Channel
        raw_img = image_tensor.squeeze().numpy()
        axes[0].imshow(raw_img, cmap='gray')
        axes[0].set_title("Input Raw Channel")
        axes[0].axis('off')

        # Saliency Visual Matrix
        axes[1].imshow(saliency_map, cmap='hot')
        axes[1].set_title("Pixel-Level Saliency")
        axes[1].axis('off')

        # Grad-CAM Overlay Integration Map
        axes[2].imshow(raw_img, cmap='gray')
        axes[2].imshow(gradcam_map, cmap='jet', alpha=0.5)
        axes[2].set_title("Grad-CAM Latent Attentions")
        axes[2].axis('off')

        st.pyplot(fig)

    with col_metrics:
        st.subheader("📋 Inference Analysis Profiles")
        
        st.metric("Predicted Target Class Name", f"{CLASS_NAMES[predicted_class]}")
        st.metric("Model Execution Certainty Score", f"{confidence:.2%}")
        
        if reference_label is not None:
            st.metric("Ground-Truth Database Metric", f"{CLASS_NAMES[reference_label]}")
            if predicted_class == reference_label:
                st.success("🟢 True Ingestion Classification Verified.")
            else:
                st.error("🔴 Model Space Misclassification Variance Documented.")
                
        # Mini probability trace tracking distribution
        st.markdown("**Output Tensor Class Distributions:**")
        chart_data = pd.DataFrame({"Activation Match": probs}, index=CLASS_NAMES)
        st.bar_chart(chart_data)
else:
    st.info("⚡ Awaiting input array matrices. Choose an engineering input strategy or index sample from the sidebar.")