"""Backtest harness for Vanta Crunch models.

Replays historical Binance candle data through a TrackerBase model, calling
tick() with each new candle and predict() at regular intervals. Scores
predictions against actual 5-minute forward returns.

Usage:
    from vanta.backtest import BacktestRunner
    from my_model import MyTracker

    result = BacktestRunner(model=MyTracker()).run(
        start="2026-01-01", end="2026-02-01"
    )
    result.predictions_df   # DataFrame in notebook
    result.metrics           # rolling window aggregates + risk metrics
    result.summary()         # formatted output
"""

from __future__ import annotations

import json
import logging
import os
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Callable

logger = logging.getLogger(__name__)


# ── BacktestResult ──


class BacktestResult:
    """Container for backtest output. Notebook-friendly."""

    def __init__(
        self,
        predictions: list[dict[str, Any]],
        metrics: dict[str, float],
        config: dict[str, Any],
    ):
        self._predictions = predictions
        self.metrics = metrics
        self.config = config

    @property
    def predictions_df(self):
        """Return predictions as a pandas DataFrame."""
        import pandas as pd
        return pd.DataFrame(self._predictions)

    def _repr_html_(self) -> str:
        """Rich HTML display for Jupyter notebooks."""
        rows = ""
        for key, value in self.metrics.items():
            rows += f"<tr><td><b>{key}</b></td><td>{value:+.6f}</td></tr>"

        pred_count = len(self._predictions)
        scored = sum(1 for p in self._predictions if p.get("score") is not None)

        return f"""
        <div style="font-family: monospace; padding: 10px;">
            <h3>🏔️ Vanta Crunch Backtest</h3>
            <p>Symbol: <b>{self.config.get('subject', 'N/A')}</b> |
               Period: {self.config.get('start', '?')} → {self.config.get('end', '?')} |
               Predictions: {pred_count} | Scored: {scored}</p>
            <table border="1" cellpadding="5" style="border-collapse: collapse;">
                <tr><th>Metric</th><th>Value</th></tr>
                {rows}
            </table>
        </div>
        """

    @property
    def diversity(self) -> dict[str, Any] | None:
        """Fetch diversity feedback from the coordinator."""
        model_id = self.config.get("model_id")
        if not model_id:
            return None
        try:
            import requests
            from vanta.config import COORDINATOR_URL
            url = f"{COORDINATOR_URL.rstrip('/')}/reports/models/{model_id}/diversity"
            resp = requests.get(url, timeout=10)
            if resp.status_code == 200:
                return resp.json()
        except Exception:
            pass
        return None

    def summary(self, show_diversity: bool = True) -> str:
        """Return a formatted summary string."""
        lines = [
            "═" * 60,
            "  VANTA CRUNCH BACKTEST",
            "═" * 60,
            f"  Symbol:     {self.config.get('subject', 'N/A')}",
            f"  Period:      {self.config.get('start', '?')} → {self.config.get('end', '?')}",
            f"  Predictions: {len(self._predictions)}",
            f"  Horizon:     {self.config.get('horizon_seconds', 300)}s",
            "─" * 60,
            "  PnL METRICS:",
        ]

        cost_bps = self.config.get("transaction_cost_bps", 0)
        if cost_bps:
            lines.append(f"  Txn Cost:    {cost_bps} bps")

        pnl_keys = {"score_recent", "score_steady", "score_anchor", "mean_return"}
        cost_keys = {"gross_return", "cost_drag", "total_txn_cost", "total_turnover", "mean_turnover"}
        risk_keys = {"hit_rate", "max_drawdown", "sortino_ratio", "ic", "ic_sharpe"}
        diversity_keys = {"model_correlation", "ensemble_correlation", "contribution", "fnc"}

        for key, value in self.metrics.items():
            if key in pnl_keys:
                lines.append(f"    {key:20s} {value:+.8f}")

        cost_metrics = {k: v for k, v in self.metrics.items() if k in cost_keys}
        if cost_metrics and cost_bps:
            lines.append("─" * 60)
            lines.append("  TRANSACTION COSTS:")
            for key, value in cost_metrics.items():
                lines.append(f"    {key:20s} {value:+.8f}")

        risk_metrics = {k: v for k, v in self.metrics.items() if k in risk_keys}
        if risk_metrics:
            lines.append("─" * 60)
            lines.append("  RISK METRICS:")
            for key, value in risk_metrics.items():
                lines.append(f"    {key:20s} {value:+.6f}")

        diversity_metrics = {k: v for k, v in self.metrics.items() if k in diversity_keys}
        if diversity_metrics:
            lines.append("─" * 60)
            lines.append("  DIVERSITY:")
            for key, value in diversity_metrics.items():
                lines.append(f"    {key:20s} {value:+.6f}")

        if show_diversity:
            div = self.diversity
            if div:
                lines.append("─" * 60)
                lines.append("  DIVERSITY FEEDBACK:")
                ds = div.get("diversity_score")
                if ds is not None:
                    lines.append(f"    {'diversity_score':20s} {ds:.4f}")
                for g in div.get("guidance", []):
                    lines.append(f"    → {g}")

        lines.append("═" * 60)
        text = "\n".join(lines)
        print(text)
        return text

    def plot(self, kind: str = "all", save: str | None = None, figsize: tuple = (14, 10)):
        """Plot backtest results. Requires matplotlib.

        Args:
            kind: "all", "equity", "drawdown", "hit_rate", or "scatter"
            save: Save to file path instead of showing
            figsize: Figure size (width, height)
        """
        try:
            import matplotlib.pyplot as plt
        except ImportError:
            print("Install matplotlib for plotting: pip install matplotlib")
            return

        scored = [p for p in self._predictions if p.get("score") is not None]
        if not scored:
            print("No scored predictions to plot.")
            return

        timestamps = [p["ts"] for p in scored]
        scores = [p["score"] for p in scored]
        cumulative = []
        running = 0.0
        for s in scores:
            running += s
            cumulative.append(running)

        plots = {
            "equity": self._plot_equity,
            "drawdown": self._plot_drawdown,
            "hit_rate": self._plot_hit_rate,
            "scatter": self._plot_scatter,
        }

        if kind == "all":
            fig, axes = plt.subplots(2, 2, figsize=figsize)
            fig.suptitle(
                f"Backtest: {self.config.get('subject', '?')}  "
                f"({self.config.get('start', '?')} → {self.config.get('end', '?')})",
                fontsize=14,
            )
            self._plot_equity(axes[0, 0], timestamps, cumulative)
            self._plot_drawdown(axes[0, 1], timestamps, cumulative)
            self._plot_hit_rate(axes[1, 0], scored)
            self._plot_scatter(axes[1, 1], scored)
            plt.tight_layout()
        elif kind in plots:
            fig, ax = plt.subplots(figsize=(figsize[0], figsize[1] // 2))
            if kind in ("equity", "drawdown"):
                plots[kind](ax, timestamps, cumulative)
            else:
                plots[kind](ax, scored)
            plt.tight_layout()
        else:
            print(f"Unknown plot kind: {kind!r}. Use: all, equity, drawdown, hit_rate, scatter")
            return

        if save:
            plt.savefig(save, dpi=150, bbox_inches="tight")
            print(f"Saved to {save}")
        else:
            plt.show()

    @staticmethod
    def _plot_equity(ax, timestamps, cumulative):
        ax.plot(timestamps, cumulative, linewidth=1)
        ax.set_title("Cumulative PnL")
        ax.set_xlabel("Time")
        ax.set_ylabel("Cumulative Score")
        ax.axhline(y=0, color="gray", linestyle="--", linewidth=0.5)
        ax.grid(True, alpha=0.3)

    @staticmethod
    def _plot_drawdown(ax, timestamps, cumulative):
        peak = 0.0
        drawdowns = []
        for c in cumulative:
            peak = max(peak, c)
            drawdowns.append(c - peak)
        ax.fill_between(timestamps, drawdowns, 0, alpha=0.4, color="red")
        ax.set_title("Drawdown")
        ax.set_xlabel("Time")
        ax.set_ylabel("Drawdown")
        ax.grid(True, alpha=0.3)

    @staticmethod
    def _plot_hit_rate(ax, scored):
        window = max(20, len(scored) // 50)
        correct = [1.0 if p.get("direction_correct") else 0.0 for p in scored]
        rolling = []
        for i in range(len(correct)):
            start = max(0, i - window + 1)
            rolling.append(sum(correct[start:i+1]) / (i - start + 1))
        timestamps = [p["ts"] for p in scored]
        ax.plot(timestamps, rolling, linewidth=1, color="green")
        ax.axhline(y=0.5, color="gray", linestyle="--", linewidth=0.5)
        ax.set_title(f"Hit Rate (rolling {window})")
        ax.set_xlabel("Time")
        ax.set_ylabel("Hit Rate")
        ax.set_ylim(0, 1)
        ax.grid(True, alpha=0.3)

    @staticmethod
    def _plot_scatter(ax, scored):
        expected = [p.get("expected_return", 0) for p in scored]
        actual = [p.get("actual_return", 0) for p in scored]
        ax.scatter(expected, actual, alpha=0.3, s=5)
        ax.axhline(y=0, color="gray", linestyle="--", linewidth=0.5)
        ax.axvline(x=0, color="gray", linestyle="--", linewidth=0.5)
        ax.set_title("Expected vs Actual Return")
        ax.set_xlabel("Expected Return")
        ax.set_ylabel("Actual Return")
        ax.grid(True, alpha=0.3)

    def __repr__(self) -> str:
        return f"BacktestResult(predictions={len(self._predictions)}, metrics={self.metrics})"


# ── BacktestClient ──


class BacktestClient:
    """Fetches backfill parquet files from a coordinator and caches locally."""

    def __init__(
        self,
        coordinator_url: str | None = None,
        cache_dir: str = ".cache/backtest",
    ):
        if coordinator_url is None:
            from vanta.config import COORDINATOR_URL
            coordinator_url = COORDINATOR_URL
        self.coordinator_url = coordinator_url.rstrip("/")
        self.cache_dir = Path(cache_dir)

    def pull(
        self,
        source: str | None = None,
        subject: str | None = None,
        kind: str | None = None,
        granularity: str | None = None,
        start: str | datetime = "2026-01-01",
        end: str | datetime = "2026-02-01",
        refresh: bool = False,
    ) -> list[Path]:
        """Download matching parquet files from coordinator, cache locally."""
        import requests
        from vanta.config import (
            DEFAULT_SOURCE, DEFAULT_SUBJECT, DEFAULT_KIND, DEFAULT_GRANULARITY,
        )

        source = source or DEFAULT_SOURCE
        subject = subject or DEFAULT_SUBJECT
        kind = kind or DEFAULT_KIND
        granularity = granularity or DEFAULT_GRANULARITY

        start_dt = _parse_date(start)
        end_dt = _parse_date(end)

        resp = requests.get(f"{self.coordinator_url}/data/backfill/index", timeout=30)
        resp.raise_for_status()
        manifest = resp.json()

        prefix = f"{source}/{subject}/{kind}/{granularity}/"
        matching = []
        for entry in manifest:
            path = entry["path"]
            if not path.startswith(prefix):
                continue
            date_str = entry.get("date", "")
            try:
                file_date = datetime.strptime(date_str, "%Y-%m-%d").replace(tzinfo=timezone.utc)
            except ValueError:
                continue
            if start_dt <= file_date <= end_dt:
                matching.append(entry)

        downloaded: list[Path] = []
        for entry in matching:
            rel_path = entry["path"]
            local_path = self.cache_dir / rel_path

            if local_path.exists() and not refresh:
                downloaded.append(local_path)
                continue

            url = f"{self.coordinator_url}/data/backfill/{rel_path}"
            logger.info("downloading: %s", url)
            file_resp = requests.get(url, timeout=120)
            file_resp.raise_for_status()

            local_path.parent.mkdir(parents=True, exist_ok=True)
            local_path.write_bytes(file_resp.content)
            downloaded.append(local_path)

        return downloaded


# ── BacktestRunner ──


class BacktestRunner:
    """Replays historical data through a TrackerBase model, scores predictions.

    Simulates the production pipeline:
      1. tick(candle_data) called for each new candle
      2. predict(symbol, horizon_seconds, step_seconds) called every step_seconds
      3. Scored against actual forward return at horizon
    """

    def __init__(
        self,
        model,
        scoring_fn: Callable[[dict, dict], dict] | None = None,
        cache_dir: str = ".cache/backtest",
        transaction_cost_bps: float = 0.0,
    ):
        self.model = model
        self.scoring_fn = scoring_fn or _default_scoring_fn()
        self.cache_dir = Path(cache_dir)
        self.transaction_cost_bps = transaction_cost_bps

    def run(
        self,
        source: str | None = None,
        subject: str | None = None,
        kind: str | None = None,
        granularity: str | None = None,
        start: str | datetime = "2026-01-01",
        end: str | datetime = "2026-02-01",
        window_size: int = 120,
        prediction_interval_seconds: int = 900,
        horizon_seconds: int = 3600,
        transaction_cost_bps: float | None = None,
    ) -> BacktestResult:
        """Replay cached data through the model, score predictions.

        Simulates the production tick() + predict() flow:
        - Each candle triggers tick() with cumulative window
        - predict() called every prediction_interval_seconds
        - Scored against actual return at horizon_seconds forward

        Args:
            transaction_cost_bps: Cost in basis points charged on signal
                change (turnover).  Applied as:
                ``cost = abs(new_signal - old_signal) * bps / 10_000``
                and deducted from the gross score.  Overrides the value
                set on the runner if provided.  Typical values:
                  - 0   = frictionless (default, current behaviour)
                  - 4   = Binance futures taker fee
                  - 10  = futures + slippage estimate
                  - 20  = spot trading, conservative
        """
        import pandas as pd
        from vanta.config import (
            DEFAULT_SOURCE, DEFAULT_SUBJECT, DEFAULT_KIND, DEFAULT_GRANULARITY,
        )

        source = source or DEFAULT_SOURCE
        subject = subject or DEFAULT_SUBJECT
        kind = kind or DEFAULT_KIND
        granularity = granularity or DEFAULT_GRANULARITY

        start_dt = _parse_date(start)
        end_dt = _parse_date(end)

        # Load cached parquet data — auto-pull if not cached
        data_dir = self.cache_dir / source / subject / kind / granularity
        if not data_dir.exists() or not list(data_dir.glob("*.parquet")):
            logger.info("No cached data found, pulling from coordinator...")
            client = BacktestClient(cache_dir=str(self.cache_dir))
            client.pull(source=source, subject=subject, kind=kind,
                        granularity=granularity, start=start, end=end)

        frames = []
        for parquet_file in sorted(data_dir.glob("*.parquet")):
            date_str = parquet_file.stem
            try:
                file_date = datetime.strptime(date_str, "%Y-%m-%d").replace(tzinfo=timezone.utc)
            except ValueError:
                continue
            if start_dt <= file_date <= end_dt:
                frames.append(pd.read_parquet(parquet_file))

        if not frames:
            raise FileNotFoundError(
                f"No data for {subject} from {start} to {end}. "
                f"Coordinator may not have backfill data for this range yet."
            )

        df = pd.concat(frames, ignore_index=True).sort_values("ts_event").reset_index(drop=True)

        if not pd.api.types.is_datetime64_any_dtype(df["ts_event"]):
            df["ts_event"] = pd.to_datetime(df["ts_event"], utc=True)

        # ── Pre-convert DataFrame to numpy arrays for fast replay ─────────
        all_candles = _df_to_candle_array(df)  # list of dicts, indexed by row
        ts_array = df["ts_event"].values  # numpy datetime64 array
        close_array = df["close"].values.astype(float)

        # Pre-compute ALL multi-TF candle arrays once
        all_5m = _aggregate_candles(all_candles, 5, len(all_candles))
        all_15m = _aggregate_candles(all_candles, 15, len(all_candles))
        all_1h = _aggregate_candles(all_candles, 60, len(all_candles))
        all_4h = _aggregate_candles(all_candles, 240, len(all_candles))
        all_1d = _aggregate_candles(all_candles, 1440, len(all_candles))

        # Build timestamp→index maps for fast lookups
        idx_5m = {c["ts"]: i for i, c in enumerate(all_5m)}
        idx_15m = {c["ts"]: i for i, c in enumerate(all_15m)}
        idx_1h = {c["ts"]: i for i, c in enumerate(all_1h)}
        idx_4h = {c["ts"]: i for i, c in enumerate(all_4h)}
        idx_1d = {c["ts"]: i for i, c in enumerate(all_1d)}

        horizon_steps = horizon_seconds // 60  # horizon in 1m candle steps

        # ── Load funding rate data (best-effort) ──────────────────────────
        funding_map: dict[int, float] = {}
        try:
            from vanta.data.funding import load_funding
            funding_map = load_funding(
                symbol=subject,
                start=str(start),
                end=str(end),
                cache_dir=str(self.cache_dir),
            )
        except Exception:
            pass  # No funding data available — models will see funding=None

        # Build sorted funding timestamps for bisect lookup
        import bisect as _bisect
        funding_ts_sorted = sorted(funding_map.keys())

        def _get_funding(as_of_ts: int) -> dict | None:
            """Get the most recent funding rate at or before as_of_ts."""
            if not funding_ts_sorted:
                return None
            idx = _bisect.bisect_right(funding_ts_sorted, as_of_ts) - 1
            if idx < 0:
                return None
            ts = funding_ts_sorted[idx]
            return {"funding_rate": funding_map[ts], "ts": ts}

        # ── Warmup: feed the initial window in bulk ───────────────────────
        #
        # The first `window_size` candles are fed as a single tick so that
        # models can bootstrap their indicators from the initial history.
        warmup_candle_ts = all_candles[window_size - 1]["ts"]

        def _slice_tf_completed(all_tf, idx_map, bucket_secs, max_bars, up_to_ts):
            """Return only fully-completed higher-TF bars up to (excluding) the
            bucket that ``up_to_ts`` falls into, so there is no overlap with
            the first incremental delivery during replay."""
            current_bucket = (up_to_ts // bucket_secs) * bucket_secs
            # Find the last completed bucket (the one before the current)
            prev_bucket = current_bucket - bucket_secs
            pos = idx_map.get(prev_bucket)
            if pos is None:
                # fallback: find nearest bucket < current_bucket
                candidates = [j for j, c in enumerate(all_tf) if c["ts"] < current_bucket]
                pos = candidates[-1] if candidates else -1
            if pos < 0:
                return []
            start = max(0, pos - max_bars + 1)
            return all_tf[start:pos + 1]

        warmup_data = {
            "symbol": subject,
            "asof_ts": warmup_candle_ts,
            "candles_1m": all_candles[:window_size],
            "candles_5m": _slice_tf_completed(all_5m, idx_5m, 300, 60, warmup_candle_ts),
            "candles_15m": _slice_tf_completed(all_15m, idx_15m, 900, 40, warmup_candle_ts),
            "candles_1h": _slice_tf_completed(all_1h, idx_1h, 3600, 24, warmup_candle_ts),
            "candles_4h": _slice_tf_completed(all_4h, idx_4h, 14400, 12, warmup_candle_ts),
            "candles_1d": _slice_tf_completed(all_1d, idx_1d, 86400, 7, warmup_candle_ts),
            "orderbook": None,
            "funding": _get_funding(warmup_candle_ts),
        }
        self.model.tick(warmup_data)

        # Track which higher-TF bucket was last delivered so we only send
        # a candle when its bucket changes (i.e. a new bar has formed).
        last_bucket_5m = warmup_candle_ts // 300
        last_bucket_15m = warmup_candle_ts // 900
        last_bucket_1h = warmup_candle_ts // 3600
        last_bucket_4h = warmup_candle_ts // 14400
        last_bucket_1d = warmup_candle_ts // 86400

        # ── Transaction cost config ────────────────────────────────────────
        cost_bps = transaction_cost_bps if transaction_cost_bps is not None else self.transaction_cost_bps
        cost_multiplier = cost_bps / 10_000.0  # bps → decimal

        # ── Replay loop ───────────────────────────────────────────────────
        #
        # Each tick delivers only NEW candles:
        #   - candles_1m: always the single new 1-minute candle
        #   - candles_5m/15m/1h: the new bar only when its bucket rolls over
        #
        # Models accumulate state incrementally via their indicators.

        predictions: list[dict[str, Any]] = []
        last_predict_ts: datetime | None = None
        prev_signal: float = 0.0  # previous expected_return for turnover cost
        interval = timedelta(seconds=prediction_interval_seconds)

        for i in range(window_size, len(df)):
            current_ts = pd.Timestamp(ts_array[i]).to_pydatetime()
            if current_ts.tzinfo is None:
                current_ts = current_ts.replace(tzinfo=timezone.utc)

            candle_ts = all_candles[i]["ts"]

            # Higher-TF candles: only deliver when the bucket changes
            new_5m = []
            bucket_5m = candle_ts // 300
            if bucket_5m != last_bucket_5m:
                # The *previous* bucket just completed — deliver it
                prev_bucket_ts = last_bucket_5m * 300
                pos = idx_5m.get(prev_bucket_ts)
                if pos is not None:
                    new_5m = [all_5m[pos]]
                last_bucket_5m = bucket_5m

            new_15m = []
            bucket_15m = candle_ts // 900
            if bucket_15m != last_bucket_15m:
                prev_bucket_ts = last_bucket_15m * 900
                pos = idx_15m.get(prev_bucket_ts)
                if pos is not None:
                    new_15m = [all_15m[pos]]
                last_bucket_15m = bucket_15m

            new_1h = []
            bucket_1h = candle_ts // 3600
            if bucket_1h != last_bucket_1h:
                prev_bucket_ts = last_bucket_1h * 3600
                pos = idx_1h.get(prev_bucket_ts)
                if pos is not None:
                    new_1h = [all_1h[pos]]
                last_bucket_1h = bucket_1h

            new_4h = []
            bucket_4h = candle_ts // 14400
            if bucket_4h != last_bucket_4h:
                prev_bucket_ts = last_bucket_4h * 14400
                pos = idx_4h.get(prev_bucket_ts)
                if pos is not None:
                    new_4h = [all_4h[pos]]
                last_bucket_4h = bucket_4h

            new_1d = []
            bucket_1d = candle_ts // 86400
            if bucket_1d != last_bucket_1d:
                prev_bucket_ts = last_bucket_1d * 86400
                pos = idx_1d.get(prev_bucket_ts)
                if pos is not None:
                    new_1d = [all_1d[pos]]
                last_bucket_1d = bucket_1d

            tick_data = {
                "symbol": subject,
                "asof_ts": candle_ts,
                "candles_1m": [all_candles[i]],
                "candles_5m": new_5m,
                "candles_15m": new_15m,
                "candles_1h": new_1h,
                "candles_4h": new_4h,
                "candles_1d": new_1d,
                "orderbook": None,
                "funding": _get_funding(candle_ts),
            }
            self.model.tick(tick_data)

            # Predict at intervals
            if last_predict_ts is None or (current_ts - last_predict_ts) >= interval:
                try:
                    output = self.model.predict(
                        symbol=subject,
                        horizon_seconds=horizon_seconds,
                        step_seconds=prediction_interval_seconds,
                    )
                    if not isinstance(output, dict):
                        output = {"expected_return": float(output)}
                except Exception as exc:
                    output = {"expected_return": 0.0, "_error": str(exc)}

                # Find actual outcome at horizon (direct array lookup)
                resolve_idx = i + horizon_steps
                if resolve_idx < len(close_array):
                    entry_price = close_array[i]
                    resolved_price = close_array[resolve_idx]
                    actual_return = (resolved_price - entry_price) / max(abs(entry_price), 1e-9)
                    actual = {
                        "entry_price": entry_price,
                        "resolved_price": resolved_price,
                        "return": actual_return,
                        "direction_up": resolved_price > entry_price,
                    }
                else:
                    actual = None

                # Score
                score_result = None
                expected = output.get("expected_return", 0.0)
                turnover = abs(expected - prev_signal)
                txn_cost = turnover * cost_multiplier

                if actual is not None:
                    try:
                        score_result = self.scoring_fn(output, actual)
                    except Exception as exc:
                        score_result = {"value": 0.0, "success": False, "failed_reason": str(exc)}

                # Deduct transaction cost from gross score
                gross_score = score_result.get("value") if score_result else None
                net_score = (gross_score - txn_cost) if gross_score is not None else None

                predictions.append({
                    "ts": current_ts,
                    "expected_return": expected,
                    "actual_return": actual.get("return", 0.0) if actual else None,
                    "direction_correct": score_result.get("direction_correct") if score_result else None,
                    "score": net_score,
                    "gross_score": gross_score,
                    "txn_cost": txn_cost,
                    "turnover": turnover,
                    "score_success": score_result.get("success", True) if score_result else None,
                })

                prev_signal = expected

                last_predict_ts = current_ts

        # Compute metrics
        metrics = _compute_metrics(predictions)

        config = {
            "source": source,
            "subject": subject,
            "kind": kind,
            "granularity": granularity,
            "start": str(start),
            "end": str(end),
            "window_size": window_size,
            "prediction_interval_seconds": prediction_interval_seconds,
            "horizon_seconds": horizon_seconds,
            "transaction_cost_bps": cost_bps,
        }

        return BacktestResult(predictions=predictions, metrics=metrics, config=config)



# ── Helpers ──


def _df_to_candle_array(df) -> list[dict[str, Any]]:
    """Pre-convert entire DataFrame to a list of candle dicts (once, up front).

    This avoids repeated iterrows() calls during the replay loop.
    """
    ts_col = df["ts_event"].values
    open_col = df["open"].values.astype(float)
    high_col = df["high"].values.astype(float)
    low_col = df["low"].values.astype(float)
    close_col = df["close"].values.astype(float)
    volume_col = df["volume"].values.astype(float)

    candles = []
    for idx in range(len(df)):
        ts = ts_col[idx]
        # Convert numpy datetime64 to unix timestamp
        ts_int = int(ts.astype("datetime64[s]").astype("int64"))
        candles.append({
            "ts": ts_int,
            "open": open_col[idx],
            "high": high_col[idx],
            "low": low_col[idx],
            "close": close_col[idx],
            "volume": volume_col[idx],
        })
    return candles


def _aggregate_candles(
    candles_1m: list[dict[str, Any]],
    target_minutes: int,
    max_output: int,
) -> list[dict[str, Any]]:
    """Roll up 1m candles into higher-timeframe OHLCV bars."""
    if not candles_1m or target_minutes <= 1:
        return candles_1m[-max_output:] if candles_1m else []

    interval_s = target_minutes * 60
    buckets: dict[int, dict[str, Any]] = {}

    for c in candles_1m:
        ts = c.get("ts", 0)
        bucket_ts = (ts // interval_s) * interval_s

        if bucket_ts not in buckets:
            buckets[bucket_ts] = {
                "ts": bucket_ts,
                "open": c["open"],
                "high": c["high"],
                "low": c["low"],
                "close": c["close"],
                "volume": c["volume"],
            }
        else:
            bar = buckets[bucket_ts]
            bar["high"] = max(bar["high"], c["high"])
            bar["low"] = min(bar["low"], c["low"])
            bar["close"] = c["close"]
            bar["volume"] += c["volume"]

    sorted_bars = sorted(buckets.values(), key=lambda b: b["ts"])
    return sorted_bars[-max_output:]


def _compute_metrics(predictions: list[dict[str, Any]]) -> dict[str, float]:
    """Compute rolling window metrics + risk metrics."""
    scored = [
        p for p in predictions
        if p.get("score") is not None and p.get("score_success", True)
    ]

    if not scored:
        return {"score_recent": 0.0, "score_steady": 0.0, "score_anchor": 0.0}

    now = scored[-1]["ts"]
    if not hasattr(now, "timestamp"):
        now = datetime.fromtimestamp(now, tz=timezone.utc)

    windows = {
        "score_recent": timedelta(hours=24),
        "score_steady": timedelta(hours=72),
        "score_anchor": timedelta(hours=168),
    }

    metrics: dict[str, float] = {}
    for name, window in windows.items():
        cutoff = now - window
        window_scores = [p["score"] for p in scored if _ts_ge(p["ts"], cutoff)]
        metrics[name] = sum(window_scores) / len(window_scores) if window_scores else 0.0

    # Mean return (overall)
    all_scores = [p["score"] for p in scored]
    metrics["mean_return"] = sum(all_scores) / len(all_scores) if all_scores else 0.0

    # Hit rate
    directional = [p for p in scored if p.get("direction_correct") is not None]
    if directional:
        correct = sum(1 for p in directional if p["direction_correct"])
        metrics["hit_rate"] = correct / len(directional)

    # Max drawdown (of cumulative score)
    cumulative = 0.0
    peak = 0.0
    max_dd = 0.0
    for p in scored:
        cumulative += p["score"]
        peak = max(peak, cumulative)
        dd = (peak - cumulative) / max(abs(peak), 1e-9) if peak > 0 else 0.0
        max_dd = max(max_dd, dd)
    metrics["max_drawdown"] = -max_dd

    # Sortino ratio
    negative_scores = [p["score"] for p in scored if p["score"] < 0]
    if negative_scores and metrics["mean_return"] != 0:
        from math import sqrt
        downside_dev = sqrt(sum(s * s for s in negative_scores) / len(negative_scores))
        metrics["sortino_ratio"] = metrics["mean_return"] / max(downside_dev, 1e-9)
    else:
        metrics["sortino_ratio"] = 0.0

    # IC (correlation of expected vs actual returns)
    expected = [p.get("expected_return", 0.0) for p in scored]
    actual = [p.get("actual_return", 0.0) for p in scored if p.get("actual_return") is not None]
    if len(expected) == len(actual) and len(expected) > 2:
        metrics["ic"] = _correlation(expected, actual)
    else:
        metrics["ic"] = 0.0

    # Transaction cost metrics
    total_txn_cost = sum(p.get("txn_cost", 0.0) for p in scored)
    total_turnover = sum(p.get("turnover", 0.0) for p in scored)
    total_gross = sum(p.get("gross_score", 0.0) for p in scored if p.get("gross_score") is not None)
    metrics["total_txn_cost"] = total_txn_cost
    metrics["total_turnover"] = total_turnover
    metrics["mean_turnover"] = total_turnover / len(scored) if scored else 0.0
    metrics["gross_return"] = total_gross / len(scored) if scored else 0.0
    metrics["cost_drag"] = total_txn_cost / len(scored) if scored else 0.0

    # Multi-metric enrichment (best-effort — coordinator_node may not be installed)
    try:
        from coordinator_node.metrics.registry import get_default_registry
        from coordinator_node.metrics.context import MetricsContext

        registry = get_default_registry()
        active_metrics = ["ic", "ic_sharpe", "hit_rate", "mean_return",
                          "max_drawdown", "sortino_ratio"]

        pred_dicts = [{"inference_output": {"expected_return": p.get("expected_return", 0.0)}} for p in scored]
        score_dicts = [
            {"result": {"value": p["score"], "actual_return": p.get("actual_return", 0.0)}}
            for p in scored
        ]
        ctx = MetricsContext(model_id="backtest")

        multi = registry.compute(active_metrics, pred_dicts, score_dicts, ctx)
        metrics.update(multi)
    except ImportError:
        pass

    return metrics


def _correlation(x: list[float], y: list[float]) -> float:
    """Pearson correlation between two lists."""
    n = len(x)
    if n < 2:
        return 0.0
    mean_x = sum(x) / n
    mean_y = sum(y) / n
    cov = sum((xi - mean_x) * (yi - mean_y) for xi, yi in zip(x, y)) / n
    std_x = (sum((xi - mean_x) ** 2 for xi in x) / n) ** 0.5
    std_y = (sum((yi - mean_y) ** 2 for yi in y) / n) ** 0.5
    if std_x < 1e-12 or std_y < 1e-12:
        return 0.0
    return cov / (std_x * std_y)


def _ts_ge(ts, cutoff: datetime) -> bool:
    if hasattr(ts, "timestamp"):
        return ts >= cutoff
    return datetime.fromtimestamp(ts, tz=timezone.utc) >= cutoff


def _default_scoring_fn() -> Callable[[dict, dict], dict]:
    try:
        from vanta.scoring import score_prediction
        return score_prediction
    except ImportError:
        pass

    def _basic_score(prediction: dict, ground_truth: dict) -> dict:
        expected = float(prediction.get("expected_return", 0.0))
        actual = float(ground_truth.get("return", 0.0))
        return {
            "value": expected * actual,
            "expected_return": expected,
            "actual_return": actual,
            "gross_pnl": expected * actual,
            "direction_correct": (expected > 0 and actual > 0) or (expected < 0 and actual < 0) or expected == 0,
            "success": True,
            "failed_reason": None,
        }
    return _basic_score


def _parse_date(value) -> datetime:
    if isinstance(value, datetime):
        return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value
    for fmt in ("%Y-%m-%d", "%Y-%m-%dT%H:%M:%S", "%Y-%m-%dT%H:%M:%SZ"):
        try:
            return datetime.strptime(value, fmt).replace(tzinfo=timezone.utc)
        except ValueError:
            continue
    raise ValueError(f"Cannot parse date: {value!r}")


def _safe_float(value) -> float | None:
    if value is None:
        return None
    try:
        return float(value)
    except (ValueError, TypeError):
        return None


# ── Compare & Sweep ───────────────────────────────────────────────────────────


def compare(
    models: list[tuple[str, Any]],
    source: str | None = None,
    subject: str | None = None,
    kind: str | None = None,
    granularity: str | None = None,
    start: str | datetime = "2025-03-01",
    end: str | datetime = "2026-02-01",
    window_size: int = 120,
    prediction_interval_seconds: int = 900,
    horizon_seconds: int = 3600,
    transaction_cost_bps: float = 0.0,
    plot: bool = True,
) -> dict[str, BacktestResult]:
    """Run multiple models on the same data and compare results.

    Args:
        models: List of (name, model_instance) tuples.
        transaction_cost_bps: Cost in basis points charged on signal turnover.
        plot: If True, shows overlaid equity curves and comparison table.

    Returns:
        Dict of name → BacktestResult.

    Example:
        from vanta.backtest import compare
        from vanta.examples import MeanReversionTracker, TrendFollowingTracker

        results = compare([
            ("Mean Reversion", MeanReversionTracker()),
            ("Trend Following", TrendFollowingTracker()),
        ], start="2025-06-01", end="2026-01-01", transaction_cost_bps=4)
    """
    results: dict[str, BacktestResult] = {}
    run_kwargs = dict(
        source=source, subject=subject, kind=kind, granularity=granularity,
        start=start, end=end, window_size=window_size,
        prediction_interval_seconds=prediction_interval_seconds,
        horizon_seconds=horizon_seconds,
        transaction_cost_bps=transaction_cost_bps,
    )

    for name, model in models:
        print(f"Running: {name}...")
        runner = BacktestRunner(model=model)
        results[name] = runner.run(**run_kwargs)

    # Print comparison table
    _print_comparison_table(results)

    if plot:
        _plot_comparison(results)

    return results


def sweep(
    model_fn,
    params: dict[str, list],
    source: str | None = None,
    subject: str | None = None,
    kind: str | None = None,
    granularity: str | None = None,
    start: str | datetime = "2025-03-01",
    end: str | datetime = "2026-02-01",
    window_size: int = 120,
    prediction_interval_seconds: int = 900,
    horizon_seconds: int = 3600,
    transaction_cost_bps: float = 0.0,
    plot: bool = True,
) -> dict[str, BacktestResult]:
    """Sweep parameters for a model and compare results.

    Args:
        model_fn: Callable that takes keyword args and returns a model instance.
        params: Dict of param_name → list of values to try.

    Example:
        from vanta.backtest import sweep
        from vanta.examples import MeanReversionTracker

        results = sweep(
            model_fn=lambda lookback: MeanReversionTracker(lookback=lookback),
            params={"lookback": [5, 10, 20, 50]},
            start="2025-06-01", end="2026-01-01",
        )
    """
    import itertools

    param_names = list(params.keys())
    param_values = list(params.values())

    models: list[tuple[str, Any]] = []
    for combo in itertools.product(*param_values):
        kwargs = dict(zip(param_names, combo))
        label = ", ".join(f"{k}={v}" for k, v in kwargs.items())
        models.append((label, model_fn(**kwargs)))

    return compare(
        models, source=source, subject=subject, kind=kind, granularity=granularity,
        start=start, end=end, window_size=window_size,
        prediction_interval_seconds=prediction_interval_seconds,
        horizon_seconds=horizon_seconds,
        transaction_cost_bps=transaction_cost_bps, plot=plot,
    )


def _print_comparison_table(results: dict[str, BacktestResult]) -> None:
    """Print a formatted comparison table of backtest results."""
    if not results:
        return

    # Include cost columns if any result has non-zero transaction costs
    has_costs = any(
        r.config.get("transaction_cost_bps", 0) > 0 for r in results.values()
    )
    metric_keys = ["mean_return", "hit_rate", "ic", "sortino_ratio", "max_drawdown",
                    "score_recent", "score_steady", "score_anchor"]
    if has_costs:
        metric_keys = ["mean_return", "gross_return", "cost_drag",
                       "hit_rate", "ic", "sortino_ratio", "max_drawdown"]

    # Header
    name_width = max(len(n) for n in results) + 2
    header = f"{'Strategy':<{name_width}}"
    for key in metric_keys:
        header += f"  {key:>14s}"
    print("─" * len(header))
    print(header)
    print("─" * len(header))

    # Rows sorted by mean_return descending
    sorted_names = sorted(results, key=lambda n: results[n].metrics.get("mean_return", 0), reverse=True)
    for name in sorted_names:
        row = f"{name:<{name_width}}"
        for key in metric_keys:
            val = results[name].metrics.get(key, 0.0)
            row += f"  {val:>14.6f}"
        print(row)
    print("─" * len(header))
    print()


def _plot_comparison(results: dict[str, BacktestResult]) -> None:
    """Plot overlaid equity curves for compared strategies."""
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        print("Install matplotlib for plotting: pip install matplotlib")
        return

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))

    for name, result in results.items():
        scored = [p for p in result._predictions if p.get("score") is not None]
        if not scored:
            continue

        timestamps = [p["ts"] for p in scored]
        cumulative = []
        running = 0.0
        for p in scored:
            running += p["score"]
            cumulative.append(running)

        ax1.plot(timestamps, cumulative, label=name, linewidth=1)

        # Drawdown
        peak = 0.0
        drawdowns = []
        for c in cumulative:
            peak = max(peak, c)
            drawdowns.append(c - peak)
        ax2.plot(timestamps, drawdowns, label=name, linewidth=1)

    ax1.set_title("Cumulative PnL Comparison")
    ax1.set_xlabel("Time")
    ax1.set_ylabel("Cumulative Score")
    ax1.legend(fontsize=8)
    ax1.grid(True, alpha=0.3)
    ax1.axhline(y=0, color="gray", linestyle="--", linewidth=0.5)

    ax2.set_title("Drawdown Comparison")
    ax2.set_xlabel("Time")
    ax2.set_ylabel("Drawdown")
    ax2.legend(fontsize=8)
    ax2.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.show()
