import streamlit as st
import pandas as pd
import numpy as np
import plotly.express as px
from datetime import datetime, timedelta
from io import StringIO
import os
import json

# Debug: Verify file paths
st.write("Debug: Checking file paths...")
files_to_check = ["new_best_improved_model.pth", "scaler.pkl", "feature_names.json", "model_config.json"]
for file in files_to_check:
    st.write(f"{file}: {'Found' if os.path.exists(file) else 'Missing'}")

try:
    from inference import load_model_and_artifacts, predict
except Exception as e:
    st.error(f"Error importing inference: {str(e)}")
    st.stop()

st.title("Store Sales Time Series Forecasting")
st.markdown("Forecast 13-week store sales using an LSTM model trained on Kaggle Store Sales data.")

# Load model and artifacts
try:
    st.write("Debug: Loading model and artifacts...")
    model, scaler, feature_names, config = load_model_and_artifacts()
    st.success("Model and artifacts loaded successfully")
except Exception as e:
    st.error(f"Error loading model or artifacts: {str(e)}")
    st.stop()

# Display model metrics
st.header("Model Performance Metrics")
metrics = {
    "MAE": 710.75,
    "RMSE": 1108.36,
    "MAPE": 7.16,
    "R2": 0.8633
}
st.markdown(f"""
- **MAE**: ${metrics['MAE']:.2f}  
- **RMSE**: ${metrics['RMSE']:.2f}  
- **MAPE**: {metrics['MAPE']:.2f}%  
- **R² Score**: {metrics['R2']:.4f}
""")

# Model architecture summary
st.header("Model Architecture")
st.markdown(f"""
- **Input Size**: {config['input_size']} features  
- **Hidden Size**: {config['hidden_size']}  
- **Number of Layers**: {config['num_layers']}  
- **Forecast Horizon**: {config['forecast_horizon']} weeks  
- **Dropout**: {config['dropout']}  
- **Attention**: {config['has_attention']}  
- **Input Projection**: {config['has_input_projection']}  
- **Parameters**: 227,441
""")

# Synthetic data generation
st.header("Generate Synthetic Test Data")
st.markdown("Create a sample dataset with 21 timesteps matching the training data distribution (sales ~$3,000–19,000).")
if st.button("Generate Synthetic Data"):
    np.random.seed(42)
    sequence_length = 21
    n_features = len(feature_names)
    synthetic_data = np.zeros((sequence_length, n_features))
    
    # Generate features based on training data characteristics
    for i, feature in enumerate(feature_names):
        if feature == "sales":
            synthetic_data[:, i] = np.random.normal(8954.97, 3307.49, sequence_length)  # Mean, std from verbose
        elif feature == "onpromotion":
            synthetic_data[:, i] = np.random.choice([0, 1], sequence_length, p=[0.8, 0.2])
        elif feature in ["dayofweek_sin", "dayofweek_cos"]:
            synthetic_data[:, i] = np.sin(np.linspace(0, 2 * np.pi, sequence_length)) if "sin" in feature else np.cos(np.linspace(0, 2 * np.pi, sequence_length))
        elif feature in ["month_sin", "month_cos"]:
            synthetic_data[:, i] = np.sin(np.linspace(0, 2 * np.pi * 12 / sequence_length, sequence_length)) if "sin" in feature else np.cos(np.linspace(0, 2 * np.pi * 12 / sequence_length, sequence_length))
        elif feature == "trend":
            synthetic_data[:, i] = np.linspace(0, sequence_length, sequence_length)
        elif feature == "is_weekend":
            synthetic_data[:, i] = np.random.choice([0, 1], sequence_length, p=[0.7, 0.3])
        elif feature == "quarter":
            synthetic_data[:, i] = np.random.choice([1, 2, 3, 4], sequence_length)
        elif "lag" in feature:
            lag = int(feature.split('_')[-1])
            synthetic_data[:, i] = np.roll(synthetic_data[:, 0], lag)
            if lag > 0:
                synthetic_data[:lag, i] = synthetic_data[:lag, 0]
        elif "ma" in feature:
            window = int(feature.split('_')[-1])
            synthetic_data[:, i] = pd.Series(synthetic_data[:, 0]).rolling(window=window, min_periods=1).mean().values
        elif "ratio" in feature:
            window = int(feature.split('_')[-1])
            ma = pd.Series(synthetic_data[:, 0]).rolling(window=window, min_periods=1).mean().values
            synthetic_data[:, i] = synthetic_data[:, 0] / (ma + 1e-8)
        elif "promo" in feature:
            synthetic_data[:, i] = np.random.choice([0, 1], sequence_length, p=[0.8, 0.2])
        elif feature == "dcoilwtico":
            synthetic_data[:, i] = np.random.normal(80, 10, sequence_length)
        elif feature == "is_holiday":
            synthetic_data[:, i] = np.random.choice([0, 1], sequence_length, p=[0.9, 0.1])
    
    # Create DataFrame with dates
    synthetic_df = pd.DataFrame(synthetic_data, columns=feature_names)
    end_date = datetime.now().date()
    dates = [end_date - timedelta(days=x) for x in range(sequence_length-1, -1, -1)]
    synthetic_df['Date'] = dates
    
    # Store in session state
    st.session_state["synthetic_df"] = synthetic_df
    
    st.subheader("Synthetic Data Preview")
    st.dataframe(synthetic_df.head())
    
    # Download synthetic data
    csv_buffer = StringIO()
    synthetic_df.to_csv(csv_buffer, index=False)
    st.download_button(
        label="Download Synthetic Data CSV",
        data=csv_buffer.getvalue(),
        file_name="synthetic_sales_data.csv",
        mime="text/csv"
    )
    
    # Generate forecast
    try:
        sequences = synthetic_df[feature_names].values.reshape(1, sequence_length, n_features)
        sequences_scaled = scaler.transform(sequences.reshape(-1, n_features)).reshape(1, sequence_length, n_features)
        predictions, uncertainties = predict(model, scaler, sequences_scaled)
        
        # Create forecast DataFrame
        forecast_dates = [end_date + timedelta(days=x*7) for x in range(1, 14)]
        forecast_df = pd.DataFrame({
            'Date': forecast_dates,
            'Predicted Sales ($)': predictions[0],
            'Uncertainty ($)': uncertainties[0]
        })
        
        st.subheader("13-Week Forecast")
        st.dataframe(forecast_df)
        
        # Plot forecast
        fig = px.line(forecast_df, x='Date', y='Predicted Sales ($)', title='13-Week Sales Forecast')
        fig.add_scatter(
            x=forecast_df['Date'],
            y=forecast_df['Predicted Sales ($)'] + forecast_df['Uncertainty ($)'],
            mode='lines', name='Upper Bound', line=dict(dash='dash', color='green')
        )
        fig.add_scatter(
            x=forecast_df['Date'],
            y=forecast_df['Predicted Sales ($)'] - forecast_df['Uncertainty ($)'],
            mode='lines', name='Lower Bound', line=dict(dash='dash', color='green'),
            fill='tonexty', fillcolor='rgba(0, 255, 0, 0.1)'
        )
        st.plotly_chart(fig)
    except Exception as e:
        st.error(f"Error generating forecast: {str(e)}")

# CSV upload for custom predictions
st.header("Upload Custom Data")
st.markdown("Upload a CSV with 21 timesteps and 20 features matching the feature names and format of the synthetic data.")
uploaded_file = st.file_uploader("Choose a CSV file", type="csv")

if uploaded_file is not None:
    try:
        data = pd.read_csv(uploaded_file)
        if set(feature_names).issubset(data.columns) and len(data) == 21:
            sequences = data[feature_names].values.reshape(1, 21, len(feature_names))
            sequences_scaled = scaler.transform(sequences.reshape(-1, len(feature_names))).reshape(1, 21, len(feature_names))
            predictions, uncertainties = predict(model, scaler, sequences_scaled)
            
            # Create forecast DataFrame
            forecast_df = pd.DataFrame({
                'Week': range(1, 14),
                'Predicted Sales ($)': predictions[0],
                'Uncertainty ($)': uncertainties[0]
            })
            
            st.subheader("13-Week Forecast")
            st.dataframe(forecast_df)
            
            # Plot forecast
            fig = px.line(forecast_df, x='Week', y='Predicted Sales ($)', title='13-Week Sales Forecast')
            fig.add_scatter(
                x=forecast_df['Week'],
                y=forecast_df['Predicted Sales ($)'] + forecast_df['Uncertainty ($)'],
                mode='lines', name='Upper Bound', line=dict(dash='dash', color='green')
            )
            fig.add_scatter(
                x=forecast_df['Week'],
                y=forecast_df['Predicted Sales ($)'] - forecast_df['Uncertainty ($)'],
                mode='lines', name='Lower Bound', line=dict(dash='dash', color='green'),
                fill='tonexty', fillcolor='rgba(0, 255, 0, 0.1)'
            )
            st.plotly_chart(fig)
        else:
            st.error(f"Invalid CSV. Expected 21 rows and columns including: {', '.join(feature_names)}")
    except Exception as e:
        st.error(f"Error processing CSV or generating forecast: {str(e)}")