from __future__ import annotations
import pandas as pd
import numpy as np
from typing import Tuple
from sklearn.metrics import mean_absolute_error, mean_squared_error

try:
    from statsmodels.tsa.holtwinters import ExponentialSmoothing
    _HAS_SM = True
except Exception:
    _HAS_SM = False
from statsmodels.tsa.statespace.sarimax import SARIMAX

try:
    from prophet import Prophet
    _HAS_PROPHET = True
except Exception:
    _HAS_PROPHET = False

def calculate_metrics(actual: pd.Series, predicted: pd.Series) -> dict:
    """
    Calculate forecast accuracy metrics
    
    Returns dict with:
    - MAE: Mean Absolute Error
    - RMSE: Root Mean Squared Error  
    - MAPE: Mean Absolute Percentage Error
    """
    # Align series by index
    common_idx = actual.index.intersection(predicted.index)
    if len(common_idx) == 0:
        return {"MAE": None, "RMSE": None, "MAPE": None}
    
    actual_aligned = actual.loc[common_idx]
    predicted_aligned = predicted.loc[common_idx]
    
    mae = mean_absolute_error(actual_aligned, predicted_aligned)
    rmse = np.sqrt(mean_squared_error(actual_aligned, predicted_aligned))
    
    # MAPE - avoid division by zero
    mask = actual_aligned != 0
    if mask.sum() > 0:
        mape = np.mean(np.abs((actual_aligned[mask] - predicted_aligned[mask]) / actual_aligned[mask])) * 100
    else:
        mape = None
    
    return {
        "MAE": round(mae, 3),
        "RMSE": round(rmse, 3),
        "MAPE": round(mape, 2) if mape else None
    }

def holt_winters_forecast(y: pd.Series, horizon: int) -> pd.DataFrame:
    idx = pd.date_range(start=y.index[-1] + pd.Timedelta(days=1), periods=horizon, freq=y.index.freq or "D")
    if _HAS_SM and len(y) >= 28:
        season_length = 7
        model = ExponentialSmoothing(
            y.astype(float), trend="add", seasonal="add", seasonal_periods=season_length
        ).fit(optimized=True)
        fc = model.forecast(horizon)
        return pd.DataFrame({"ds": idx, "yhat": np.maximum(fc.values, 0.0)})
    if len(y) == 0:
        vals = [0] * horizon
    elif len(y) < 7:
        last_val = y.iloc[-1]
        vals = [last_val] * horizon
    else:
        season = 7
        last_season = y.tail(season).to_list()
        vals = [last_season[i % season] for i in range(horizon)]
    return pd.DataFrame({"ds": idx, "yhat": vals})

# ---------- ETS ----------
def exponential_smoothing_forecast(
    series: "pd.Series",
    steps: int = 31,
    seasonal_periods: int = 7,
    trend: str = "add",
    seasonal: str = "add",
    use_boxcox: bool | None = None,
    return_conf_int: bool = True,
):
    if not isinstance(series, pd.Series):
        raise TypeError("series must be a pandas Series indexed by datetime")
    s = series.copy()
    if not isinstance(s.index, pd.DatetimeIndex):
        s.index = pd.to_datetime(s.index)
    s = s.asfreq("D").fillna(0)
    model = ExponentialSmoothing(
        s, trend=trend, seasonal=seasonal, seasonal_periods=seasonal_periods, initialization_method="estimated"
    )
    fit = model.fit(optimized=True, use_boxcox=use_boxcox)
    future_index = pd.date_range(start=s.index[-1] + pd.Timedelta(days=1), periods=steps, freq="D")
    forecast = pd.Series(fit.forecast(steps), index=future_index)
    
    if return_conf_int:
        # Simulate confidence intervals for ETS
        # Use historical residuals to estimate uncertainty
        fitted_values = fit.fittedvalues
        residuals = s - fitted_values
        std_error = residuals.std()
        
        # 95% confidence interval (1.96 * std error)
        lower_bound = forecast - 1.96 * std_error
        upper_bound = forecast + 1.96 * std_error
        
        # Ensure non-negative bounds
        lower_bound = lower_bound.clip(lower=0)
        upper_bound = upper_bound.clip(lower=0)
        
        return forecast, fit, lower_bound, upper_bound
    
    return forecast, fit

# ---------- SARIMA ----------
def sarima_forecast(
    series: "pd.Series",
    steps: int = 31,
    order: tuple[int, int, int] = (1, 1, 1),
    seasonal_order: tuple[int, int, int, int] = (1, 1, 1, 7),
    return_conf_int: bool = True,
):
    if not isinstance(series, pd.Series):
        raise TypeError("series must be a pandas Series indexed by datetime")
    s = series.copy()
    if not isinstance(s.index, pd.DatetimeIndex):
        s.index = pd.to_datetime(s.index)
    s = s.asfreq("D").fillna(0)
    model = SARIMAX(s, order=order, seasonal_order=seasonal_order, enforce_stationarity=False, enforce_invertibility=False)
    fit = model.fit(disp=False)
    future_index = pd.date_range(start=s.index[-1] + pd.Timedelta(days=1), periods=steps, freq="D")
    
    if return_conf_int:
        # Get forecast with confidence intervals
        forecast_result = fit.get_forecast(steps=steps)
        fc = pd.Series(forecast_result.predicted_mean, index=future_index)
        conf_int = forecast_result.conf_int(alpha=0.05)  # 95% confidence
        
        # Create confidence interval series
        lower_bound = pd.Series(conf_int.iloc[:, 0].values, index=future_index).clip(lower=0)
        upper_bound = pd.Series(conf_int.iloc[:, 1].values, index=future_index).clip(lower=0)
        
        return fc, fit, lower_bound, upper_bound
    else:
        fc = pd.Series(fit.forecast(steps=steps), index=future_index)
        return fc, fit

# ---------- PROPHET (NEW!) ----------
def prophet_forecast(
    series: "pd.Series",
    steps: int = 31,
    yearly_seasonality: bool = True,
    weekly_seasonality: bool = True,
    daily_seasonality: bool = False,
    return_conf_int: bool = True,
):
    """
    Facebook Prophet forecasting model
    Better at handling trends and multiple seasonalities
    
    Args:
        series: Time series data (datetime index)
        steps: Number of days to forecast
        yearly_seasonality: Include yearly patterns
        weekly_seasonality: Include weekly patterns
        daily_seasonality: Include daily patterns
        return_conf_int: Return confidence intervals
    
    Returns:
        forecast: Predicted values
        model: Fitted Prophet model
        lower_bound: Lower confidence interval (if return_conf_int=True)
        upper_bound: Upper confidence interval (if return_conf_int=True)
    """
    if not _HAS_PROPHET:
        raise ImportError("Prophet not installed. Run: pip install prophet")
    
    if not isinstance(series, pd.Series):
        raise TypeError("series must be a pandas Series indexed by datetime")
    
    # Prepare data in Prophet format (needs 'ds' and 'y' columns)
    df = pd.DataFrame({
        'ds': series.index,
        'y': series.values
    })
    
    # Initialize Prophet model
    model = Prophet(
        yearly_seasonality=yearly_seasonality,
        weekly_seasonality=weekly_seasonality,
        daily_seasonality=daily_seasonality,
        seasonality_mode='additive',
        interval_width=0.95  # 95% confidence intervals
    )
    
    # Suppress Prophet's verbose output
    import logging
    logging.getLogger('prophet').setLevel(logging.WARNING)
    
    # Fit model
    model.fit(df)
    
    # Create future dataframe
    future = model.make_future_dataframe(periods=steps, freq='D')
    
    # Make predictions
    forecast_df = model.predict(future)
    
    # Extract only the future predictions
    future_index = pd.date_range(start=series.index[-1] + pd.Timedelta(days=1), periods=steps, freq='D')
    forecast = pd.Series(
        forecast_df['yhat'].tail(steps).values,
        index=future_index
    )
    
    # Ensure non-negative forecasts
    forecast = forecast.clip(lower=0)
    
    if return_conf_int:
        lower_bound = pd.Series(
            forecast_df['yhat_lower'].tail(steps).values,
            index=future_index
        ).clip(lower=0)
        upper_bound = pd.Series(
            forecast_df['yhat_upper'].tail(steps).values,
            index=future_index
        ).clip(lower=0)
        
        return forecast, model, lower_bound, upper_bound
    
    return forecast, model