from __future__ import annotations

import warnings
from datetime import date
from dateutil.relativedelta import relativedelta

import numpy as np
import pandas as pd

warnings.filterwarnings("ignore")

MIN_POINTS = 3
FORECAST_MONTHS = 3


def _next_month_starts(from_date: date, n: int) -> list[str]:
    months = []
    d = (from_date.replace(day=1) + relativedelta(months=1))
    for _ in range(n):
        months.append(d.strftime("%Y-%m-%d"))
        d += relativedelta(months=1)
    return months


def _fit_holt(values: list[float], n: int) -> tuple[list[float], list[float], list[float]]:
    from statsmodels.tsa.holtwinters import ExponentialSmoothing

    try:
        if len(values) >= 12:
            model = ExponentialSmoothing(values, trend="add", seasonal="add", seasonal_periods=12)
        elif len(values) >= 4:
            model = ExponentialSmoothing(values, trend="add", seasonal=None)
        else:
            model = ExponentialSmoothing(values, trend=None, seasonal=None)

        fit = model.fit(optimized=True, disp=False)
        forecast = fit.forecast(n)
        sigma = float(np.std(fit.resid)) if len(fit.resid) > 1 else float(np.mean(values) * 0.15)
        lower = np.maximum(0, forecast - 1.28 * sigma)
        upper = forecast + 1.28 * sigma
        return forecast.tolist(), lower.tolist(), upper.tolist()
    except Exception:
        avg = float(np.mean(values))
        sigma = float(np.std(values)) if len(values) > 1 else avg * 0.15
        return [avg] * n, [max(0, avg - 1.28 * sigma)] * n, [(avg + 1.28 * sigma)] * n


def forecast_spending(df: pd.DataFrame) -> list[dict]:
    """
    df columns: category_id, category_name, ds (monthly), y (amount)
    Returns list of category forecast dicts.
    """
    if df.empty:
        return []

    today = date.today()
    future_dates = _next_month_starts(today, FORECAST_MONTHS)
    results = []

    for (cat_id, cat_name), group in df.groupby(["category_id", "category_name"]):
        group = group.sort_values("ds")
        values = group["y"].tolist()
        actuals = [
            {"date": row["ds"].strftime("%Y-%m-%d"), "amount": row["y"]}
            for _, row in group.iterrows()
        ]

        if len(values) < MIN_POINTS:
            avg = float(np.mean(values))
            forecast_pts = [
                {"date": d, "amount": round(avg, 2), "lower": round(avg * 0.7, 2), "upper": round(avg * 1.3, 2)}
                for d in future_dates
            ]
        else:
            fcast, lower, upper = _fit_holt(values, FORECAST_MONTHS)
            forecast_pts = [
                {"date": d, "amount": round(max(0, f), 2), "lower": round(l, 2), "upper": round(u, 2)}
                for d, f, l, u in zip(future_dates, fcast, lower, upper)
            ]

        results.append({
            "category_id": cat_id,
            "category_name": cat_name,
            "monthly_avg": round(float(np.mean(values)), 2),
            "actuals": actuals[-6:],  # last 6 months for display
            "forecast": forecast_pts,
        })

    # Sort by monthly_avg descending (highest spend first)
    results.sort(key=lambda x: x["monthly_avg"], reverse=True)
    return results