"""Signal Quality Analysis: Score-bucket monotonicity check.

Does higher composite score predict better forward returns?

Reads snapshot parquet, computes compute_entry_score for each row,
buckets into quintiles, and checks Spearman rank correlation between
score bucket and mean forward return.

Usage:
  python -m dev.analysis.signal_quality_analysis \
    --snapshot-dir data/datasets/snapshots/snapshot_2026_03_20 \
    --split train [--output-dir ./data/analysis]
"""
from __future__ import annotations

import argparse
import csv
import statistics
from pathlib import Path
from typing import Any

import pyarrow.parquet as pq

from libs.backtest.scoring import compute_entry_score


def _load_rows(snapshot_dir: Path, split: str) -> list[dict[str, Any]]:
    """Load parquet snapshot into list of dicts."""
    parquet_path = snapshot_dir / f"{split}.parquet"
    table = pq.read_table(parquet_path)
    return table.to_pylist()


def _bucket_label(score: float) -> str:
    """Assign score to a quintile bucket."""
    if score < 0.2:
        return "0.0-0.2"
    if score < 0.4:
        return "0.2-0.4"
    if score < 0.6:
        return "0.4-0.6"
    if score < 0.8:
        return "0.6-0.8"
    return "0.8-1.0"


BUCKET_ORDER = ["0.0-0.2", "0.2-0.4", "0.4-0.6", "0.6-0.8", "0.8-1.0"]
FWD_HORIZONS = ["fwd_return_1d", "fwd_return_3d", "fwd_return_5d", "fwd_return_10d"]


def analyze(rows: list[dict[str, Any]], output_dir: Path) -> None:
    """Score all rows, bucket, compute statistics, write output."""
    output_dir.mkdir(parents=True, exist_ok=True)

    # Score each row and assign bucket
    scored: list[dict[str, Any]] = []
    for row in rows:
        score = compute_entry_score(row)
        bucket = _bucket_label(score)
        scored.append({"score": score, "bucket": bucket, **row})

    # Group by bucket
    by_bucket: dict[str, list[dict[str, Any]]] = {b: [] for b in BUCKET_ORDER}
    for s in scored:
        by_bucket[s["bucket"]].append(s)

    # Compute per-bucket stats
    csv_rows: list[dict[str, Any]] = []
    for bucket in BUCKET_ORDER:
        items = by_bucket[bucket]
        n = len(items)
        row_out: dict[str, Any] = {"bucket": bucket, "count": n}

        for horizon in FWD_HORIZONS:
            vals = [float(r[horizon]) for r in items if r.get(horizon) is not None]
            if vals:
                row_out[f"{horizon}_mean"] = statistics.mean(vals)
                row_out[f"{horizon}_win_rate"] = sum(1 for v in vals if v > 0) / len(vals)
            else:
                row_out[f"{horizon}_mean"] = None
                row_out[f"{horizon}_win_rate"] = None

        csv_rows.append(row_out)

    # Spearman rank correlation (bucket rank vs mean return)
    spearman_results: dict[str, float | None] = {}
    for horizon in FWD_HORIZONS:
        means = [r[f"{horizon}_mean"] for r in csv_rows if r[f"{horizon}_mean"] is not None]
        if len(means) >= 3:
            # Rank correlation: do the means increase with bucket rank?
            ranks_score = list(range(len(means)))
            ranks_return = _rank(means)
            spearman_results[horizon] = _spearman(ranks_score, ranks_return)
        else:
            spearman_results[horizon] = None

    # Write CSV
    csv_path = output_dir / "signal_quality_analysis.csv"
    if csv_rows:
        fieldnames = list(csv_rows[0].keys())
        with open(csv_path, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            writer.writerows(csv_rows)

    # Console output
    print(f"\n{'='*80}")
    print(f"Signal Quality Analysis — {len(scored)} events")
    print(f"{'='*80}")
    print(f"\n{'Bucket':<12} {'Count':>6}", end="")
    for h in FWD_HORIZONS:
        label = h.replace("fwd_return_", "")
        print(f"  {label+'_mean':>10} {label+'_wr':>8}", end="")
    print()
    print("-" * 80)

    for row_out in csv_rows:
        print(f"{row_out['bucket']:<12} {row_out['count']:>6}", end="")
        for h in FWD_HORIZONS:
            mean_val = row_out[f"{h}_mean"]
            wr_val = row_out[f"{h}_win_rate"]
            mean_str = f"{mean_val:+.4f}" if mean_val is not None else "    N/A"
            wr_str = f"{wr_val:.1%}" if wr_val is not None else "  N/A"
            print(f"  {mean_str:>10} {wr_str:>8}", end="")
        print()

    print(f"\nSpearman rank correlation (bucket vs mean return):")
    for h in FWD_HORIZONS:
        label = h.replace("fwd_return_", "")
        val = spearman_results.get(h)
        print(f"  {label}: {val:+.3f}" if val is not None else f"  {label}: N/A")

    print(f"\nCSV written to: {csv_path}")


def _rank(values: list[float]) -> list[float]:
    """Compute ranks (0-indexed) for a list of values."""
    indexed = sorted(enumerate(values), key=lambda x: x[1])
    ranks = [0.0] * len(values)
    for rank, (idx, _) in enumerate(indexed):
        ranks[idx] = float(rank)
    return ranks


def _spearman(x: list[float], y: list[float]) -> float:
    """Compute Spearman rank correlation coefficient."""
    n = len(x)
    if n < 2:
        return 0.0
    d_sq = sum((xi - yi) ** 2 for xi, yi in zip(x, y))
    return 1 - (6 * d_sq) / (n * (n**2 - 1))


def main() -> None:
    parser = argparse.ArgumentParser(description="Signal Quality Analysis")
    parser.add_argument("--snapshot-dir", required=True, help="Path to snapshot directory")
    parser.add_argument("--split", default="train", help="Split name (train/valid/test)")
    parser.add_argument("--output-dir", default="./data/analysis", help="Output directory")
    args = parser.parse_args()

    rows = _load_rows(Path(args.snapshot_dir), args.split)
    if not rows:
        print("No data found.")
        return

    analyze(rows, Path(args.output_dir))


if __name__ == "__main__":
    main()