"""MFE/MAE Distribution Analysis by Event Type and Horizon.

Reads EventLabel records from the database and analyzes:
- MFE/MAE distributions per event_type per horizon (3d, 5d, 10d, 20d)
- Optimal holding period per event type (where MFE peaks)
- Forward return distributions

Outputs:
- Matplotlib plots (saved to output directory)
- Summary CSV with statistics per event_type × horizon

Usage:
  python -m dev.analysis.label_horizon_analysis [--output-dir ./data/analysis]
"""
from __future__ import annotations

import argparse
import asyncio
import csv
import datetime as dt
from pathlib import Path
from typing import Any

from libs.common.logging import configure_logging, get_logger

logger = get_logger(__name__)

HORIZONS = ["3d", "5d", "10d", "20d"]


async def _load_label_data() -> list[dict[str, Any]]:
    """Load EventLabel + Event data from DB."""
    from sqlalchemy import select
    from libs.db.models import Event, EventLabel
    from libs.db.session import get_session

    rows: list[dict[str, Any]] = []
    async with get_session() as session:
        result = await session.execute(
            select(EventLabel, Event.event_type)
            .join(Event, EventLabel.event_id == Event.event_id)
            .where(EventLabel.label_status == "ok")
            .where(EventLabel.invalid_event_for_labeling.is_(False))
        )
        for lbl, event_type in result.all():
            rows.append({
                "event_id": lbl.event_id,
                "event_type": event_type,
                "fwd_return_3d": float(lbl.fwd_return_3d) if lbl.fwd_return_3d is not None else None,
                "fwd_return_5d": float(lbl.fwd_return_5d) if lbl.fwd_return_5d is not None else None,
                "fwd_return_10d": float(lbl.fwd_return_10d) if lbl.fwd_return_10d is not None else None,
                "fwd_return_20d": float(lbl.fwd_return_20d) if lbl.fwd_return_20d is not None else None,
                "mfe_3d": float(lbl.mfe_3d) if lbl.mfe_3d is not None else None,
                "mae_3d": float(lbl.mae_3d) if lbl.mae_3d is not None else None,
                "mfe_5d": float(lbl.mfe_5d) if lbl.mfe_5d is not None else None,
                "mae_5d": float(lbl.mae_5d) if lbl.mae_5d is not None else None,
                "mfe_10d": float(lbl.mfe_10d) if lbl.mfe_10d is not None else None,
                "mae_10d": float(lbl.mae_10d) if lbl.mae_10d is not None else None,
                "mfe_20d": float(lbl.mfe_20d) if lbl.mfe_20d is not None else None,
                "mae_20d": float(lbl.mae_20d) if lbl.mae_20d is not None else None,
            })
    return rows


def _compute_stats(values: list[float]) -> dict[str, float | None]:
    """Compute summary statistics for a list of values."""
    if not values:
        return {"count": 0, "mean": None, "median": None, "std": None, "min": None, "max": None}
    import statistics
    return {
        "count": len(values),
        "mean": statistics.mean(values),
        "median": statistics.median(values),
        "std": statistics.stdev(values) if len(values) > 1 else 0.0,
        "min": min(values),
        "max": max(values),
    }


def analyze_and_write(rows: list[dict[str, Any]], output_dir: Path) -> None:
    """Compute statistics and generate plots."""
    output_dir.mkdir(parents=True, exist_ok=True)

    # Group by event_type
    by_type: dict[str, list[dict[str, Any]]] = {}
    for r in rows:
        et = r["event_type"]
        by_type.setdefault(et, []).append(r)

    # Summary CSV
    csv_rows: list[dict[str, Any]] = []
    optimal_holding: dict[str, str] = {}

    for event_type, type_rows in sorted(by_type.items()):
        mfe_means: dict[str, float] = {}
        for horizon in HORIZONS:
            mfe_vals = [r[f"mfe_{horizon}"] for r in type_rows if r.get(f"mfe_{horizon}") is not None]
            mae_vals = [r[f"mae_{horizon}"] for r in type_rows if r.get(f"mae_{horizon}") is not None]
            fwd_vals = [r[f"fwd_return_{horizon}"] for r in type_rows if r.get(f"fwd_return_{horizon}") is not None]

            mfe_stats = _compute_stats(mfe_vals)
            mae_stats = _compute_stats(mae_vals)
            fwd_stats = _compute_stats(fwd_vals)

            if mfe_stats["mean"] is not None:
                mfe_means[horizon] = mfe_stats["mean"]

            csv_rows.append({
                "event_type": event_type,
                "horizon": horizon,
                "n": mfe_stats["count"],
                "mfe_mean": mfe_stats["mean"],
                "mfe_median": mfe_stats["median"],
                "mae_mean": mae_stats["mean"],
                "mae_median": mae_stats["median"],
                "fwd_return_mean": fwd_stats["mean"],
                "fwd_return_median": fwd_stats["median"],
                "fwd_return_std": fwd_stats["std"],
                "reward_risk_ratio": (
                    mfe_stats["mean"] / abs(mae_stats["mean"])
                    if mfe_stats["mean"] is not None and mae_stats["mean"] and abs(mae_stats["mean"]) > 0
                    else None
                ),
            })

        # Optimal holding period: horizon with highest mean MFE
        if mfe_means:
            best_h = max(mfe_means, key=lambda h: mfe_means[h])
            optimal_holding[event_type] = best_h

    # Write CSV
    csv_path = output_dir / "mfe_mae_summary.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)
        logger.info("summary_csv_written", path=str(csv_path), rows=len(csv_rows))

    # Print summary table
    print(f"\n{'='*80}")
    print(f"MFE/MAE Distribution Analysis — {len(rows)} events, {len(by_type)} event types")
    print(f"{'='*80}")
    for event_type in sorted(by_type):
        n = len(by_type[event_type])
        best_h = optimal_holding.get(event_type, "N/A")
        print(f"\n  {event_type} (n={n}, optimal holding={best_h})")
        for horizon in HORIZONS:
            matching = [r for r in csv_rows if r["event_type"] == event_type and r["horizon"] == horizon]
            if matching:
                m = matching[0]
                mfe_str = f"{m['mfe_mean']:+.2%}" if m["mfe_mean"] is not None else "N/A"
                mae_str = f"{m['mae_mean']:+.2%}" if m["mae_mean"] is not None else "N/A"
                fwd_str = f"{m['fwd_return_mean']:+.2%}" if m["fwd_return_mean"] is not None else "N/A"
                rr_str = f"{m['reward_risk_ratio']:.2f}" if m["reward_risk_ratio"] is not None else "N/A"
                print(f"    {horizon}: MFE={mfe_str}  MAE={mae_str}  FwdRet={fwd_str}  R:R={rr_str}")

    # Generate plots (if matplotlib available)
    try:
        _generate_plots(by_type, output_dir)
    except ImportError:
        logger.info("matplotlib_not_available", msg="Skipping plot generation")
        print("\n(matplotlib not available — skipping plot generation)")


def _generate_plots(
    by_type: dict[str, list[dict[str, Any]]],
    output_dir: Path,
) -> None:
    """Generate MFE/MAE distribution plots using matplotlib."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    event_types = sorted(by_type.keys())

    # Plot 1: MFE by horizon per event type
    fig, axes = plt.subplots(
        len(event_types), len(HORIZONS),
        figsize=(4 * len(HORIZONS), 3 * len(event_types)),
        squeeze=False,
    )
    fig.suptitle("MFE Distribution by Event Type and Horizon", fontsize=14)

    for i, et in enumerate(event_types):
        for j, h in enumerate(HORIZONS):
            ax = axes[i][j]
            vals = [r[f"mfe_{h}"] for r in by_type[et] if r.get(f"mfe_{h}") is not None]
            if vals:
                ax.hist(vals, bins=max(5, len(vals) // 3), alpha=0.7, color="steelblue", edgecolor="white")
                ax.axvline(x=sum(vals) / len(vals), color="red", linestyle="--", linewidth=1)
            ax.set_title(f"{et}\n{h}" if i == 0 else h, fontsize=8)
            if j == 0:
                ax.set_ylabel(et[:15], fontsize=8)
            ax.tick_params(labelsize=6)

    plt.tight_layout()
    mfe_path = output_dir / "mfe_distributions.png"
    fig.savefig(str(mfe_path), dpi=150)
    plt.close(fig)
    print(f"\nMFE plot saved: {mfe_path}")

    # Plot 2: Mean MFE across horizons (optimal holding period)
    fig2, ax2 = plt.subplots(figsize=(10, 6))
    horizon_days = [3, 5, 10, 20]
    for et in event_types:
        means = []
        for h in HORIZONS:
            vals = [r[f"mfe_{h}"] for r in by_type[et] if r.get(f"mfe_{h}") is not None]
            means.append(sum(vals) / len(vals) * 100 if vals else 0)
        ax2.plot(horizon_days, means, marker="o", label=et)
    ax2.set_xlabel("Horizon (trading days)")
    ax2.set_ylabel("Mean MFE (%)")
    ax2.set_title("Optimal Holding Period — Mean MFE by Horizon")
    ax2.legend(fontsize=8, loc="upper left")
    ax2.grid(True, alpha=0.3)
    plt.tight_layout()
    optimal_path = output_dir / "optimal_holding_period.png"
    fig2.savefig(str(optimal_path), dpi=150)
    plt.close(fig2)
    print(f"Optimal holding plot saved: {optimal_path}")


def main() -> None:
    parser = argparse.ArgumentParser(description="MFE/MAE Distribution Analysis")
    parser.add_argument(
        "--output-dir",
        default="./data/analysis",
        help="Output directory for plots and CSV (default: ./data/analysis)",
    )
    args = parser.parse_args()

    configure_logging("info")
    rows = asyncio.run(_load_label_data())

    if not rows:
        print("No label data found. Run the labeler pipeline first.")
        return

    analyze_and_write(rows, Path(args.output_dir))


if __name__ == "__main__":
    main()
