"""Build a lightweight bucket-prior ranking model from snapshot labels."""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import pandas as pd

from libs.backtest.ranking_models import _resolve_feature_key


FEATURE_SPECS = [
    ("event_type", 0.08),
    ("direction_guidance_combo", 0.27),
    ("reaction_bucket", 0.22),
    ("close_bucket", 0.15),
    ("volume_bucket", 0.12),
    ("gap_bucket", 0.08),
    ("document_bucket", 0.05),
    ("confidence_bucket", 0.03),
]


def build_model(
    snapshot_path: Path,
    cutoff_event_date: str,
    min_bucket_count: int,
) -> dict:
    df = pd.read_parquet(snapshot_path)
    df["event_date"] = pd.to_datetime(df["event_date"]).dt.date
    df["reaction_date"] = pd.to_datetime(df["reaction_date"], errors="coerce").dt.date
    cutoff_date = pd.to_datetime(cutoff_event_date).date()
    df = df[
        (df["event_date"] <= cutoff_date)
        & (df["reaction_date"] == df["event_date"])
        & (df["reaction_day_return"] > 0)
        & df["event_type"].isin(["earnings_release", "guidance_update"])
        & df["fwd_return_20d"].notna()
    ].copy()
    global_mean = float(df["fwd_return_20d"].mean()) if not df.empty else 0.0

    features = []
    for name, weight in FEATURE_SPECS:
        bucket_map: dict[str, list[float]] = {}
        for row in df.to_dict("records"):
            key = _resolve_feature_key(name, row)
            if key is None:
                continue
            bucket_map.setdefault(key, []).append(float(row["fwd_return_20d"]))
        values = {
            key: sum(values) / len(values)
            for key, values in bucket_map.items()
            if len(values) >= min_bucket_count
        }
        features.append(
            {
                "name": name,
                "weight": weight,
                "values": values,
            }
        )

    return {
        "model_type": "bucket_blend_v1",
        "target": "fwd_return_20d",
        "snapshot_id": snapshot_path.parent.name,
        "cutoff_event_date": cutoff_event_date,
        "min_bucket_count": min_bucket_count,
        "training_rows": int(len(df)),
        "global_mean": global_mean,
        "features": features,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Build a bucket-prior ranking model")
    parser.add_argument("--snapshot", required=True, help="Path to the parquet snapshot")
    parser.add_argument("--output", required=True, help="Output JSON path")
    parser.add_argument("--cutoff-event-date", default="2024-12-31")
    parser.add_argument("--min-bucket-count", type=int, default=8)
    args = parser.parse_args()

    model = build_model(
        snapshot_path=Path(args.snapshot),
        cutoff_event_date=args.cutoff_event_date,
        min_bucket_count=args.min_bucket_count,
    )
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(model, indent=2))
    print(output)


if __name__ == "__main__":
    main()