#!/usr/bin/env python3
"""C20: Sector-conditional model — pool training data within sector.

Instead of training per-stock in isolation, each stock's walk-forward window
is augmented with date-aligned rows from its sector peers.  The sector model
then predicts only the target stock's test window.

Sector groups:
  semis:       2330, 2454, 2303
  electronics: 2317, 2382, 2308
  financials:  2881, 2882, 2886
  etfs:        0050, 0056
  telecom:     2412  (solo — no peer benefit, same as baseline)

Hypothesis: pooling intra-sector reduces label noise from cross-sector mixing
(C9 failure) while giving the model more examples of sector-specific patterns
that hurt the worst stocks (2454 semi, 2412 telecom, 2886 financial).
"""
import sys
from pathlib import Path

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

import warnings; warnings.filterwarnings("ignore")
import json
import argparse
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler

from scripts.improvement_harness import (
    fetch_df, build_triple_barrier_labels, compute_metrics,
    CURRENT_FEATURES, DEFAULT_STOCKS, EXTENDED_STOCKS,
    PASS_DIR_ACC, PASS_UP_PREC,
    MIN_TRAIN, STEP, LABEL_HORIZON, RF_PARAMS,
)
from models.predictor import _build_features

# All stocks needed for sector pools (superset of both eval sets)
ALL_STOCKS = list(dict.fromkeys(DEFAULT_STOCKS + EXTENDED_STOCKS))

SECTOR_MAP = {
    "semis":       ["2330", "2454", "2303"],
    "electronics": ["2317", "2382", "2308"],
    "financials":  ["2881", "2882", "2886"],
    "etfs":        ["0050", "0056"],
    "telecom":     ["2412"],
}

# Sectors where intra-sector businesses are too heterogeneous to benefit from pooling.
# electronics: Foxconn (contract mfg) / Delta (power components) / Quanta (laptop ODM)
# are structurally different — pooling adds noise, not signal.
NO_POOL_SECTORS = {"electronics", "telecom"}

# Reverse map: stock → sector name
STOCK_SECTOR = {s: sec for sec, stocks in SECTOR_MAP.items() for s in stocks}


def _get_close(df):
    return (df.set_index("date")["close"] if "date" in df.columns else df["close"]).values


def _get_dates(df):
    if "date" in df.columns:
        return df["date"].values
    return None


def walk_forward_sector(target_no, target_feat, target_labels, target_dates,
                        peers: list[tuple], cols):
    """Walk-forward on target stock, training augmented with sector peer rows."""
    avail = [c for c in cols if c in target_feat.columns]
    X_target = target_feat[avail].fillna(0).values
    n = len(target_feat)

    y_true_all, y_pred_all = [], []

    cutoff = MIN_TRAIN
    while cutoff + STEP + LABEL_HORIZON <= n:
        train_end = cutoff - LABEL_HORIZON
        if train_end < MIN_TRAIN - LABEL_HORIZON:
            cutoff += STEP; continue

        y_tr  = target_labels[:train_end]
        valid = ~np.isnan(y_tr)
        y_v   = y_tr[valid].astype(int)
        if len(y_v) < 10 or len(np.unique(y_v)) < 2:
            cutoff += STEP; continue

        X_train_list = [X_target[:train_end][valid]]
        y_train_list = [y_v]

        # Cutoff date for peer alignment (avoid future leakage)
        cutoff_date = target_dates[train_end - 1] if (target_dates is not None and train_end > 0) else None

        for (peer_feat, peer_labels, peer_dates) in peers:
            peer_avail = [c for c in cols if c in peer_feat.columns]

            if cutoff_date is not None and peer_dates is not None:
                peer_end = int((peer_dates <= cutoff_date).sum())
            else:
                peer_end = int(train_end * len(peer_feat) / n)

            peer_end = min(peer_end, len(peer_feat) - LABEL_HORIZON)
            if peer_end < 15:
                continue

            y_p   = peer_labels[:peer_end]
            valid_p = ~np.isnan(y_p)
            y_vp  = y_p[valid_p].astype(int)
            if len(y_vp) < 5 or len(np.unique(y_vp)) < 2:
                continue

            X_p = peer_feat[peer_avail].fillna(0).values[:peer_end][valid_p]
            X_train_list.append(X_p)
            y_train_list.append(y_vp)

        X_train = np.vstack(X_train_list)
        y_train = np.concatenate(y_train_list)
        if len(np.unique(y_train)) < 2:
            cutoff += STEP; continue

        scaler = StandardScaler()
        X_train_s = scaler.fit_transform(X_train)

        rf = RandomForestClassifier(**RF_PARAMS)
        rf.fit(X_train_s, y_train)

        test_end = min(cutoff + STEP, n - LABEL_HORIZON)
        y_te     = target_labels[cutoff:test_end]
        valid_te = ~np.isnan(y_te)
        if valid_te.sum() == 0:
            cutoff += STEP; continue

        X_te_s = scaler.transform(X_target[cutoff:test_end][valid_te])
        y_pred = rf.predict(X_te_s)

        y_true_all.extend(y_te[valid_te].astype(int).tolist())
        y_pred_all.extend(y_pred.tolist())
        cutoff += STEP

    if not y_true_all:
        return {}
    return compute_metrics(np.array(y_true_all), np.array(y_pred_all))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--extended", action="store_true")
    args = parser.parse_args()

    eval_stocks = EXTENDED_STOCKS if args.extended else DEFAULT_STOCKS
    tag = "12-stock" if args.extended else "5-stock"

    print(f"\n=== C20: Sector-conditional model [{tag}] ===")
    print("  Loading all stock data...", flush=True)

    # Load all stocks upfront (needed for sector pools)
    stock_data: dict[str, dict] = {}
    for s in ALL_STOCKS:
        df = fetch_df(s)
        if df is None or df.empty:
            print(f"    {s}: no data")
            continue
        feat   = _build_features(df)
        close  = _get_close(df)
        labels = build_triple_barrier_labels(close)
        dates  = _get_dates(feat) if "date" in feat.columns else _get_dates(df)
        stock_data[s] = {"feat": feat, "labels": labels, "dates": dates}
        print(f"    {s}: {len(feat)} rows, sector={STOCK_SECTOR.get(s, '?')}")

    print()
    per_stock, metrics_list = {}, []

    for target_no in eval_stocks:
        if target_no not in stock_data:
            print(f"  {target_no}: missing data, skip")
            continue

        sector = STOCK_SECTOR.get(target_no, "unknown")
        if sector in NO_POOL_SECTORS:
            peer_stocks = []
        else:
            peer_stocks = [s for s in SECTOR_MAP.get(sector, []) if s != target_no and s in stock_data]

        peers = [
            (stock_data[p]["feat"], stock_data[p]["labels"], stock_data[p]["dates"])
            for p in peer_stocks
        ]

        td = stock_data[target_no]
        print(f"  {target_no} [{sector}, peers={peer_stocks}]...", end=" ", flush=True)

        m = walk_forward_sector(
            target_no,
            td["feat"], td["labels"], td["dates"],
            peers, CURRENT_FEATURES,
        )

        per_stock[target_no] = {**m, "sector": sector, "peers": peer_stocks}
        metrics_list.append(m)
        if m:
            print(f"dir={m['dir_accuracy']}%  ↑prec={m['up_precision']}%")
        else:
            print("no output")

    def _avg(key):
        vals = [m[key] for m in metrics_list
                if m and not np.isnan(m.get(key, float("nan")))]
        return round(float(np.mean(vals)), 1) if vals else float("nan")

    avg_dir = _avg("dir_accuracy")
    avg_up  = _avg("up_precision")
    passed  = avg_dir >= PASS_DIR_ACC and avg_up >= PASS_UP_PREC

    print(f"\n  Avg: dir={avg_dir}%  ↑prec={avg_up}%")
    print(f"  Gate (dir≥{PASS_DIR_ACC}% AND ↑prec≥{PASS_UP_PREC}%): {'PASS ✓' if passed else 'FAIL ✗'}")

    result = {
        "experiment":  "C20",
        "description": "Sector-conditional model: pool training within sector",
        "stocks":      eval_stocks,
        "aggregate":   {"dir_accuracy": avg_dir, "up_precision": avg_up},
        "per_stock":   per_stock,
        "passed":      passed,
        "pass_gate":   {"dir_accuracy": PASS_DIR_ACC, "up_precision": PASS_UP_PREC},
    }

    suffix = "_12stock" if args.extended else ""
    out    = ROOT / f"docs/c20_result{suffix}.json"
    out.parent.mkdir(exist_ok=True)
    out.write_text(json.dumps(result, indent=2))
    print(f"\n  Saved: {out}")


if __name__ == "__main__":
    main()