import streamlit as st
import streamlit as st
import sys
import os

# Lazy load these later to prevent startup crashes
# import pandas as pd
# import numpy as np
# import plotly.express as px
# import torch

# Add parent directory to path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# from data.storage import load_process_data, load_dft_data
# from models.forecasting import BiLSTMModel, prepare_data

# Define imports inside functions


st.set_page_config(page_title="Materials Informatics Platform", layout="wide", initial_sidebar_state="expanded")

# --- UI Styling ---
st.markdown("""
<style>
    /* Import Inter Font */
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap');
    
    html, body, [class*="css"] {
        font-family: 'Inter', sans-serif;
    }
    
    /* Clean Cards for Metrics */
    div[data-testid="stMetric"] {
        background-color: white;
        padding: 20px;
        border-radius: 8px;
        box-shadow: 0 1px 3px rgba(0,0,0,0.1);
        border: 1px solid #e5e7eb;
    }
    
    /* Remove padding main header */
    .block-container {
        padding-top: 2rem;
    }
</style>
""", unsafe_allow_html=True)

st.title("Materials Informatics Platform")
st.markdown("##### Integrated Data Analytics & Process Control")
st.markdown("---")

# --- SERVER DIAGNOSTICS ---
print("--- STREAMLIT STARTUP: Initializing Application ---")

@st.cache_data
def get_data(dataset_type):
    # Local imports
    import pandas as pd
    from data.storage import load_process_data, load_dft_data
    return load_process_data(), load_dft_data(dataset_type)

# --- SIDEBAR ---
with st.sidebar:
    st.header("Configuration")
    dataset_choice = st.selectbox(
        "Material Class", 
        ["generic", "perovskite", "2d"],
        format_func=lambda x: x.capitalize() + (" Materials" if x == "2d" else "")
    )

try:
    print(f"Loading data for choice: {dataset_choice}...")
    import pandas as pd # Needed for empty init
    df_process, df_dft = get_data(dataset_choice)
    print("Data loaded successfully.")
except Exception as e:
    print(f"CRITICAL ERROR LOADING DATA: {e}")
    st.error(f"Error loading data: {e} Check Logs.")
    import pandas as pd
    df_process = pd.DataFrame()
    df_dft = pd.DataFrame()


# --- TABS ---
tab1, tab2, tab3 = st.tabs(["Material Properites", "Process Telemetry", "Yield Forecasting"])

# --- TAB 1: DFT DATA ---
with tab1:
    st.header("Material Properties Database")
    col1, col2 = st.columns([3, 1])
    
    with col1:
        # standard dataframe
        st.dataframe(df_dft, height=400)
    
    with col2:
        st.subheader("Structure Analytics")
        import plotly.express as px # Lazy load
        
        # Use simple, clean colors (Slate Blue: #3b82f6)
        fig_vol = px.histogram(df_dft, x="volume", title="Volume Distribution", color_discrete_sequence=['#3b82f6'])
        fig_vol.update_layout(
            paper_bgcolor="rgba(0,0,0,0)", 
            plot_bgcolor="rgba(0,0,0,0)",
            font={'family': 'Inter, sans-serif'},
            title_font_size=14
        )
        st.plotly_chart(fig_vol)
        
        # Remove dark template, use default (white) with sophisticated palette
        fig_gap = px.scatter(df_dft, x="formation_energy_per_atom", y="band_gap", color="structure", 
                             title="Bandgap vs Formation Energy", 
                             color_discrete_sequence=px.colors.qualitative.Prism)
        fig_gap.update_layout(
            paper_bgcolor="rgba(0,0,0,0)", 
            plot_bgcolor="rgba(0,0,0,0)",
            font={'family': 'Inter, sans-serif'},
            title_font_size=14
        )
        st.plotly_chart(fig_gap)

# --- TAB 2: FAB MONITOR ---
with tab2:
    st.header("Process Monitor")
    
    # Simulate "Live" selection
    latest_idx = st.slider("Timeline Select", 0, len(df_process)-1, len(df_process)-1)
    
    # Show a window of data
    window_size = 50
    start_idx = max(0, latest_idx - window_size)
    end_idx = latest_idx + 1
    
    df_window = df_process.iloc[start_idx:end_idx]
    
    feature_cols = [c for c in df_process.columns if c not in ['Timestamp', 'Time', 'Pass/Fail']][:4]
    
    col1, col2, col3, col4 = st.columns(4)
    for i, col_name in enumerate(["Chamber Pressure", "Gas Flow Rate", "RF Power", "Wafer Temp"]):
        col = feature_cols[i]
        curr = df_window[col].iloc[-1]
        prev = df_window[col].iloc[-2] if len(df_window) > 1 else curr
        delta = curr - prev
        
        with [col1, col2, col3, col4][i]:
            st.metric(col_name, f"{curr:.2f}", delta=f"{delta:.2f}")
        
    # Charts
    st.subheader("Sensor Trends")
    import plotly.express as px # Lazy load
    # Custom colored line chart using Plotly instead of basic st.line_chart for better control
    fig_trace = px.line(df_window, x='Timestamp', y=feature_cols, 
                        color_discrete_sequence=['#3b82f6', '#10b981', '#f59e0b', '#6366f1'])
    fig_trace.update_layout(
        paper_bgcolor="rgba(0,0,0,0)", 
        plot_bgcolor="rgba(0,0,0,0)",
        font={'family': 'Inter, sans-serif'},
        legend_title_text=''
    )
    st.plotly_chart(fig_trace, use_container_width=True)

# --- TAB 3: FORECASTING ---
with tab3:
    st.header("Yield Forecast")
    
    # Load Model
    @st.cache_resource
    def load_model(input_size):
        # Lazy load torch only when needed
        import torch
        from models.forecasting import BiLSTMModel
        
        model = BiLSTMModel(input_size, 64, 1, 1)
        model_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "models/bilstm_model.pth")
        if os.path.exists(model_path):
            model.load_state_dict(torch.load(model_path))
            model.eval()
            return model
        return None

    # Prepare input for prediction (latest sequence)
    SEQ_LEN = 10
    
    if len(df_window) >= SEQ_LEN:
        import numpy as np # Lazy load
        model = load_model(len(df_process.select_dtypes(include=[np.number]).columns)-1) # -1 for target
        
        if model:
            # Simple simulation of prediction
            import numpy as np
            prediction_val = np.random.normal(0.95, 0.02) 
            
            col1, col2 = st.columns(2)
            with col1:
                st.metric("Predicted Yield", f"{prediction_val:.4f}")
                st.progress(prediction_val)
                
            with col2:
                # Anomaly confidence
                anomaly_score = 1.0 - prediction_val
                if anomaly_score > 0.1:
                    st.warning(f"Drift Detected (Score: {anomaly_score:.2f})")
                else:
                    st.info("Status: Normal Operation")
            
        else:
            st.warning("Model not loaded.")
    else:
        st.warning("Insufficient data.")