import pandas as pd
from typing import Dict, List, Optional, Tuple
from .models import (
    exponential_smoothing_forecast,
    sarima_forecast,
    prophet_forecast,
    calculate_metrics
)
from statsmodels.tsa.seasonal import seasonal_decompose

def prepare_daily_series(order_items_df: pd.DataFrame, sku: str) -> pd.Series:
    """
    Create a daily demand time series for a given SKU.
    """
    subset = order_items_df[order_items_df["sku"] == sku]
    if subset.empty:
        return pd.Series(dtype=float)
    series = subset.groupby(subset["order_date"].dt.date)["qty"].sum()
    series.index = pd.to_datetime(series.index)
    return series.asfreq("D").fillna(0.0)

def decompose_seasonality(series: pd.Series, period: int = 7) -> Dict[str, pd.Series]:
    """
    Decompose time series into trend, seasonal, and residual components.
    
    Args:
        series: Time series data
        period: Seasonality period (7 for weekly, 30 for monthly)
    
    Returns:
        Dict with 'trend', 'seasonal', 'residual' components
    """
    if len(series) < 2 * period:
        return None
    
    try:
        # Perform seasonal decomposition
        decomposition = seasonal_decompose(
            series, 
            model='additive',
            period=period,
            extrapolate_trend='freq'
        )
        
        return {
            'trend': decomposition.trend,
            'seasonal': decomposition.seasonal,
            'residual': decomposition.resid,
            'observed': series
        }
    except Exception as e:
        print(f"Decomposition error: {e}")
        return None

def calculate_reorder_metrics(forecast: pd.Series, avg_daily_demand: float, lead_time_days: int = 7) -> dict:
    """
    Calculate business metrics from forecast
    
    Args:
        forecast: Predicted demand series
        avg_daily_demand: Average daily demand
        lead_time_days: Supplier lead time in days
    
    Returns:
        Dict with reorder_point, safety_stock, stockout_risk
    """
    # Calculate average forecasted demand
    avg_forecast = forecast.mean()
    
    # Standard deviation of forecast (measure of uncertainty)
    forecast_std = forecast.std()
    
    # Safety stock = Z-score (1.96 for 95% service level) * std * sqrt(lead_time)
    safety_stock = int(1.96 * forecast_std * (lead_time_days ** 0.5))
    
    # Reorder point = (Average daily demand * Lead time) + Safety stock
    reorder_point = int((avg_daily_demand * lead_time_days) + safety_stock)
    
    # Stockout risk (simplified): if forecast variance is high, risk is higher
    cv = (forecast_std / avg_forecast * 100) if avg_forecast > 0 else 0  # Coefficient of variation
    if cv < 15:
        stockout_risk = "LOW"
    elif cv < 30:
        stockout_risk = "MEDIUM"
    else:
        stockout_risk = "HIGH"
    
    # Suggested order quantity (Economic Order Quantity approximation)
    suggested_order_qty = int(avg_daily_demand * 30)  # ~1 month supply
    
    return {
        "reorder_point": reorder_point,
        "safety_stock": safety_stock,
        "stockout_risk": stockout_risk,
        "suggested_order_qty": suggested_order_qty,
        "avg_daily_demand": round(avg_daily_demand, 2),
        "forecast_uncertainty": round(cv, 2)
    }

def forecast_for_sku(series: pd.Series,
                     horizon: int = 14,
                     models: Optional[List[str]] = None,
                     calculate_business_metrics: bool = True,
                     include_decomposition: bool = False
                     ) -> Tuple[Dict[str, pd.Series], Dict[str, dict], Dict[str, dict], Dict]:
    """
    Run forecasts on a single SKU using selected models.
    
    Returns:
        - forecasts: Dict of forecast series by model
        - conf_intervals: Dict of (lower, upper) bounds by model
        - metrics: Dict of performance metrics by model
        - decomposition: Seasonality decomposition (if requested)
    """
    if models is None:
        models = ["sarima", "ets", "prophet"]
    
    forecasts = {}
    conf_intervals = {}
    metrics = {}
    decomposition = None
    
    if len(series) < 7:
        # Not enough data for real forecasting
        return forecasts, conf_intervals, metrics, decomposition
    
    # Calculate historical metrics (use last 30 days as validation if available)
    validation_size = min(30, len(series) // 4)  # Use 25% of data or 30 days
    train_series = series.iloc[:-validation_size] if validation_size > 0 else series
    validation_series = series.iloc[-validation_size:] if validation_size > 0 else None
    
    # SARIMA
    if "sarima" in models:
        try:
            fc, _fit, lower, upper = sarima_forecast(train_series, steps=horizon, return_conf_int=True)
            forecasts["SARIMA"] = fc
            conf_intervals["SARIMA"] = {"lower": lower, "upper": upper}
            
            # Calculate accuracy if we have validation data
            if validation_series is not None and len(validation_series) > 0:
                val_fc, _ = sarima_forecast(train_series, steps=len(validation_series), return_conf_int=False)
                metrics["SARIMA"] = calculate_metrics(validation_series, val_fc)
            else:
                metrics["SARIMA"] = {"MAE": None, "RMSE": None, "MAPE": None}
        except Exception as e:
            print(f"SARIMA error: {e}")
    
    # ETS
    if "ets" in models:
        try:
            fc, _fit, lower, upper = exponential_smoothing_forecast(train_series, steps=horizon, return_conf_int=True)
            forecasts["ETS"] = fc
            conf_intervals["ETS"] = {"lower": lower, "upper": upper}
            
            if validation_series is not None and len(validation_series) > 0:
                val_fc, _ = exponential_smoothing_forecast(train_series, steps=len(validation_series), return_conf_int=False)
                metrics["ETS"] = calculate_metrics(validation_series, val_fc)
            else:
                metrics["ETS"] = {"MAE": None, "RMSE": None, "MAPE": None}
        except Exception as e:
            print(f"ETS error: {e}")
    
    # Prophet
    if "prophet" in models:
        try:
            fc, _fit, lower, upper = prophet_forecast(train_series, steps=horizon, return_conf_int=True)
            forecasts["PROPHET"] = fc
            conf_intervals["PROPHET"] = {"lower": lower, "upper": upper}
            
            if validation_series is not None and len(validation_series) > 0:
                val_fc, _ = prophet_forecast(train_series, steps=len(validation_series), return_conf_int=False)
                metrics["PROPHET"] = calculate_metrics(validation_series, val_fc)
            else:
                metrics["PROPHET"] = {"MAE": None, "RMSE": None, "MAPE": None}
        except Exception as e:
            print(f"Prophet error: {e}")
    
    # Add business metrics
    if calculate_business_metrics and forecasts:
        avg_daily_demand = series.mean()
        best_model = min(metrics.keys(), key=lambda k: metrics[k]["MAE"] if metrics[k]["MAE"] else float('inf'))
        best_forecast = forecasts.get(best_model)
        
        if best_forecast is not None:
            business_metrics = calculate_reorder_metrics(best_forecast, avg_daily_demand)
            metrics["BUSINESS"] = business_metrics
    
    # Add seasonality decomposition
    if include_decomposition and len(series) >= 14:
        decomposition = decompose_seasonality(series, period=7)
    
    return forecasts, conf_intervals, metrics, decomposition

def generate_forecasts(order_items_df: pd.DataFrame,
                       sku_list: Optional[List[str]] = None,
                       horizon: int = 14,
                       models: Optional[List[str]] = None,
                       include_decomposition: bool = False
                       ) -> Tuple[Dict[str, Dict[str, pd.Series]], Dict[str, Dict[str, dict]], Dict[str, dict], Dict[str, dict]]:
    """
    Main entry point for forecasting inside the dashboard.
    
    Input:
        - DataFrame of order_items
        - SKU list (optional)
        - forecast horizon
        - list of model names
        - include_decomposition: whether to add seasonality analysis
    
    Output:
        - forecasts: {SKU: {model: series}}
        - conf_intervals: {SKU: {model: {lower, upper}}}
        - metrics: {SKU: {model: metrics_dict}}
        - decompositions: {SKU: decomposition_dict}
    """
    if sku_list is None:
        sku_list = order_items_df["sku"].unique().tolist()
    
    all_forecasts = {}
    all_conf_intervals = {}
    all_metrics = {}
    all_decompositions = {}
    
    for sku in sku_list:
        series = prepare_daily_series(order_items_df, sku)
        if series.empty:
            continue
        
        forecasts, conf_intervals, metrics, decomposition = forecast_for_sku(
            series,
            horizon=horizon,
            models=models,
            include_decomposition=include_decomposition
        )
        
        if forecasts:
            all_forecasts[sku] = forecasts
            all_conf_intervals[sku] = conf_intervals
            all_metrics[sku] = metrics
            if decomposition:
                all_decompositions[sku] = decomposition
    
    return all_forecasts, all_conf_intervals, all_metrics, all_decompositions