#!/usr/bin/env python3
"""C10: grid-search asymmetric triple barrier pt_ratio to maximize UP precision."""
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

try:
    from dotenv import load_dotenv; load_dotenv(ROOT / ".env")
except ImportError:
    pass

import warnings; warnings.filterwarnings("ignore")
import json
import numpy as np
import pandas as pd

from scripts.improvement_harness import (
    DEFAULT_STOCKS, BASELINE_FEATURES, LABEL_HORIZON, VOL_LOOKBACK,
    RF_PARAMS, fetch_df, walk_forward,
)
from models.predictor import _build_features

PT_RATIOS = [1.0, 1.5, 2.0, 2.5, 3.0]
SL_RATIO  = 1.0


def build_asymmetric_labels(close, pt_ratio=1.0, sl_ratio=1.0,
                             horizon=5, vol_lookback=20):
    """pt_ratio: multiplier for take-profit barrier (UP label)
       sl_ratio: multiplier for stop-loss barrier (DOWN label)
    """
    n = len(close)
    log_ret = np.diff(np.log(close + 1e-9))
    labels = np.full(n, np.nan)
    for i in range(n - horizon):
        start_v = max(0, i - vol_lookback)
        window = log_ret[start_v:i]
        vol = float(np.std(window)) if len(window) >= 5 else 0.015
        vol = max(vol, 0.001)
        upper = close[i] * (1.0 + vol * pt_ratio)
        lower = close[i] * (1.0 - vol * sl_ratio)
        label = 0
        for j in range(1, horizon + 1):
            c = close[i + j]
            if c >= upper: label = 1; break
            if c <= lower: label = -1; break
        labels[i] = label
    return labels


def _mean(rows, field):
    vals = [r[field] for r in rows
            if isinstance(r.get(field), (int, float)) and not np.isnan(r.get(field, float("nan")))]
    return round(sum(vals) / len(vals), 1) if vals else float("nan")


def main():
    # Pre-fetch all stock data once
    print("Fetching stock data...")
    stock_data = {}
    for stock_no in DEFAULT_STOCKS:
        print(f"  {stock_no}...", end=" ", flush=True)
        df = fetch_df(stock_no)
        if df is None or df.empty:
            print("no data")
            continue
        feat = _build_features(df)
        close = (df.set_index("date")["close"] if "date" in df.columns else df["close"]).values
        stock_data[stock_no] = (feat, close)
        print("ok")

    grid_results = {}

    hdr = f"{'pt_ratio':>8}  {'stock':>6}  {'dir%':>5}  {'↑prec%':>7}  {'signals':>7}  {'n_pred':>7}"
    print(f"\n{hdr}\n{'-'*len(hdr)}")

    for pt_ratio in PT_RATIOS:
        ratio_key = str(pt_ratio)
        per_stock_results = {}
        rows = []

        for stock_no, (feat, close) in stock_data.items():
            labels = build_asymmetric_labels(
                close, pt_ratio=pt_ratio, sl_ratio=SL_RATIO,
                horizon=LABEL_HORIZON, vol_lookback=VOL_LOOKBACK,
            )
            r = walk_forward(feat, labels, BASELINE_FEATURES)
            per_stock_results[stock_no] = r
            if r:
                rows.append(r)
                print(f"{pt_ratio:>8.1f}  {stock_no:>6}  {r['dir_accuracy']:>5.1f}  "
                      f"{r['up_precision']:>7.1f}  {r.get('n_signals',0):>7}  "
                      f"{r.get('n_predictions',0):>7}")

        agg = {k: _mean(rows, k) for k in ("dir_accuracy", "up_precision", "n_signals", "n_predictions")}
        grid_results[ratio_key] = {"results": per_stock_results, "aggregate": agg}
        print(f"  {'→ avg':>14}  {agg['dir_accuracy']:>5.1f}  {agg['up_precision']:>7.1f}  "
              f"{agg['n_signals']:>7}\n")

    # Aggregate table
    print("=== AGGREGATE BY PT_RATIO ===")
    agg_hdr = f"{'pt_ratio':>8}  {'avg_dir%':>8}  {'avg_↑prec%':>10}  {'avg_signals':>11}"
    print(f"{agg_hdr}\n{'-'*len(agg_hdr)}")
    for pt_ratio in PT_RATIOS:
        a = grid_results[str(pt_ratio)]["aggregate"]
        print(f"{pt_ratio:>8.1f}  {a['dir_accuracy']:>8.1f}  {a['up_precision']:>10.1f}  {a['n_signals']:>11.1f}")

    # Best: highest up_precision where avg_signals >= 15
    baseline_agg = grid_results["1.0"]["aggregate"]
    best_ratio = None
    best_up_prec = -1.0
    for pt_ratio in PT_RATIOS:
        a = grid_results[str(pt_ratio)]["aggregate"]
        if a["n_signals"] >= 15 and a["up_precision"] > best_up_prec:
            best_up_prec = a["up_precision"]
            best_ratio = pt_ratio

    best_agg = grid_results[str(best_ratio)]["aggregate"] if best_ratio is not None else {}
    dir_regression = baseline_agg["dir_accuracy"] - best_agg.get("dir_accuracy", 0) if best_ratio else 999.0

    passed = (
        best_ratio is not None
        and best_up_prec >= 57.0
        and best_agg.get("n_signals", 0) >= 15
        and dir_regression <= 3.0
    )

    output = {
        "grid": grid_results,
        "best_pt_ratio": best_ratio,
        "best_aggregate": best_agg,
        "baseline_aggregate": baseline_agg,
        "dir_regression_vs_baseline": round(dir_regression, 1) if best_ratio else None,
        "passed": passed,
        "pass_criterion": "up_precision >= 57.0 AND avg_signals >= 15 AND dir_regression <= 3.0pp",
    }

    def _clean(obj):
        if isinstance(obj, dict):
            return {k: _clean(v) for k, v in obj.items()}
        if isinstance(obj, list):
            return [_clean(v) for v in obj]
        if isinstance(obj, float) and (np.isnan(obj) or np.isinf(obj)):
            return None
        if isinstance(obj, (np.integer,)):
            return int(obj)
        if isinstance(obj, (np.floating,)):
            return None if np.isnan(obj) else float(obj)
        if isinstance(obj, (np.bool_,)):
            return bool(obj)
        if hasattr(obj, 'item'):  # any remaining numpy scalar
            return obj.item()
        return obj

    out_path = ROOT / "docs" / "c10_barrier_result.json"
    with open(out_path, "w") as f:
        json.dump(_clean(output), f, indent=2)

    print(f"\n=== RESULT ===")
    print(f"  best_pt_ratio={best_ratio}  up_precision={best_up_prec}  "
          f"n_signals={best_agg.get('n_signals','?')}  dir_regression={round(dir_regression,1) if best_ratio else 'N/A'}")
    print(f"  PASSED: {passed}")
    print(f"  Output: {out_path}")


if __name__ == "__main__":
    main()