import streamlit as st
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import pandas as pd
import sys, os
from pathlib import Path

# Backend API
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from api_client import WarehouseAPI

# Add project root for AI modules
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

# Import AI modules
try:
    from forecasting.smart_suggestions import get_smart_suggestions
    from forecasting.forecasting_core import generate_forecasts
    from utils.suggestions_panel import display_smart_suggestions
    imports_successful = True
except ImportError as e:
    imports_successful = False
    import_error = str(e)

st.set_page_config(
    page_title="AI Warehouse Management",
    page_icon="🏭",
    layout="wide"
)

# ================================
# DATA LOADING
# ================================
@st.cache_data(show_spinner=False)
def load_data():
    api = WarehouseAPI()
    try:
        inventory = api.get_inventory()
        orders = api.get_orders()
        order_items = api.get_order_items()
        product = api.get_products()
        spare_parts = api.get_spare_parts()
        warehouses = api.get_warehouses()

        # Ensure datetime conversion
        if "order_date" in orders.columns:
            orders["order_date"] = pd.to_datetime(orders["order_date"], errors="coerce")

        if "estimated_delivery" in orders.columns:
            orders["estimated_delivery"] = pd.to_datetime(
                orders["estimated_delivery"], errors="coerce"
            )

        if "order_date" in order_items.columns:
            order_items["order_date"] = pd.to_datetime(
                order_items["order_date"], errors="coerce"
            )

        return inventory, orders, order_items, product, spare_parts, warehouses

    except Exception as e:
        st.error(f"API Error: {e}")
        return [pd.DataFrame()] * 6


# Load data
inventory, orders, order_items, product, spare_parts, warehouses = load_data()

# Store in session for reuse
st.session_state["inventory"] = inventory
st.session_state["orders"] = orders
st.session_state["order_items"] = order_items
st.session_state["product"] = product
st.session_state["spare_parts"] = spare_parts
st.session_state["warehouses"] = warehouses

# Create merged order_data for AI
if not orders.empty and not order_items.empty:
    duplicate_cols = [
        c for c in order_items.columns if c in orders.columns and c != "order_id"
    ]
    order_items_clean = order_items.drop(columns=duplicate_cols, errors="ignore")

    order_data = order_items_clean.merge(
        orders.drop_duplicates(subset=["order_id"]),
        how="left",
        on="order_id",
    )
    st.session_state["order_data"] = order_data
else:
    st.session_state["order_data"] = pd.DataFrame()


# ================================
# ENHANCED FORECAST PANEL
# ================================
st.header("🔮 Forecasting")
st.divider()

if not order_items.empty:

    # User selections
    col1, col2 = st.columns([2, 1])
    with col1:
        available_skus = sorted(order_items["sku"].unique().tolist())
        selected_sku = st.selectbox("Select SKU to Forecast", available_skus)
    
    with col2:
        horizon = st.slider("Forecast Horizon (days)", 7, 60, 17)

    model_options = ["SARIMA", "ETS", "PROPHET"]
    selected_models = st.multiselect(
        "Select Forecasting Models",
        model_options,
        default=["SARIMA", "ETS"]
    )
    
    # Phase 2: Add option for seasonality decomposition
    include_decomposition = st.checkbox("📊 Include Seasonality Analysis", value=True, 
                                        help="Break down demand into trend, seasonal, and random components")

    # Run forecast
    if st.button("Run Forecast", type="primary"):
        with st.spinner(f"Running {', '.join(selected_models)} forecast(s)..."):

            forecasts, conf_intervals, metrics, decompositions = generate_forecasts(
                order_items_df=order_items,
                sku_list=[selected_sku],
                horizon=horizon,
                models=[m.lower() for m in selected_models],
                include_decomposition=include_decomposition
            )

        if selected_sku not in forecasts:
            st.warning("Not enough data to generate a forecast for this SKU.")
        else:
            st.success("Forecast completed!")
            
            # Store in session state for persistence
            st.session_state['forecasts'] = forecasts
            st.session_state['conf_intervals'] = conf_intervals
            st.session_state['metrics'] = metrics
            st.session_state['decompositions'] = decompositions
            st.session_state['selected_sku'] = selected_sku
            st.session_state['include_decomposition'] = include_decomposition

    # Display results if available
    if 'forecasts' in st.session_state and st.session_state['selected_sku'] == selected_sku:
        forecasts = st.session_state['forecasts']
        conf_intervals = st.session_state['conf_intervals']
        metrics = st.session_state['metrics']
        decompositions = st.session_state.get('decompositions', {})
        
        sku_forecasts = forecasts[selected_sku]
        sku_conf_intervals = conf_intervals[selected_sku]
        sku_metrics = metrics[selected_sku]
        
        # ============================
        # PHASE 1.1: MODEL PERFORMANCE METRICS PANEL
        # ============================
        st.markdown("---")
        st.subheader("📊 Model Performance Comparison")
        
        # Create metrics dataframe
        metrics_data = []
        for model_name in sku_forecasts.keys():
            model_metrics = sku_metrics.get(model_name, {})
            metrics_data.append({
                "Model": model_name,
                "MAE": model_metrics.get("MAE", "N/A"),
                "RMSE": model_metrics.get("RMSE", "N/A"),
                "MAPE": model_metrics.get("MAPE", "N/A")
            })
        
        metrics_df = pd.DataFrame(metrics_data)
        
        # Find best model (lowest MAE)
        valid_mae = [(row["Model"], row["MAE"]) for _, row in metrics_df.iterrows() if row["MAE"] != "N/A"]
        if valid_mae:
            best_model = min(valid_mae, key=lambda x: x[1])[0]
        else:
            best_model = list(sku_forecasts.keys())[0]
        
        # Display metrics in columns
        metric_cols = st.columns(len(metrics_data))
        for idx, (col, row) in enumerate(zip(metric_cols, metrics_data)):
            with col:
                is_best = row["Model"] == best_model
                border_color = "#28a745" if is_best else "#6c757d"
                
                st.markdown(f"""
                <div style="border: 2px solid {border_color}; border-radius: 10px; padding: 15px; background-color: #1e1e1e;">
                    <h4 style="color: {border_color}; margin: 0;">
                        {row["Model"]} {"✓" if is_best else ""}
                    </h4>
                    <hr style="margin: 10px 0; border-color: {border_color};">
                    <p style="margin: 5px 0;"><strong>MAE:</strong> {row["MAE"]}</p>
                    <p style="margin: 5px 0;"><strong>RMSE:</strong> {row["RMSE"]}</p>
                    <p style="margin: 5px 0;"><strong>MAPE:</strong> {f'{row["MAPE"]}%' if row["MAPE"] != "N/A" else "N/A"}</p>
                    {"<p style='color: #28a745; font-weight: bold; margin-top: 10px;'>✓ Best Model</p>" if is_best else ""}
                </div>
                """, unsafe_allow_html=True)
        
        st.markdown(f"**Recommendation:** {best_model} shows the best accuracy for this SKU (lowest MAE).")
        
        # ============================
        # PHASE 1.3: BUSINESS ACTION PANEL
        # ============================
        if "BUSINESS" in sku_metrics:
            st.markdown("---")
            st.subheader("💼 Recommended Actions")
            
            biz_metrics = sku_metrics["BUSINESS"]
            
            action_col1, action_col2, action_col3 = st.columns(3)
            
            with action_col1:
                st.metric("Reorder Point", f"{biz_metrics['reorder_point']} units")
                st.metric("Safety Stock", f"{biz_metrics['safety_stock']} units")
            
            with action_col2:
                st.metric("Suggested Order Qty", f"{biz_metrics['suggested_order_qty']} units")
                st.metric("Avg Daily Demand", f"{biz_metrics['avg_daily_demand']} units/day")
            
            with action_col3:
                risk_color = {"LOW": "#28a745", "MEDIUM": "#ffc107", "HIGH": "#dc3545"}[biz_metrics['stockout_risk']]
                st.markdown(f"**Stockout Risk:** <span style='color: {risk_color}; font-size: 20px; font-weight: bold;'>{biz_metrics['stockout_risk']}</span>", unsafe_allow_html=True)
                
                # Display uncertainty with capped value and explanation
                uncertainty = min(biz_metrics['forecast_uncertainty'], 100)
                uncertainty_label = "Very High" if biz_metrics['forecast_uncertainty'] > 100 else f"{uncertainty:.1f}%"
                st.metric("Demand Variability", uncertainty_label, help="How unpredictable this SKU's demand is. Higher = more volatile.")
            
            # Action recommendations
            st.info(f"""
            **📋 Action Plan:**
            - Order {biz_metrics['suggested_order_qty']} units when inventory drops to {biz_metrics['reorder_point']} units
            - Maintain safety stock of {biz_metrics['safety_stock']} units to handle demand variability
            - Expected stockout risk: **{biz_metrics['stockout_risk']}**
            """)
        
        # ============================
        # PHASE 2.1: SEASONALITY DECOMPOSITION
        # ============================
        if st.session_state.get('include_decomposition', False) and selected_sku in decompositions:
            st.markdown("---")
            st.subheader("📊 Demand Pattern Analysis")
            
            decomp = decompositions[selected_sku]
            
            st.write("Understanding what drives demand changes:")
            
            # Create 4 subplots for decomposition
            fig = make_subplots(
                rows=4, cols=1,
                subplot_titles=('📈 Observed Demand', '📉 Trend Component', 
                              '🔄 Seasonal Pattern (Weekly)', '🎲 Random Variation'),
                vertical_spacing=0.08,
                row_heights=[0.25, 0.25, 0.25, 0.25]
            )
            
            # Observed data
            fig.add_trace(
                go.Scatter(
                    x=decomp['observed'].index, 
                    y=decomp['observed'].values,
                    mode='lines', 
                    name='Observed', 
                    line=dict(color='#3498db', width=2),
                    hovertemplate='<b>Date:</b> %{x}<br><b>Demand:</b> %{y}<extra></extra>'
                ),
                row=1, col=1
            )
            
            # Trend
            fig.add_trace(
                go.Scatter(
                    x=decomp['trend'].index, 
                    y=decomp['trend'].values,
                    mode='lines', 
                    name='Trend', 
                    line=dict(color='#e74c3c', width=2),
                    hovertemplate='<b>Date:</b> %{x}<br><b>Trend:</b> %{y:.2f}<extra></extra>'
                ),
                row=2, col=1
            )
            
            # Seasonal
            fig.add_trace(
                go.Scatter(
                    x=decomp['seasonal'].index, 
                    y=decomp['seasonal'].values,
                    mode='lines', 
                    name='Seasonal', 
                    line=dict(color='#2ecc71', width=2),
                    hovertemplate='<b>Date:</b> %{x}<br><b>Effect:</b> %{y:.2f}<extra></extra>'
                ),
                row=3, col=1
            )
            
            # Residual
            fig.add_trace(
                go.Scatter(
                    x=decomp['residual'].index, 
                    y=decomp['residual'].values,
                    mode='lines', 
                    name='Residual', 
                    line=dict(color='#95a5a6', width=1),
                    hovertemplate='<b>Date:</b> %{x}<br><b>Residual:</b> %{y:.2f}<extra></extra>'
                ),
                row=4, col=1
            )
            
            fig.update_layout(
                height=900,
                template='plotly_dark',
                showlegend=False,
                title_text="Time Series Decomposition - Understanding Demand Patterns"
            )
            
            fig.update_xaxes(title_text="Date", row=4, col=1)
            fig.update_yaxes(title_text="Units", row=1, col=1)
            fig.update_yaxes(title_text="Units", row=2, col=1)
            fig.update_yaxes(title_text="Effect", row=3, col=1)
            fig.update_yaxes(title_text="Variance", row=4, col=1)
            
            st.plotly_chart(fig, use_container_width=True)
            
            # Interpretation guide
            col1, col2 = st.columns(2)
            
            with col1:
                st.info("""
                **📖 How to Read This Analysis:**
                
                **Observed:** Your actual demand data
                
                **Trend:** Long-term direction - Is demand growing, declining, or stable?
                
                **Seasonal:** Repeating weekly patterns - Do certain days have consistently higher/lower demand?
                
                **Residual:** Random noise that can't be explained by trend or seasonality
                """)
            
            with col2:
                # Calculate insights
                trend_direction = "increasing" if decomp['trend'].iloc[-1] > decomp['trend'].iloc[0] else "decreasing" if decomp['trend'].iloc[-1] < decomp['trend'].iloc[0] else "stable"
                seasonal_strength = decomp['seasonal'].std()
                residual_strength = decomp['residual'].std()
                
                st.success(f"""
                **🔍 Key Insights for This SKU:**
                
                **Trend:** Demand is **{trend_direction}** over time
                
                **Seasonality Strength:** {"Strong" if seasonal_strength > 0.5 else "Weak"} weekly pattern detected
                
                **Predictability:** {"High variance in residuals - external factors may be influencing demand" if residual_strength > 1 else "Low residual variance - demand is predictable"}
                """)
        
        # ============================
        # PHASE 1.2: CHARTS WITH CONFIDENCE INTERVALS
        # ============================
        st.markdown("---")
        st.subheader("📈 Forecast Visualizations")
        
        # Prepare historical data
        history_df = order_items[order_items["sku"] == selected_sku].copy()
        history_df = history_df.groupby(history_df["order_date"].dt.date)["qty"].sum()
        history_df.index = pd.to_datetime(history_df.index)
        
        # Display each model's forecast
        for model_name, forecast_series in sku_forecasts.items():
            st.subheader(f"📈 {model_name} Forecast")
            
            # Get confidence intervals
            conf_int = sku_conf_intervals.get(model_name, {})
            lower_bound = conf_int.get("lower")
            upper_bound = conf_int.get("upper")
            
            # Create Plotly figure
            fig = go.Figure()
            
            # Historical data
            fig.add_trace(go.Scatter(
                x=history_df.index,
                y=history_df.values,
                mode='lines',
                name='Historical Demand',
                line=dict(color='#3498db', width=2),
                hovertemplate='<b>Date:</b> %{x}<br><b>Demand:</b> %{y}<extra></extra>'
            ))
            
            # Forecast
            fig.add_trace(go.Scatter(
                x=forecast_series.index,
                y=forecast_series.values,
                mode='lines',
                name='Forecast',
                line=dict(color='#e74c3c', width=2, dash='dash'),
                hovertemplate='<b>Date:</b> %{x}<br><b>Forecast:</b> %{y:.2f}<extra></extra>'
            ))
            
            # Confidence intervals (shaded area)
            if lower_bound is not None and upper_bound is not None:
                # Upper bound
                fig.add_trace(go.Scatter(
                    x=upper_bound.index,
                    y=upper_bound.values,
                    mode='lines',
                    name='95% Upper Bound',
                    line=dict(width=0),
                    showlegend=False,
                    hoverinfo='skip'
                ))
                
                # Lower bound with fill
                fig.add_trace(go.Scatter(
                    x=lower_bound.index,
                    y=lower_bound.values,
                    mode='lines',
                    name='95% Confidence Interval',
                    line=dict(width=0),
                    fillcolor='rgba(231, 76, 60, 0.2)',
                    fill='tonexty',
                    hovertemplate='<b>Date:</b> %{x}<br><b>Lower:</b> %{y:.2f}<extra></extra>'
                ))
            
            # Update layout
            fig.update_layout(
                title=f"{model_name} Demand Forecast with 95% Confidence Interval",
                xaxis_title="Date",
                yaxis_title="Demand (units)",
                hovermode='x unified',
                template='plotly_dark',
                height=500,
                showlegend=True,
                legend=dict(
                    orientation="h",
                    yanchor="bottom",
                    y=1.02,
                    xanchor="right",
                    x=1
                )
            )
            
            st.plotly_chart(fig, use_container_width=True)
            
            # Display forecast data table
            with st.expander(f"📋 {model_name} Forecast Data"):
                forecast_df = pd.DataFrame({
                    "Date": forecast_series.index,
                    "Forecast": forecast_series.values.round(2),
                    "Lower Bound (95%)": lower_bound.values.round(2) if lower_bound is not None else None,
                    "Upper Bound (95%)": upper_bound.values.round(2) if upper_bound is not None else None
                })
                st.dataframe(forecast_df, hide_index=True, use_container_width=True)
                
                # Export button
                csv = forecast_df.to_csv(index=False)
                st.download_button(
                    label=f"📥 Download {model_name} Forecast CSV",
                    data=csv,
                    file_name=f"{selected_sku}_{model_name}_forecast.csv",
                    mime="text/csv"
                )

else:
    st.info("❗ Need order_items data to run forecasting tests.")