"""Walk-forward evaluation, parameter sensitivity maps and holdout handling."""

from __future__ import annotations

import itertools
import math
from typing import Any

import numpy as np
import pandas as pd

from ktrade.backtest.costs import CostModel
from ktrade.backtest.engine import run_backtest
from ktrade.backtest.metrics import compute_metrics
from ktrade.strategy.base import Strategy


def _score(m: dict[str, Any]) -> float:
    s = m.get("sharpe")
    return s if s is not None and math.isfinite(s) else -1e9


def param_grid(strategy_cls: type[Strategy], grid: dict[str, list[Any]]) -> list[dict[str, Any]]:
    keys = list(grid)
    combos = []
    for values in itertools.product(*(grid[k] for k in keys)):
        p = dict(zip(keys, values, strict=True))
        try:
            strategy_cls.validate_params(p)
            combos.append(p)
        except ValueError:
            continue
    return combos


def walk_forward(
    strategy_cls: type[Strategy], df: pd.DataFrame, *, interval_min: int, grid: dict[str, list[Any]], n_folds: int = 4,
    train_frac: float = 0.6, anchored: bool = True, cost_model: CostModel | None = None, initial_cash: float = 10_000.0, embargo_bars: int = 0, **bt_kwargs,
) -> dict[str, Any]:
    """Split the series into n_folds consecutive test windows; for each, pick the grid point with the best IS Sharpe on
    the preceding train window (anchored = growing train, else rolling of fixed length), then evaluate OOS.
    An embargo of `embargo_bars` between train and test avoids leakage through overlapping feature windows."""
    n = len(df)
    combos = param_grid(strategy_cls, grid)
    if n < 300 or not combos:
        return {"folds": [], "error": "not enough data or empty grid", "n_trials": len(combos)}
    first_test = int(n * train_frac)
    test_len = (n - first_test) // n_folds
    folds = []
    oos_returns: list[np.ndarray] = []
    for k in range(n_folds):
        test_start = first_test + k * test_len
        test_end = n if k == n_folds - 1 else test_start + test_len
        train_start = 0 if anchored else max(0, test_start - first_test)
        train_end = max(train_start + 50, test_start - embargo_bars)
        train = df.iloc[train_start:train_end]
        # the test window gets a warm-up prefix (features need history) but metrics are computed on the test range only
        warm = min(train_end, strategy_cls.meta.min_bars * 3)
        test_with_warm = df.iloc[max(0, test_start - warm):test_end]
        best, best_m = None, None
        for p in combos:
            r = run_backtest(strategy_cls(p), train, interval_min=interval_min, initial_cash=initial_cash, cost_model=cost_model, **bt_kwargs)
            m = compute_metrics(r.equity, r.trades, interval_min, initial_cash, r.fees_paid, r.slippage_paid, r.turnover)
            if best_m is None or _score(m) > _score(best_m):
                best, best_m = p, m
        r = run_backtest(strategy_cls(best), test_with_warm, interval_min=interval_min, initial_cash=initial_cash, cost_model=cost_model, **bt_kwargs)
        eq = r.equity[r.equity["ts"] >= int(df.iloc[test_start]["ts"])].reset_index(drop=True)
        if len(eq) > 2:
            eq = eq.assign(equity=eq["equity"] / eq["equity"].iloc[0] * initial_cash)
        m_oos = compute_metrics(eq, [t for t in r.trades if t.ts >= int(df.iloc[test_start]["ts"])], interval_min, initial_cash, 0.0, 0.0, r.turnover) if len(eq) > 2 else {"error": "short"}
        oos_returns.append(np.diff(np.log(np.maximum(eq["equity"].to_numpy(), 1e-12))) if len(eq) > 2 else np.array([]))
        folds.append({
            "fold": k, "train": [int(df.iloc[train_start]["ts"]), int(df.iloc[train_end - 1]["ts"])], "test": [int(df.iloc[test_start]["ts"]), int(df.iloc[test_end - 1]["ts"])],
            "best_params": best, "is_sharpe": best_m.get("sharpe") if best_m else None, "oos_sharpe": m_oos.get("sharpe"), "oos_return": m_oos.get("total_return"), "oos_max_dd": m_oos.get("max_drawdown"),
        })
    all_oos = np.concatenate([x for x in oos_returns if len(x)]) if any(len(x) for x in oos_returns) else np.array([])
    bars_per_year = 365 * 24 * 60 / interval_min
    stitched = {"sharpe": float(all_oos.mean() / all_oos.std(ddof=1) * math.sqrt(bars_per_year)) if len(all_oos) > 10 and all_oos.std(ddof=1) > 0 else None, "total_return": float(math.exp(all_oos.sum()) - 1) if len(all_oos) else None, "bars": int(len(all_oos))}
    return {"folds": folds, "anchored": anchored, "n_trials": len(combos), "grid": grid, "embargo_bars": embargo_bars, "stitched_oos": stitched}


def sensitivity_map(strategy_cls: type[Strategy], df: pd.DataFrame, *, interval_min: int, x: tuple[str, list[Any]], y: tuple[str, list[Any]], base_params: dict[str, Any], cost_model: CostModel | None = None, initial_cash: float = 10_000.0, metric: str = "sharpe", **bt_kwargs) -> dict[str, Any]:
    xs, ys = x[1], y[1]
    z: list[list[float | None]] = []
    for yv in ys:
        row: list[float | None] = []
        for xv in xs:
            p = {**base_params, x[0]: xv, y[0]: yv}
            try:
                r = run_backtest(strategy_cls(p), df, interval_min=interval_min, initial_cash=initial_cash, cost_model=cost_model, **bt_kwargs)
                m = compute_metrics(r.equity, r.trades, interval_min, initial_cash, r.fees_paid, r.slippage_paid, r.turnover)
                val = m.get(metric)
                row.append(float(val) if val is not None and math.isfinite(val) else None)
            except ValueError:
                row.append(None)
        z.append(row)
    return {"x": {"name": x[0], "values": xs}, "y": {"name": y[0], "values": ys}, "metric": metric, "z": z}