# forecasting/run_forecast.py
from __future__ import annotations

import argparse
from dataclasses import dataclass
from pathlib import Path
import io
import sys

import numpy as np
import pandas as pd
import requests
# --- Import models: prefer package import, fallback to relative when running as module ---
try:
    from forecasting.models import exponential_smoothing_forecast, sarima_forecast, prophet_forecast
except Exception:  # pragma: no cover
    from .models import exponential_smoothing_forecast, sarima_forecast, prophet_forecast  # type: ignore

# -------------------------- Utilities --------------------------

@dataclass
class Columns:
    dt: str
    sku: str
    qty: str


def detect_columns(df: pd.DataFrame) -> Columns:
    cols = {c.lower().strip(): c for c in df.columns}
    # date/timestamp
    for cand in ("order_ts", "order_date", "date", "timestamp", "ts"):
        if cand in cols:
            dt = cols[cand]
            break
    else:
        raise ValueError("Could not find a date/timestamp column. "
                         "Expected one of: order_ts, order_date, date, timestamp, ts")

    # sku
    for cand in ("sku_id", "sku", "item_id", "product_id"):
        if cand in cols:
            sku = cols[cand]
            break
    else:
        raise ValueError("Could not find a SKU column. "
                         "Expected one of: sku_id, sku, item_id, product_id")

    # quantity
    for cand in ("quantity", "qty", "units", "order_qty"):
        if cand in cols:
            qty = cols[cand]
            break
    else:
        raise ValueError("Could not find a quantity column. "
                         "Expected one of: quantity, qty, units, order_qty")

    return Columns(dt=dt, sku=sku, qty=qty)


def read_orders(path: Path) -> tuple[pd.DataFrame, Columns]:
    if not path.exists():
        raise FileNotFoundError(f"Orders file not found: {path}")
    # Robust CSV read (sniff delimiter, skip bad lines)
    raw = path.read_bytes()
    df = pd.read_csv(
        io.BytesIO(raw),
        sep=None,
        engine="python",
        on_bad_lines="skip",
        quotechar='"',
    )
    df.columns = [str(c).strip().lower() for c in df.columns]
    cols = detect_columns(df)
    # Normalize names
    df = df.rename(columns={cols.dt: "order_ts", cols.sku: "sku_id", cols.qty: "quantity"})
    # Parse dates
    df["order_ts"] = pd.to_datetime(df["order_ts"], errors="coerce")
    df = df.dropna(subset=["order_ts", "sku_id"])
    # Ensure numeric quantity
    df["quantity"] = pd.to_numeric(df["quantity"], errors="coerce").fillna(0)
    return df, Columns(dt="order_ts", sku="sku_id", qty="quantity")


def daily_series(df: pd.DataFrame, sku: str) -> pd.Series:
    """Return daily demand series for a SKU (DateIndex, daily sum)."""
    g = (df.loc[df["sku_id"] == sku]
           .groupby(df["order_ts"].dt.date)["quantity"]
           .sum())
    s = pd.Series(g, dtype=float)
    s.index = pd.to_datetime(s.index)
    return s.asfreq("D").fillna(0.0)


def rolling_last_k_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> dict:
    from sklearn.metrics import mean_absolute_error, mean_squared_error
    mae = float(mean_absolute_error(y_true, y_pred))
    rmse = float(np.sqrt(mean_squared_error(y_true, y_pred)))
    mape = float(np.mean(np.abs((y_true - y_pred) / (y_true + 1e-9))) * 100.0)
    return {"MAE": mae, "RMSE": rmse, "MAPE": mape}


# -------------------------- CLI --------------------------

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Run SKU demand forecasts and export metrics/plots")
    p.add_argument("--orders", default="app/data/mock_orders.csv",
                   help="Path to orders CSV (default: app/data/mock_orders.csv)")
    p.add_argument("--sku", required=True,
                   help="SKU id to forecast OR 'all' to run every SKU")
    p.add_argument("--start", default=None,
                   help="Train start date (YYYY-MM-DD). Default: first available date")
    p.add_argument("--end", default=None,
                   help="Train end date (YYYY-MM-DD). Default: last available date")
    p.add_argument("--horizon", type=int, default=14,
                   help="Forecast horizon in days (default: 14)")
    p.add_argument("--model", choices=["ets", "sarima", "prophet", "all"], default="sarima",
                   help="Model to run (default: sarima)")
    p.add_argument("--outdir", default="outputs",
                   help="Output directory (default: outputs)")
    return p.parse_args()


# -------------------------- Main --------------------------

def main() -> None:
    args = parse_args()
    root = Path.cwd()

    # 1) Try CWD-relative path
    # orders_path = Path(args.orders).expanduser().resolve()
    orders_path = Path(args.orders).expanduser().resolve()

    # 2) Fallback: project root (folder containing the package)
    if not orders_path.exists():
        pkg_root = Path(__file__).resolve().parents[1]
        alt_path = (pkg_root / args.orders).resolve()
        if alt_path.exists():
            orders_path = alt_path

    # 3) Final check
    if not orders_path.exists():
        raise FileNotFoundError(f"Orders file not found: {orders_path}")

    # Make outputs dir (always)
    outdir = (root / args.outdir).resolve()
    outdir.mkdir(parents=True, exist_ok=True)

    print(f"[i] Reading: {orders_path}")
    df, _cols = read_orders(orders_path)

    # All SKUs or a single one
    if args.sku.lower() == "all":
        sku_list = sorted(df["sku_id"].dropna().unique().tolist())
    else:
        sku_list = [args.sku]

    all_rows = []     # for outputs/forecast.csv (aggregate)
    metric_rows = []  # for outputs/metrics.csv

    for sku in sku_list:
        s = daily_series(df, sku)
        if s.empty:
            print(f"[skip] No data for {sku}")
            continue

        # Trim train window
        if args.start:
            s = s[s.index >= pd.to_datetime(args.start)]
        if args.end:
            s = s[s.index <= pd.to_datetime(args.end)]

        if len(s) < 7:
            print(f"[skip] Not enough history for {sku} (len={len(s)})")
            continue

        # Train set = all data (simple pipeline). For metrics, compare last k days if present.
        k = min(args.horizon, max(1, len(s) // 5))  # small holdout
        train = s.iloc[:-k] if len(s) > k else s
        holdout = s.iloc[-k:] if len(s) > k else pd.Series(dtype=float)

        forecasts: dict[str, pd.Series] = {}

        if args.model in ("ets", "all"):
            fc_ets, _fit_ets = exponential_smoothing_forecast(
                series=train.squeeze(), steps=args.horizon,
                seasonal_periods=7, trend="add", seasonal="add", use_boxcox=None,
            )
            forecasts["ETS"] = fc_ets

        if args.model in ("sarima", "all"):
            fc_sarima, _fit_sarima = sarima_forecast(
                series=train.squeeze(), steps=args.horizon,
                order=(1, 1, 1), seasonal_order=(1, 1, 1, 7),
            )
            forecasts["SARIMA"] = fc_sarima

        if args.model in ("prophet", "all"):
            try:
                fc_prophet, _fit_prophet = prophet_forecast(
                    series=train.squeeze(), steps=args.horizon,
                    yearly_seasonality=True, weekly_seasonality=True, daily_seasonality=False,
                )
                forecasts["PROPHET"] = fc_prophet
            except Exception as e:
                print(f"[warn] Prophet forecast failed for {sku}: {e}")

        # Save per-SKU outputs and accumulate combined forecast
        for model_name, fc in forecasts.items():
            # Per-SKU CSV
            sku_csv = outdir / f"forecast_{model_name}_SKU-{sku}.csv"
            pd.DataFrame({"date": fc.index, "forecast": fc.values}).to_csv(sku_csv, index=False)
            print(f"[ok] wrote {sku_csv.name}")

            # Per-SKU PNG (optional quick plot)
            try:
                import matplotlib.pyplot as plt
                plt.figure(figsize=(8, 3.5))
                s.plot(alpha=0.25, linewidth=1, label="History")
                s.rolling(7).mean().plot(linewidth=2, label="7-day avg")
                plt.plot(fc.index, fc.values, linewidth=2, label=f"{model_name} fc")
                if args.end:
                    plt.axvline(pd.to_datetime(args.end), linestyle="--", linewidth=1)
                plt.title(f"{model_name} • {sku}")
                plt.xlabel("date"); plt.ylabel("demand"); plt.legend(); plt.tight_layout()
                fig_path = outdir / f"forecast_{model_name}_SKU-{sku}.png"
                plt.savefig(fig_path, dpi=150)
                plt.close()
                print(f"[ok] wrote {fig_path.name}")
            except Exception as e:  # plotting optional
                print(f"[warn] could not plot {sku}: {e}")

            # Append to combined forecast rows
            all_rows.append(pd.DataFrame({
                "sku_id": sku,
                "model": model_name,
                "date": fc.index,
                "forecast": fc.values
            }))

            # Metrics vs holdout (if any)
            if not holdout.empty:
                n = min(len(holdout), len(fc))
                m = rolling_last_k_metrics(holdout.iloc[:n].to_numpy(dtype=float),
                                           fc.iloc[:n].to_numpy(dtype=float))
                metric_rows.append({
                    "sku_id": sku,
                    "model": model_name,
                    "MAE": m["MAE"],
                    "RMSE": m["RMSE"],
                    "MAPE": m["MAPE"],
                    "train_len": int(len(train)),
                    "holdout_len": int(len(holdout)),
                })

    # Write combined artifacts
    if all_rows:
        all_fc = pd.concat(all_rows, ignore_index=True)
        all_fc.to_csv(outdir / "forecast.csv", index=False)
        print(f"[ok] wrote forecast.csv ({len(all_fc)} rows)")

    if metric_rows:
        metrics_df = pd.DataFrame(metric_rows)
        metrics_df.to_csv(outdir / "metrics.csv", index=False)
        print(f"[ok] wrote metrics.csv ({len(metrics_df)} rows)")

    if not all_rows:
        print("[done] No forecasts created (check --start/--end, sku filter, or data).")


if __name__ == "__main__":
    # Allow running both as a script and via `python -m forecasting.run_forecast`
    # Ensure project root is on sys.path when executed directly
    here = Path(__file__).resolve()
    pkg_root = here.parents[1]
    if str(pkg_root) not in sys.path:
        sys.path.insert(0, str(pkg_root))
    main()