#!/usr/bin/env python3
from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path
from typing import Any, Dict, List, Tuple

import numpy as np

from forecast_diagnostics import (
    diebold_mariano_squared_error,
    rolling_mean_forecast,
    rolling_rmse,
)


def _load_history(path: Path) -> Tuple[np.ndarray, np.ndarray, List[Dict[str, Any]]]:
    """
    Loads history containing actual and stochastic forecasts.

    Supported:
      - JSON array: [{"sprint":"...", "actual":..., "sc_forecast":...}, ...]
      - CSV with columns including: actual, sc_forecast (and optional sprint)
    """
    if path.suffix.lower() == ".json":
        rows = json.loads(path.read_text(encoding="utf-8"))
    elif path.suffix.lower() in {".csv", ".tsv"}:
        delim = "\t" if path.suffix.lower() == ".tsv" else ","
        rows = []
        with path.open("r", encoding="utf-8", newline="") as f:
            reader = csv.DictReader(f, delimiter=delim)
            for r in reader:
                rows.append(r)
    else:
        raise ValueError("Unsupported history file type; use .json, .csv, or .tsv")

    actual = np.asarray([float(r["actual"]) for r in rows], dtype=float)
    sc = np.asarray([float(r["sc_forecast"]) for r in rows], dtype=float)
    return actual, sc, rows


def main() -> None:
    ap = argparse.ArgumentParser(description="Forecast-accuracy comparison via Diebold–Mariano test.")
    ap.add_argument("--history", required=True, help="Path to history .json/.csv with columns actual, sc_forecast")
    ap.add_argument("--window", type=int, default=5, help="Rolling mean window K for baseline")
    ap.add_argument("--rmse-window", type=int, default=20, help="Window for rolling RMSE series (optional output)")
    ap.add_argument("--out", default="forecast_accuracy.json", help="Output JSON path")
    args = ap.parse_args()

    history_path = Path(args.history)
    actual, sc, rows = _load_history(history_path)

    idx, rm = rolling_mean_forecast(actual, window=args.window)
    actual_aligned = actual[idx]
    sc_aligned = sc[idx]

    dm = diebold_mariano_squared_error(actual_aligned, forecast_sc=sc_aligned, forecast_rm=rm)

    out: Dict[str, Any] = {
        "history_path": str(history_path),
        "window_k": args.window,
        "aligned_indices": idx.tolist(),
        "dm_test": dm.to_json(),
    }

    # Optional rolling RMSE curves (useful for plots)
    if args.rmse_window is not None and args.rmse_window >= 2 and actual_aligned.size >= args.rmse_window:
        rmse_idx_sc, rmse_sc = rolling_rmse(actual_aligned, sc_aligned, window=args.rmse_window)
        rmse_idx_rm, rmse_rm = rolling_rmse(actual_aligned, rm, window=args.rmse_window)
        out["rolling_rmse"] = {
            "rmse_window": int(args.rmse_window),
            "indices": rmse_idx_sc.tolist(),
            "sc": rmse_sc.tolist(),
            "rm": rmse_rm.tolist(),
        }

    Path(args.out).write_text(json.dumps(out, indent=2), encoding="utf-8")
    print(f"Wrote {args.out}")


if __name__ == "__main__":
    main()
