import streamlit as st
import plotly.graph_objects as go
from pipeline import MechanisticEngine

# Configure asynchronous widescreen UI layout
st.set_page_config(layout="wide", page_title="Mechanistic Interpretability Dashboard")

st.title("🧠 Deep Inference & Mechanistic Interpretability Engine")
st.markdown("""
    This interface tracks internal **Residual Stream Vector Dynamics** inside transformers. 
    By splitting operations between an decoupled engine pipeline and frontend visualization layers, 
    we can probe features without injecting inference execution latency.
""")

# Cache the Engine instantiation to prevent re-initializing model weights on slider changes
@st.cache_resource
def get_ml_engine():
    return MechanisticEngine(model_name="tiny-stories-1M", device="cpu")

with st.spinner("Initializing Mechanistic Pipeline Engine (Caching Weights)..."):
    engine = get_ml_engine()

# Configure global control panel sidebar
st.sidebar.header("Pipeline Hardware & Model Controls")
st.sidebar.markdown(f"**Loaded Model Architecture:** `{engine.model_name}`")

input_text = st.sidebar.text_area(
    "Prompt Text Sequence input:", 
    value="The brave knight rode his horse into the dark forest and found a hidden"
)

# Populate sliders dynamically by inspecting instantiated Engine specifications
selected_layer = st.sidebar.slider(
    "Target Residual Layer Hook:", 
    min_value=0, 
    max_value=engine.n_layers - 1, 
    value=0
)

if st.button("Trigger Telemetry Engine Probe"):
    if not input_text.strip():
        st.error("Invalid sequence input. String sequence cannot be null.")
    else:
        # Step 1: Tokenize sequence through the engine pipeline
        tokens, str_tokens = engine.tokenize_sequence(input_text)
        
        # Step 2: Extract low-level telemetry dictionary from the target hooks
        with st.spinner(f"Intercepting activations along layer {selected_layer} residual highway..."):
            telemetry_data = engine.run_telemetry_pass(tokens, selected_layer)
            
        st.success("Internal vector calculations complete.")
        
        # --- UI VISUALIZATION GRID ---
        col1, col2 = st.columns([1, 1])
        
        with col1:
            st.subheader(f"📊 Layer {selected_layer} Residual L2 Vector Norms")
            st.caption("Tracks the scalar structural magnitude of each specific token sequence slice.")
            
            # Draw spatial metrics bar graph via Plotly
            fig_norm = go.Figure(data=[go.Bar(
                x=str_tokens, 
                y=telemetry_data["l2_norms"], 
                marker_color='rgb(115, 103, 240)'
            )])
            fig_norm.update_layout(
                margin=dict(l=20, r=20, t=20, b=20),
                xaxis_title="Token Sub-Sequence",
                yaxis_title="L2 Norm Magnitude (||v||₂)",
                template="plotly_white"
            )
            st.plotly_chart(fig_norm, use_container_width=True)

        with col2:
            st.subheader("🎯 Unembedding Next-Token Softmax Projections")
            st.caption("Top 5 matrix projections mapped out of the absolute last position of the attention loop.")
            
            # Draw horizontal predictive distribution metrics
            fig_probs = go.Figure(data=[go.Bar(
                x=telemetry_data["top_probabilities"], 
                y=telemetry_data["top_tokens"], 
                orientation='h',
                marker_color='rgb(40, 199, 111)'
            )])
            fig_probs.update_layout(
                margin=dict(l=20, r=20, t=20, b=20),
                xaxis_title="Softmax Probability Distribution",
                yaxis_title="Vocabulary Tokens",
                template="plotly_white"
            )
            st.plotly_chart(fig_probs, use_container_width=True)
            
        # --- DEEP STRUCTURAL DATA VIEW ---
        st.markdown("---")
        st.subheader("⚙️ High-Density Internal Activation Metadata Matrix")
        
        with st.expander("Expand Activation Tensor Blueprints"):
            st.json({
                "Engine Targeted Module": f"blocks.{selected_layer}.hook_resid_post",
                "Computational Tensor Dimension Architecture": telemetry_data["tensor_shape"],
                "Model Structural Width (d_model)": engine.d_model,
                "Configured Multi-Head Multiplicity": engine.n_heads,
                "Precision Quantization Mapping Matrix": "Float32 (Standard Uncompressed Precision CPU Space)"
            })