import streamlit as st
import pandas as pd
import re
import plotly.express as px
import plotly.graph_objects as go
from typing import Optional, Tuple

# ✅ MUST be the first Streamlit command
st.set_page_config(
    page_title="LLM Compatibility Advisor", 
    layout="wide",
    page_icon="🧠",
    initial_sidebar_state="expanded"
)

# Enhanced data loading with error handling
@st.cache_data
def load_data():
    try:
        df = pd.read_excel("ICFAI.xlsx", sheet_name="Form Responses 1")
        df.columns = df.columns.str.strip()
        return df, None
    except FileNotFoundError:
        return None, "Excel file 'ICFAI.xlsx' not found. Please upload the file."
    except Exception as e:
        return None, f"Error loading data: {str(e)}"

# Enhanced RAM extraction with better parsing
def extract_numeric_ram(ram) -> Optional[int]:
    if pd.isna(ram):
        return None
    
    ram_str = str(ram).lower().replace(" ", "")
    
    # Handle various formats: "8GB", "8 GB", "8gb", "8192MB", etc.
    gb_match = re.search(r"(\d+(?:\.\d+)?)(?:gb|g)", ram_str)
    if gb_match:
        return int(float(gb_match.group(1)))
    
    # Handle MB format
    mb_match = re.search(r"(\d+)(?:mb|m)", ram_str)
    if mb_match:
        return max(1, int(int(mb_match.group(1)) / 1024))  # Convert MB to GB
    
    # Handle plain numbers (assume GB)
    plain_match = re.search(r"(\d+)", ram_str)
    if plain_match:
        return int(plain_match.group(1))
    
    return None

# Enhanced LLM recommendation with performance tiers
def recommend_llm(ram_str) -> Tuple[str, str, str]:
    """Returns (recommendation, performance_tier, additional_info)"""
    ram = extract_numeric_ram(ram_str)
    
    if ram is None:
        return "⚪ Check exact specs or test with quantized models.", "Unknown", "Verify RAM specifications"
    
    if ram <= 2:
        return ("🔸 DistilBERT, MobileBERT, TinyBERT", 
                "Basic", 
                "Suitable for simple NLP tasks, limited context")
    elif ram <= 4:
        return ("🔸 MiniLM, TinyLLaMA, DistilRoBERTa", 
                "Basic", 
                "Good for text classification, basic chat")
    elif ram <= 6:
        return ("🟠 Phi-1.5, Alpaca.cpp (3B), Mistral Tiny", 
                "Moderate", 
                "Decent reasoning, short conversations")
    elif ram <= 8:
        return ("🟠 Phi-2, Gemma 2B, LLaMA 2 7B (4-bit)", 
                "Moderate", 
                "Good general purpose, coding assistance")
    elif ram <= 12:
        return ("🟢 LLaMA 2 7B (GGUF), Mistral 7B (int4)", 
                "Good", 
                "Strong performance, longer contexts")
    elif ram <= 16:
        return ("🟢 Mixtral 8x7B (4-bit), Command R+", 
                "Good", 
                "Excellent reasoning, complex tasks")
    elif ram <= 24:
        return ("🔵 LLaMA 2 13B, Mistral 7B FP16, Gemma 7B", 
                "High", 
                "Professional grade, high accuracy")
    else:
        return ("🔵 Mixtral 8x7B Full, LLaMA 2 70B (quantized)", 
                "High", 
                "Top-tier performance, enterprise ready")

# Enhanced OS detection with better icons
def get_os_info(os_name) -> Tuple[str, str]:
    """Returns (icon, clean_name)"""
    if pd.isna(os_name):
        return "💻", "Not specified"
    
    os = str(os_name).lower()
    if "windows" in os:
        return "🪟", os_name
    elif "mac" in os or "darwin" in os:
        return "🍎", os_name
    elif "linux" in os or "ubuntu" in os:
        return "🐧", os_name
    elif "android" in os:
        return "🤖", os_name
    elif "ios" in os:
        return "📱", os_name
    else:
        return "💻", os_name

# Performance visualization
def create_performance_chart(df):
    """Create a performance distribution chart"""
    laptop_rams = df["Laptop RAM"].apply(extract_numeric_ram).dropna()
    mobile_rams = df["Mobile RAM"].apply(extract_numeric_ram).dropna()
    
    fig = go.Figure()
    
    fig.add_trace(go.Histogram(
        x=laptop_rams,
        name="Laptop RAM",
        opacity=0.7,
        nbinsx=10
    ))
    
    fig.add_trace(go.Histogram(
        x=mobile_rams,
        name="Mobile RAM",
        opacity=0.7,
        nbinsx=10
    ))
    
    fig.update_layout(
        title="RAM Distribution Across Devices",
        xaxis_title="RAM (GB)",
        yaxis_title="Number of Students",
        barmode='overlay',
        height=400
    )
    
    return fig

# Main App
st.title("🧠 LLM Compatibility Advisor")
st.markdown("Get personalized, device-based suggestions for running LLMs efficiently!")

# Load data
df, error = load_data()

if error:
    st.error(error)
    st.info("Please ensure the Excel file 'BITS_INTERNS.xlsx' is in the same directory as this script.")
    st.stop()

if df is None or df.empty:
    st.error("No data found in the Excel file.")
    st.stop()

# Sidebar filters and info
with st.sidebar:
    st.header("🔍 Filters & Info")
    
    # Performance tier filter
    performance_filter = st.multiselect(
        "Filter by Performance Tier:",
        ["Basic", "Moderate", "Good", "High", "Unknown"],
        default=["Basic", "Moderate", "Good", "High", "Unknown"]
    )
    
    # RAM range filter
    st.subheader("RAM Range Filter")
    min_ram = st.slider("Minimum RAM (GB)", 0, 32, 0)
    max_ram = st.slider("Maximum RAM (GB)", 0, 64, 64)
    
    st.markdown("---")
    st.markdown("### 📊 Quick Stats")
    st.metric("Total Students", len(df))
    
    # Calculate average RAM
    avg_laptop_ram = df["Laptop RAM"].apply(extract_numeric_ram).mean()
    avg_mobile_ram = df["Mobile RAM"].apply(extract_numeric_ram).mean()
    
    if not pd.isna(avg_laptop_ram):
        st.metric("Avg Laptop RAM", f"{avg_laptop_ram:.1f} GB")
    if not pd.isna(avg_mobile_ram):
        st.metric("Avg Mobile RAM", f"{avg_mobile_ram:.1f} GB")

# User selection with search
st.subheader("👤 Individual Student Analysis")
selected_user = st.selectbox(
    "Choose a student:",
    options=[""] + list(df["Full Name"].unique()),
    format_func=lambda x: "Select a student..." if x == "" else x
)

if selected_user:
    user_data = df[df["Full Name"] == selected_user].iloc[0]
    
    # Enhanced user display
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("### 💻 Laptop Configuration")
        laptop_os_icon, laptop_os_name = get_os_info(user_data.get('Laptop Operating System'))
        laptop_ram = user_data.get('Laptop RAM', 'Not specified')
        laptop_rec, laptop_tier, laptop_info = recommend_llm(laptop_ram)
        
        st.markdown(f"**OS:** {laptop_os_icon} {laptop_os_name}")
        st.markdown(f"**RAM:** {laptop_ram}")
        st.markdown(f"**Performance Tier:** {laptop_tier}")
        
        st.success(f"**💡 Recommendation:** {laptop_rec}")
        st.info(f"**ℹ️ Notes:** {laptop_info}")
    
    with col2:
        st.markdown("### 📱 Mobile Configuration")
        mobile_os_icon, mobile_os_name = get_os_info(user_data.get('Mobile Operating System'))
        mobile_ram = user_data.get('Mobile RAM', 'Not specified')
        mobile_rec, mobile_tier, mobile_info = recommend_llm(mobile_ram)
        
        st.markdown(f"**OS:** {mobile_os_icon} {mobile_os_name}")
        st.markdown(f"**RAM:** {mobile_ram}")
        st.markdown(f"**Performance Tier:** {mobile_tier}")
        
        st.success(f"**💡 Recommendation:** {mobile_rec}")
        st.info(f"**ℹ️ Notes:** {mobile_info}")

# Batch Analysis Section
st.markdown("---")
st.header("📊 Batch Analysis & Insights")

# Create enhanced batch table
df_display = df[["Full Name", "Laptop RAM", "Mobile RAM"]].copy()

# Add recommendations and performance tiers
laptop_recommendations = df["Laptop RAM"].apply(lambda x: recommend_llm(x)[0])
mobile_recommendations = df["Mobile RAM"].apply(lambda x: recommend_llm(x)[0])
laptop_tiers = df["Laptop RAM"].apply(lambda x: recommend_llm(x)[1])
mobile_tiers = df["Mobile RAM"].apply(lambda x: recommend_llm(x)[1])

df_display["Laptop LLM"] = laptop_recommendations
df_display["Mobile LLM"] = mobile_recommendations
df_display["Laptop Tier"] = laptop_tiers
df_display["Mobile Tier"] = mobile_tiers

# Filter based on sidebar selections
laptop_ram_numeric = df["Laptop RAM"].apply(extract_numeric_ram)
mobile_ram_numeric = df["Mobile RAM"].apply(extract_numeric_ram)

# Apply filters
mask = (
    (laptop_tiers.isin(performance_filter) | mobile_tiers.isin(performance_filter)) &
    ((laptop_ram_numeric.between(min_ram, max_ram)) | (mobile_ram_numeric.between(min_ram, max_ram)))
)

df_filtered = df_display[mask]

# Display filtered table
st.subheader(f"📋 Student Recommendations ({len(df_filtered)} students)")
st.dataframe(
    df_filtered, 
    use_container_width=True,
    column_config={
        "Full Name": st.column_config.TextColumn("Student Name", width="medium"),
        "Laptop RAM": st.column_config.TextColumn("Laptop RAM", width="small"),
        "Mobile RAM": st.column_config.TextColumn("Mobile RAM", width="small"),
        "Laptop LLM": st.column_config.TextColumn("Laptop Recommendation", width="large"),
        "Mobile LLM": st.column_config.TextColumn("Mobile Recommendation", width="large"),
        "Laptop Tier": st.column_config.TextColumn("L-Tier", width="small"),
        "Mobile Tier": st.column_config.TextColumn("M-Tier", width="small"),
    }
)

# Performance distribution chart
if len(df) > 1:
    st.subheader("📈 RAM Distribution Analysis")
    fig = create_performance_chart(df)
    st.plotly_chart(fig, use_container_width=True)

# Performance tier summary
st.subheader("🎯 Performance Tier Summary")
tier_col1, tier_col2 = st.columns(2)

with tier_col1:
    st.markdown("**Laptop Performance Tiers:**")
    laptop_tier_counts = laptop_tiers.value_counts()
    for tier, count in laptop_tier_counts.items():
        percentage = (count / len(laptop_tiers)) * 100
        st.write(f"• {tier}: {count} students ({percentage:.1f}%)")

with tier_col2:
    st.markdown("**Mobile Performance Tiers:**")
    mobile_tier_counts = mobile_tiers.value_counts()
    for tier, count in mobile_tier_counts.items():
        percentage = (count / len(mobile_tiers)) * 100
        st.write(f"• {tier}: {count} students ({percentage:.1f}%)")

# Enhanced reference table
with st.expander("📘 Comprehensive LLM Reference Guide"):
    st.markdown("""
    ## RAM-to-LLM Compatibility Matrix
    
    | RAM Size | Performance Tier | Recommended Models | Use Cases | Additional Notes |
    |----------|------------------|-------------------|-----------|------------------|
    | ≤2GB | 🔸 Basic | DistilBERT, MobileBERT, TinyBERT | Simple NLP, text classification | Limited context, fast inference |
    | 4GB | 🔸 Basic | MiniLM, TinyLLaMA, DistilRoBERTa | Basic chat, QA systems | Good for mobile deployment |
    | 6GB | 🟠 Moderate | Phi-1.5, Alpaca.cpp (3B), Mistral Tiny | Reasoning tasks, short conversations | Balanced performance/memory |
    | 8GB | 🟠 Moderate | Phi-2, Gemma 2B, LLaMA 2 7B (4-bit) | General purpose, coding help | Popular sweet spot |
    | 12GB | 🟢 Good | LLaMA 2 7B (GGUF), Mistral 7B (int4) | Professional tasks, longer context | Production ready |
    | 16GB | 🟢 Good | Mixtral 8x7B (4-bit), Command R+ | Complex reasoning, analysis | Excellent capabilities |
    | 24GB | 🔵 High | LLaMA 2 13B, Mistral 7B FP16, Gemma 7B | High-accuracy tasks, research | Professional grade |
    | >24GB | 🔵 High | Mixtral 8x7B Full, LLaMA 2 70B (quantized) | Enterprise applications | Top-tier performance |
    
    ### 🔧 Optimization Tips:
    - **Quantization**: Use 4-bit or 8-bit quantization to reduce memory usage
    - **GGUF Format**: Optimized format for CPU inference with lower memory overhead
    - **Context Length**: Longer contexts require significantly more memory
    - **Batch Size**: Reduce batch size for inference to save memory
    """)

# Footer with additional resources
st.markdown("---")
st.markdown("""
### 🔗 Additional Resources
- **Quantization Tools**: GPTQ, GGML, bitsandbytes
- **Inference Engines**: llama.cpp, vLLM, TensorRT-LLM
- **Model Repositories**: Hugging Face, Ollama, LM Studio
""")