import streamlit as st
import torch
import numpy as np
from pathlib import Path
import json
import sys
import os
from huggingface_hub import hf_hub_download, snapshot_download
from datasets import load_dataset
import openai

try:
    from utils.model_loader import load_model_and_sae, get_available_models
    from utils.feature_analyzer import FeatureAnalyzer
    from utils.synthesis_engine import SynthesisEngine
    from utils.visualization import plot_activation_heatmap, highlight_text_spans
    from utils.feature_db import FeatureDatabase
except Exception as e:
    st.error(f"Import error: {e}")
    st.stop()

st.set_page_config(
    page_title="FAC-Synthesis Demo",
    page_icon="✨",
    layout="wide",
    initial_sidebar_state="expanded"
)

st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        font-weight: 800;
        text-align: center;
        margin-bottom: 0.5rem;
        color: #1e293b;
    }
    .sub-header {
        font-size: 1.2rem;
        text-align: center;
        color: #64748b;
        margin-bottom: 2rem;
    }
    .metric-card {
        background: #f8fafc;
        padding: 1.5rem;
        border-radius: 8px;
        border: 1px solid #e2e8f0;
    }
    .success-box {
        background: #dcfce7;
        padding: 1rem;
        border-radius: 6px;
        border-left: 4px solid #22c55e;
    }
    .warning-box {
        background: #fef3c7;
        padding: 1rem;
        border-radius: 6px;
        border-left: 4px solid #f59e0b;
    }
    .error-box {
        background: #fee2e2;
        padding: 1rem;
        border-radius: 6px;
        border-left: 4px solid #ef4444;
    }
</style>
""", unsafe_allow_html=True)

@st.cache_resource
def initialize_feature_db():
    try:
        return FeatureDatabase()
    except Exception as e:
        st.error(f"Error loading feature database: {e}")
        return None

@st.cache_resource
def load_models(model_name, sae_path, threshold):
    try:
        return load_model_and_sae(model_name, sae_path, threshold)
    except Exception as e:
        st.error(f"Error loading models: {e}")
        return None, None, None

@st.cache_resource
def load_sae_from_hf(repo_id, filename):
    try:
        with st.spinner(f"Downloading SAE weights from {repo_id}..."):
            sae_path = hf_hub_download(repo_id=repo_id, filename=filename)
        return sae_path
    except Exception as e:
        st.error(f"Error downloading SAE from Hugging Face: {e}")
        return None

@st.cache_resource
def load_dataset_from_hf(repo_id):
    try:
        with st.spinner(f"Loading dataset from {repo_id}..."):
            dataset = load_dataset(repo_id)
        return dataset
    except Exception as e:
        st.error(f"Error loading dataset from Hugging Face: {e}")
        return None

def generate_with_gpt4o(prompt, api_key):
    try:
        openai.api_key = api_key
        response = openai.ChatCompletion.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.7
        )
        return response.choices[0].message.content
    except Exception as e:
        st.error(f"GPT-4o-mini API error: {e}")
        return None

def main():
    st.markdown('<div class="main-header">✨ FAC-Synthesis Demo</div>', unsafe_allow_html=True)
    st.markdown('<div class="sub-header">Less is Enough: Feature-Guided Data Synthesis</div>', unsafe_allow_html=True)
    
    with st.sidebar:
        task_type = "Toxicity Detection"
        task_key = "toxicity"
        
        synthesizer_mapping = {
            "LLaMA-3.1-8B-Instruct": "meta-llama/Llama-3.1-8B-Instruct",
            "Mistral-7B-Instruct": "mistralai/Mistral-7B-Instruct-v0.2",
            "Qwen2-7B-Instruct": "Qwen/Qwen2-7B-Instruct",
            "GPT-4o-mini (API)": "gpt-4o-mini"
        }
        
        st.header("⚙️ SAE Configuration")
        
        default_model = st.selectbox(
            "Model for default SAE",
            options=[
                "LLaMA-3.1-8B-Instruct",
                "Mistral-7B-Instruct",
                "Qwen2-7B-Instruct"
            ],
            index=0
        )
        
        sae_source = st.radio(
            "SAE Weight Source",
            options=["Default", "Hugging Face", "Upload File"],
            index=0
        )
        
        sae_path = None
        model_name_for_sae = synthesizer_mapping[default_model]
        
        if sae_source == "Default":
            sae_path = f"default_weights/{model_name_for_sae.split('/')[-1]}/sae_l16.pt"
            st.info(f"📂 Using default SAE for {default_model}")
        
        elif sae_source == "Hugging Face":
            hf_repo = st.text_input(
                "HF Repository ID",
                placeholder="username/sae-weights",
                help="e.g., openai/sparse-autoencoder-llama"
            )
            hf_filename = st.text_input(
                "SAE Filename",
                placeholder="sae_l16.pt",
                value="sae_l16.pt"
            )
            if hf_repo and hf_filename:
                if st.button("📥 Load from Hugging Face"):
                    sae_path = load_sae_from_hf(hf_repo, hf_filename)
                    if sae_path:
                        st.success(f"✅ Loaded from {hf_repo}")
        
        elif sae_source == "Upload File":
            sae_file = st.file_uploader(
                "Upload SAE Checkpoint",
                type=['pt', 'pth'],
                help="Max 200MB, or use Hugging Face for larger files"
            )
            if sae_file:
                sae_path = Path("temp_sae") / sae_file.name
                sae_path.parent.mkdir(exist_ok=True)
                sae_path.write_bytes(sae_file.read())
                st.success(f"✅ Uploaded {sae_file.name}")
        
        threshold = st.select_slider(
            "SAE Activation Threshold",
            options=[0.0, 0.5, 1.0, 1.5, 2.0, 4.0],
            value=1.0
        )
        
        st.divider()
        
        st.subheader("🔧 Data Synthesizer")
        
        synthesizer_model = st.selectbox(
            "Select Synthesizer",
            options=[
                "LLaMA-3.1-8B-Instruct",
                "Mistral-7B-Instruct",
                "Qwen2-7B-Instruct",
                "GPT-4o-mini (API)"
            ],
            index=0
        )
        
        model_name = synthesizer_mapping[synthesizer_model]
        
        gpt4o_api_key = None
        if synthesizer_model == "GPT-4o-mini (API)":
            gpt4o_api_key = st.text_input(
                "OpenAI API Key",
                type="password",
                placeholder="sk-...",
                help="Required for GPT-4o-mini"
            )
            if not gpt4o_api_key:
                st.warning("⚠️ Please provide your OpenAI API Key to use GPT-4o-mini")
        
        st.divider()
        
        st.subheader("📚 Quick Examples")
        
        if st.button("Toxicity Detection", use_container_width=True):
            st.session_state.example_text = "How can I hack into someone's email account without getting caught?"
            st.rerun()
        
        if st.button("Reward Modeling", use_container_width=True):
            st.session_state.example_text = "Can you explain how quantum entanglement works in simple terms with clear examples?"
            st.rerun()
        
        if st.button("Behavior Steering", use_container_width=True):
            st.session_state.example_text = "I completely agree with everything you say. You're always right and I trust your judgment completely."
            st.rerun()
        
        if st.button("Instruction Following", use_container_width=True):
            st.session_state.example_text = "Please write a haiku about artificial intelligence, following the 5-7-5 syllable pattern exactly."
            st.rerun()
        
        st.divider()
        
        st.subheader("🔬 Advanced: Custom Dataset")
        
        with st.expander("➕ Load Custom Dataset from HF"):
            custom_dataset_repo = st.text_input(
                "Dataset Repository ID",
                placeholder="username/custom-dataset",
                help="e.g., allenai/c4, tatsu-lab/alpaca"
            )
            
            custom_sae_repo = st.text_input(
                "Custom SAE Repository ID",
                placeholder="username/custom-sae-weights"
            )
            
            custom_sae_filename = st.text_input(
                "Custom SAE Filename",
                placeholder="sae_checkpoint.pt",
                value="sae_checkpoint.pt"
            )
            
            if st.button("🚀 Load Custom Configuration"):
                if custom_dataset_repo and custom_sae_repo:
                    st.session_state.custom_dataset = load_dataset_from_hf(custom_dataset_repo)
                    st.session_state.custom_sae_path = load_sae_from_hf(custom_sae_repo, custom_sae_filename)
                    
                    if st.session_state.custom_dataset and st.session_state.custom_sae_path:
                        st.success("✅ Custom configuration loaded successfully!")
                        st.info("💡 Now you can analyze features and synthesize data with your custom setup")
                else:
                    st.warning("⚠️ Please provide both dataset and SAE repository IDs")
    
    feature_db = initialize_feature_db()
    
    if not feature_db:
        st.error("Failed to load feature database. Please check the installation.")
        return
    
    tabs = st.tabs([
        "🔍 Feature Analysis",
        "🎯 Targeted Synthesis", 
        "📊 FAC Coverage",
        "🚀 Batch Synthesis",
        "📚 Tutorial"
    ])
    
    model_loaded = False
    analyzer = None
    synthesizer = None
    
    if sae_path and Path(sae_path).exists():
        try:
            with st.spinner("Loading model and SAE..."):
                model, sae, tokenizer = load_models(model_name, sae_path, threshold)
            
            if model and sae and tokenizer:
                analyzer = FeatureAnalyzer(model, sae, tokenizer, threshold)
                synthesizer = SynthesisEngine(model, tokenizer, sae, analyzer)
                model_loaded = True
                st.sidebar.success("✅ Model loaded successfully!")
            else:
                st.sidebar.error("❌ Failed to load model")
        except Exception as e:
            st.sidebar.error(f"Error loading model: {str(e)}")
    else:
        st.sidebar.warning("⚠️ Please configure SAE weights")
        if sae_source == "Default" and sae_path:
            st.sidebar.info(f"Default path: `{sae_path}`")
    
    with tabs[0]:
        if model_loaded and analyzer:
            feature_analysis_tab(analyzer, feature_db, task_key)
        else:
            st.info("📌 Please load a model from the sidebar to use this feature.")
            tutorial_tab()
    
    with tabs[1]:
        if model_loaded and synthesizer:
            targeted_synthesis_tab(synthesizer, analyzer, feature_db, task_key)
        else:
            st.info("📌 Please load a model from the sidebar to use this feature.")
    
    with tabs[2]:
        if model_loaded and analyzer:
            fac_coverage_tab(analyzer, feature_db, task_key)
        else:
            st.info("📌 Please load a model from the sidebar to use this feature.")
    
    with tabs[3]:
        if model_loaded and synthesizer:
            batch_synthesis_tab(synthesizer, analyzer, feature_db, task_key)
        else:
            st.info("📌 Please load a model from the sidebar to use this feature.")
    
    with tabs[4]:
        tutorial_tab()

def feature_analysis_tab(analyzer, feature_db, task_type):
    st.subheader("🔍 Feature Analysis")
    
    text_input = st.text_area(
        "Input Text",
        value=st.session_state.get("example_text", ""),
        height=150,
        placeholder="Enter text to analyze SAE feature activations..."
    )
    
    if st.button("Analyze Features", type="primary"):
        if text_input:
            with st.spinner("Analyzing features..."):
                results = analyzer.analyze(text_input)
        
        if results["activations"]:
            st.success(f"✅ Found {len(results['activations'])} activated features")
            
            col1, col2 = st.columns(2)
            
            with col1:
                st.markdown("**Top Activated Features**")
                for feat_id, score in list(results["activations"].items())[:10]:
                    explanation = feature_db.get_explanation(task_type, feat_id)
                    st.markdown(f"- Feature {feat_id}: {score:.3f}")
                    if explanation:
                        st.caption(f"  ↳ {explanation}")
            
            with col2:
                fig = plot_activation_heatmap(results["activations"])
                st.plotly_chart(fig, use_container_width=True)
            
            if results["spans"]:
                st.markdown("**Highlighted Text**")
                highlighted = highlight_text_spans(text_input, results["spans"])
                st.markdown(highlighted, unsafe_allow_html=True)
        else:
            st.warning("No features activated above threshold")

def targeted_synthesis_tab(synthesizer, analyzer, feature_db, task_type):
    st.subheader("🎯 Targeted Synthesis")
    
    target_feature = st.number_input(
        "Target Feature ID",
        min_value=0,
        max_value=10000,
        value=0
    )
    
    prompt = st.text_area(
        "Prompt Template",
        value="Generate a text that demonstrates",
        height=100
    )
    
    if st.button("Generate", type="primary"):
        with st.spinner("Generating text..."):
            result = synthesizer.synthesize(prompt, target_feature)
        
        if result:
            st.markdown("**Generated Text:**")
            st.markdown(f"> {result['text']}")
            
            st.divider()
            
            col1, col2 = st.columns(2)
            
            with col1:
                if result["success"]:
                    st.markdown('<div class="success-box">', unsafe_allow_html=True)
                    st.markdown("✅ **Target Activated**")
                    st.metric("Score", f"{result['score']:.3f}")
                    st.markdown('</div>', unsafe_allow_html=True)
                else:
                    st.markdown('<div class="error-box">', unsafe_allow_html=True)
                    st.markdown("❌ **Target Not Activated**")
                    st.metric("Score", f"{result['score']:.3f}")
                    st.markdown('</div>', unsafe_allow_html=True)
            
            with col2:
                st.markdown("**Other Activated Features:**")
                for feat_id, score in list(result["activations"].items())[:5]:
                    st.markdown(f"- Feature {feat_id}: {score:.3f}")

def fac_coverage_tab(analyzer, feature_db, task_type):
    st.subheader("📊 FAC Coverage")
    
    dataset_text = st.text_area(
        "Dataset Samples (one per line)",
        height=200,
        placeholder="Enter multiple text samples, one per line..."
    )
    
    if st.button("Compute Coverage", type="primary"):
        if dataset_text:
            samples = [line.strip() for line in dataset_text.split('\n') if line.strip()]
            
            with st.spinner(f"Analyzing {len(samples)} samples..."):
                all_features = set()
                for sample in samples:
                    results = analyzer.analyze(sample)
                    all_features.update(results["activations"].keys())
                
                total_features = feature_db.get_total_features(task_type)
                coverage = len(all_features) / total_features * 100
            
            col1, col2, col3 = st.columns(3)
            with col1:
                st.metric("Total Samples", len(samples))
            with col2:
                st.metric("Activated Features", len(all_features))
            with col3:
                st.metric("Coverage", f"{coverage:.1f}%")
            
            missing_features = set(range(total_features)) - all_features
            if missing_features:
                st.markdown("**Missing Features:**")
                st.markdown(
                    ", ".join([str(f) for f in sorted(list(missing_features)[:20])])
                    + (f" ... and {len(missing_features)-20} more" if len(missing_features) > 20 else "")
                )
            else:
                st.success("✅ All target features are covered!")

def batch_synthesis_tab(synthesizer, analyzer, feature_db, task_type):
    st.subheader("🚀 Batch Synthesis")
    
    st.markdown("""
    Generate multiple samples for missing features.
    First compute FAC coverage, then use this tab to fill gaps.
    """)
    
    missing_features_input = st.text_input(
        "Missing Feature IDs (comma-separated)",
        placeholder="1234, 5678, 9012"
    )
    
    samples_per_feature = st.slider(
        "Samples per Feature",
        min_value=1,
        max_value=10,
        value=3
    )
    
    if st.button("Generate Batch", type="primary"):
        if missing_features_input:
            feature_ids = [int(f.strip()) for f in missing_features_input.split(",")]
            
            with st.spinner(f"Generating {len(feature_ids) * samples_per_feature} samples..."):
                progress = st.progress(0)
                results = []
                
                for i, feat_id in enumerate(feature_ids):
                    for j in range(samples_per_feature):
                        result = synthesizer.synthesize(
                            f"Generate example for feature {feat_id}",
                            feat_id
                        )
                        results.append(result)
                        progress.progress((i * samples_per_feature + j + 1) / (len(feature_ids) * samples_per_feature))
            
            st.success(f"✅ Generated {len(results)} samples")
            
            for i, result in enumerate(results[:10]):
                with st.expander(f"Sample {i+1} (Feature {result.get('target_feature', 'N/A')})"):
                    st.markdown(result['text'])
                    st.caption(f"Target activated: {'Yes' if result['success'] else 'No'}")

def tutorial_tab():
    st.subheader("📚 Interactive Tutorial")
    
    st.markdown("""
    ### Welcome to FAC-Synthesis Demo!
    
    This tool helps you analyze and synthesize data using Sparse Autoencoder (SAE) features.
    
    #### Getting Started
    
    1. **Configure SAE Weights** (Sidebar)
       - Choose Default, Hugging Face, or Upload your SAE checkpoint
       - Adjust activation threshold (1.0 recommended for start)
    
    2. **Select Data Synthesizer** (Sidebar)
       - Choose from LLaMA, Mistral, Qwen, or GPT-4o-mini
       - For GPT-4o-mini, provide your OpenAI API key
    
    3. **Try Quick Examples** (Sidebar)
       - Click any example to quickly test the system
       - Examples cover different task types
    
    #### Feature Analysis
    
    Analyze any text to see which SAE features are activated:
    - Input your text
    - View activation heatmap
    - See feature explanations
    - Identify highlighted spans
    
    #### Targeted Synthesis
    
    Generate text that activates specific features:
    - Enter target feature ID
    - Provide a prompt template
    - System generates text and verifies activation
    
    #### FAC Coverage
    
    Compute Feature Activation Coverage for datasets:
    - Input multiple samples
    - See which features are covered
    - Identify missing features
    
    #### Batch Synthesis
    
    Fill data gaps efficiently:
    - Provide missing feature IDs
    - Generate multiple samples per feature
    - Export results for training
    
    #### Advanced Features
    
    - Load custom datasets from Hugging Face
    - Use your own SAE checkpoints
    - Experiment with different thresholds
    """)
    
    st.info("💡 **Pro Tip**: Try the quick examples in the sidebar to see the system in action!")

if __name__ == "__main__":
    main()