#!/usr/bin/env python3
"""C21: dry-run Old Wang style rules.

Rules under test:
- 三陽開泰 / 三聲無奈
- 投信看 10MA, 外資看 20MA
- 爆大量 K 棒高低點防守
- 跳空缺口守住 / 回補

No production FEATURE_COLUMNS are modified by this script.
"""

from __future__ import annotations

import argparse
import json
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 numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier

from models.predictor import FEATURE_COLUMNS, _build_features
from scripts.backtest_c20_targeted_signal_gates import _safe_div, _true_range, _zscore_prior, fetch_validation_df
from scripts.improvement_harness import (
    DEFAULT_STOCKS,
    EXTENDED_STOCKS,
    LABEL_HORIZON,
    MIN_TRAIN,
    RF_PARAMS,
    STEP,
    build_triple_barrier_labels,
    compute_metrics,
)


PASS_PER_STOCK_ACCURACY_DELTA = 3.0
MIN_SIGNAL_RATIO = 0.70

OLDWANG_FEATURES = [
    "oldwang_triple_bull",
    "oldwang_triple_bear",
    "oldwang_ma5_hold",
    "oldwang_trust_ma10_guard",
    "oldwang_trust_ma10_broken",
    "oldwang_foreign_ma20_guard",
    "oldwang_foreign_ma20_broken",
    "oldwang_volume_spike",
    "oldwang_volume_high_break",
    "oldwang_volume_low_guard",
    "oldwang_volume_low_break",
    "oldwang_gap_guard",
    "oldwang_gap_filled",
    "oldwang_bull_score",
    "oldwang_bear_score",
]

STRATEGIES = [
    "oldwang_features",
    "oldwang_veto",
    "oldwang_long_gate",
    "oldwang_features_veto",
]


def _rolling_recent_flag(flag: pd.Series, window: int) -> pd.Series:
    return flag.astype(float).rolling(window, min_periods=1).max().fillna(0.0)


def _last_event_level(level: pd.Series, event: pd.Series) -> pd.Series:
    return level.where(event).shift(1).ffill()


def build_oldwang_features(df: pd.DataFrame) -> pd.DataFrame:
    out = pd.DataFrame(index=df.index)
    open_ = df["open"].astype(float)
    high = df["high"].astype(float)
    low = df["low"].astype(float)
    close = df["close"].astype(float)
    volume = df.get("volume", pd.Series(0.0, index=df.index)).astype(float)
    prev_high = high.shift(1)
    prev_close = close.shift(1)

    ma5 = df.get("ma5", close.rolling(5, min_periods=3).mean()).astype(float)
    ma10 = df.get("ma10", close.rolling(10, min_periods=5).mean()).astype(float)
    ma20 = df.get("ma20", close.rolling(20, min_periods=10).mean()).astype(float)
    ma5_slope = ma5.pct_change(3).fillna(0.0)
    ma10_slope = ma10.pct_change(3).fillna(0.0)
    ma20_slope = ma20.pct_change(3).fillna(0.0)

    triple_bull = (
        (close > ma5)
        & (ma5 > ma10)
        & (ma10 > ma20)
        & (ma5_slope > 0)
        & (ma10_slope > 0)
        & (ma20_slope > 0)
    )
    triple_bear = (
        (close < ma5)
        & (ma5 < ma10)
        & (ma10 < ma20)
        & (ma5_slope < 0)
        & (ma10_slope < 0)
        & (ma20_slope < 0)
    )
    out["oldwang_triple_bull"] = triple_bull.astype(float)
    out["oldwang_triple_bear"] = triple_bear.astype(float)
    out["oldwang_ma5_hold"] = (close >= ma5).astype(float)

    trust_5d = df.get("trust_net", pd.Series(0.0, index=df.index)).astype(float).rolling(5, min_periods=1).sum()
    foreign_5d = df.get("foreign_net", pd.Series(0.0, index=df.index)).astype(float).rolling(5, min_periods=1).sum()
    out["oldwang_trust_ma10_guard"] = ((trust_5d > 0) & (close >= ma10)).astype(float)
    out["oldwang_trust_ma10_broken"] = ((trust_5d > 0) & (close < ma10)).astype(float)
    out["oldwang_foreign_ma20_guard"] = ((foreign_5d > 0) & (close >= ma20)).astype(float)
    out["oldwang_foreign_ma20_broken"] = ((foreign_5d > 0) & (close < ma20)).astype(float)

    log_volume = np.log1p(volume.clip(lower=0.0))
    volume_z = _zscore_prior(log_volume, 20, 10)
    volume_spike = volume_z >= 2.0
    spike_high = _last_event_level(high, volume_spike)
    spike_low = _last_event_level(low, volume_spike)
    out["oldwang_volume_spike"] = volume_spike.astype(float)
    out["oldwang_volume_high_break"] = ((close > spike_high) & spike_high.notna()).astype(float)
    out["oldwang_volume_low_guard"] = ((close >= spike_low) & spike_low.notna()).astype(float)
    out["oldwang_volume_low_break"] = ((close < spike_low) & spike_low.notna()).astype(float)

    tr = _true_range(df)
    atr = df.get("atr", tr.rolling(14, min_periods=5).mean()).astype(float)
    atr_prior = atr.shift(1).replace(0, np.nan).bfill().fillna(close * 0.02)
    gap_up = open_ > prev_high
    gap_support = _last_event_level(prev_high, gap_up)
    recent_gap_up = _rolling_recent_flag(gap_up, 5) > 0
    out["oldwang_gap_guard"] = (recent_gap_up & (low >= gap_support - 0.1 * atr_prior)).astype(float)
    out["oldwang_gap_filled"] = (recent_gap_up & (close < gap_support)).astype(float)

    out["oldwang_bull_score"] = (
        out["oldwang_triple_bull"]
        + out["oldwang_ma5_hold"]
        + out["oldwang_trust_ma10_guard"]
        + out["oldwang_foreign_ma20_guard"]
        + out["oldwang_volume_high_break"]
        + out["oldwang_gap_guard"]
    )
    out["oldwang_bear_score"] = (
        out["oldwang_triple_bear"]
        + out["oldwang_trust_ma10_broken"]
        + out["oldwang_foreign_ma20_broken"]
        + out["oldwang_volume_low_break"]
        + out["oldwang_gap_filled"]
    )

    return out[OLDWANG_FEATURES].replace([np.inf, -np.inf], np.nan).fillna(0.0)


def _apply_veto(pred: np.ndarray, features: pd.DataFrame) -> np.ndarray:
    out = pred.copy()
    veto = features["oldwang_bear_score"].to_numpy() >= 1
    out[(out == 1) & veto] = 0
    return out


def _apply_long_gate(pred: np.ndarray, features: pd.DataFrame) -> np.ndarray:
    out = pred.copy()
    allow_long = features["oldwang_bull_score"].to_numpy() >= 2
    out[(out == 1) & ~allow_long] = 0
    return out


def walk_forward_oldwang(feat_df: pd.DataFrame, oldwang_df: pd.DataFrame, labels: np.ndarray) -> tuple[dict, dict[str, dict]]:
    base_avail = [
        col for col in FEATURE_COLUMNS
        if col in feat_df.columns and col not in OLDWANG_FEATURES
    ]
    y_true_all: list[float] = []
    y_pred_base_all: list[int] = []
    y_pred_by_strategy: dict[str, list[int]] = {strategy: [] for strategy in STRATEGIES}
    n = len(feat_df)

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

        y_train_full = labels[:train_end]
        valid_train = np.where(~np.isnan(y_train_full))[0]
        if len(valid_train) < MIN_TRAIN or len(np.unique(y_train_full[valid_train])) < 2:
            cutoff += STEP
            continue

        test_end = min(cutoff + STEP, n - LABEL_HORIZON)
        y_test = labels[cutoff:test_end]
        valid_test_mask = ~np.isnan(y_test)
        if valid_test_mask.sum() == 0:
            cutoff += STEP
            continue
        test_idx = np.arange(cutoff, test_end)[valid_test_mask]

        X_train_base = feat_df.loc[valid_train, base_avail].fillna(0.0)
        X_test_base = feat_df.loc[test_idx, base_avail].fillna(0.0)
        X_train_oldwang = pd.concat(
            [X_train_base.reset_index(drop=True), oldwang_df.loc[valid_train, OLDWANG_FEATURES].reset_index(drop=True)],
            axis=1,
        )
        X_test_oldwang = pd.concat(
            [X_test_base.reset_index(drop=True), oldwang_df.loc[test_idx, OLDWANG_FEATURES].reset_index(drop=True)],
            axis=1,
        )

        y_train = y_train_full[valid_train].astype(int)
        base_model = RandomForestClassifier(**RF_PARAMS)
        oldwang_model = RandomForestClassifier(**RF_PARAMS)
        base_model.fit(X_train_base.values, y_train)
        oldwang_model.fit(X_train_oldwang.values, y_train)

        y_pred_base = base_model.predict(X_test_base.values)
        y_pred_oldwang = oldwang_model.predict(X_test_oldwang.values)
        test_oldwang = oldwang_df.loc[test_idx, OLDWANG_FEATURES]

        y_true_all.extend(y_test[valid_test_mask].tolist())
        y_pred_base_all.extend(y_pred_base.tolist())
        y_pred_by_strategy["oldwang_features"].extend(y_pred_oldwang.tolist())
        y_pred_by_strategy["oldwang_veto"].extend(_apply_veto(y_pred_base, test_oldwang).tolist())
        y_pred_by_strategy["oldwang_long_gate"].extend(_apply_long_gate(y_pred_base, test_oldwang).tolist())
        y_pred_by_strategy["oldwang_features_veto"].extend(_apply_veto(y_pred_oldwang, test_oldwang).tolist())
        cutoff += STEP

    if not y_true_all:
        return {}, {}

    y_true = np.array(y_true_all)
    baseline = compute_metrics(y_true, np.array(y_pred_base_all))
    candidates = {strategy: compute_metrics(y_true, np.array(preds)) for strategy, preds in y_pred_by_strategy.items()}
    for row in candidates.values():
        row["added_features"] = len(OLDWANG_FEATURES)
    return baseline, candidates


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


def _aggregate(rows: list[dict]) -> dict:
    return {
        "accuracy": _mean(rows, "accuracy"),
        "dir_accuracy": _mean(rows, "dir_accuracy"),
        "up_precision": _mean(rows, "up_precision"),
        "dn_precision": _mean(rows, "dn_precision"),
        "n_signals": _mean(rows, "n_signals"),
        "n_predictions": _mean(rows, "n_predictions"),
    }


def _json_default(value):
    if isinstance(value, np.integer):
        return int(value)
    if isinstance(value, np.floating):
        return float(value)
    if isinstance(value, np.bool_):
        return bool(value)
    raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable")


def run(stocks: list[str], output_path: Path, *, with_institutional: bool = True) -> dict:
    per_stock: dict[str, dict] = {}
    agg_base: list[dict] = []
    agg_by_strategy: dict[str, list[dict]] = {strategy: [] for strategy in STRATEGIES}

    hdr = f"{'Stock':>6}  {'Model':>21}  {'Acc%':>5}  {'Dir%':>5}  {'Up%':>6}  {'Signals':>7}  {'Delta':>6}"
    print(f"\n{hdr}\n{'-' * len(hdr)}")
    for stock_no in stocks:
        print(f"  computing {stock_no}...", end="\r", flush=True)
        df = fetch_validation_df(stock_no, with_institutional=with_institutional)
        if df is None or df.empty:
            print(f"{stock_no:>6}  no data")
            continue

        feat = _build_features(df)
        oldwang = build_oldwang_features(df)
        close = (df.set_index("date")["close"] if "date" in df.columns else df["close"]).values
        labels = build_triple_barrier_labels(close)
        baseline, candidates = walk_forward_oldwang(feat, oldwang, labels)
        if not baseline:
            print(f"{stock_no:>6}  insufficient folds")
            continue

        agg_base.append(baseline)
        stock_row: dict[str, dict] = {"baseline": baseline}
        print(
            f"{stock_no:>6}  {'baseline':>21}  {baseline['accuracy']:>5.1f}  "
            f"{baseline['dir_accuracy']:>5.1f}  {baseline['up_precision']:>6.1f}  "
            f"{baseline['n_signals']:>7}  {'':>6}"
        )
        for strategy in STRATEGIES:
            row = candidates[strategy]
            delta = round(row["accuracy"] - baseline["accuracy"], 1)
            signal_ratio = (row["n_signals"] / baseline["n_signals"]) if baseline["n_signals"] else 1.0
            passed = delta >= PASS_PER_STOCK_ACCURACY_DELTA and signal_ratio >= MIN_SIGNAL_RATIO
            row = {
                **row,
                "accuracy_delta_pp": delta,
                "signal_ratio_vs_baseline": round(signal_ratio, 3),
                "passed": passed,
            }
            stock_row[strategy] = row
            agg_by_strategy[strategy].append(row)
            print(
                f"{'':>6}  {strategy:>21}  {row['accuracy']:>5.1f}  "
                f"{row['dir_accuracy']:>5.1f}  {row['up_precision']:>6.1f}  "
                f"{row['n_signals']:>7}  {delta:>+6.1f} {'PASS' if passed else 'FAIL'}"
            )
        per_stock[stock_no] = stock_row

    aggregate = {"baseline": _aggregate(agg_base)}
    for strategy, rows in agg_by_strategy.items():
        aggregate[strategy] = _aggregate(rows)
        aggregate[strategy]["accuracy_delta_pp"] = round(aggregate[strategy]["accuracy"] - aggregate["baseline"]["accuracy"], 1)
        base_signals = aggregate["baseline"].get("n_signals", 0) or 0
        strategy_signals = aggregate[strategy].get("n_signals", 0) or 0
        aggregate[strategy]["signal_ratio_vs_baseline"] = round(strategy_signals / base_signals, 3) if base_signals else 1.0

    passing_strategies = [
        strategy
        for strategy in STRATEGIES
        if per_stock
        and all(row[strategy]["passed"] for row in per_stock.values())
        and aggregate[strategy]["signal_ratio_vs_baseline"] >= MIN_SIGNAL_RATIO
    ]
    best_strategy = max(
        STRATEGIES,
        key=lambda strategy: (
            aggregate[strategy].get("accuracy_delta_pp", float("-inf")),
            aggregate[strategy].get("dir_accuracy", float("-inf")),
        ),
        default=None,
    )

    print("\n=== AGGREGATE ===")
    print(
        f"  {'baseline':>21}: acc={aggregate['baseline']['accuracy']}%  "
        f"dir={aggregate['baseline']['dir_accuracy']}%  up={aggregate['baseline']['up_precision']}%  "
        f"signals={aggregate['baseline']['n_signals']}"
    )
    for strategy in STRATEGIES:
        row = aggregate[strategy]
        print(
            f"  {strategy:>21}: acc={row['accuracy']}%  dir={row['dir_accuracy']}%  "
            f"up={row['up_precision']}%  signals={row['n_signals']}  "
            f"delta={row['accuracy_delta_pp']:+.1f}pp  signal_ratio={row['signal_ratio_vs_baseline']}"
        )

    result = {
        "experiment": "C21_oldwang_rules",
        "description": "Dry run of Old Wang style MA, institutional guard, volume-bar, and gap-guard rules.",
        "stocks": stocks,
        "with_institutional": with_institutional,
        "pass_criterion": {
            "every_stock_accuracy_delta_pp_at_least": PASS_PER_STOCK_ACCURACY_DELTA,
            "min_signal_ratio_vs_baseline": MIN_SIGNAL_RATIO,
        },
        "candidate_oldwang_features": OLDWANG_FEATURES,
        "strategies": STRATEGIES,
        "aggregate": aggregate,
        "best_strategy": best_strategy,
        "passing_strategies": passing_strategies,
        "per_stock": per_stock,
        "passed": bool(passing_strategies),
        "promotion_decision": "integrate" if passing_strategies else "do_not_integrate",
        "validation": {
            "mode": "walk_forward_with_embargo",
            "label_horizon": LABEL_HORIZON,
            "train_end": "cutoff - LABEL_HORIZON",
            "step": STEP,
            "min_train": MIN_TRAIN,
            "leakage_guard": "rules use current-close-known and prior rolling/event levels; training excludes the label horizon before each test slice.",
        },
    }

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(result, indent=2, ensure_ascii=False, default=_json_default))
    print(f"\nSaved -> {output_path}")
    return result


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--stocks", default="", help="Comma-separated stock list")
    parser.add_argument("--extended", action="store_true", help="Use 12-stock extended validation")
    parser.add_argument("--no-institutional", action="store_true", help="Skip institutional flow fetch")
    parser.add_argument("--output", default="docs/c21_oldwang_rules_result.json")
    args = parser.parse_args()

    if args.stocks:
        stocks = [code.strip() for code in args.stocks.split(",") if code.strip()]
    else:
        stocks = EXTENDED_STOCKS if args.extended else DEFAULT_STOCKS

    result = run(stocks, ROOT / args.output, with_institutional=not args.no_institutional)
    return 0 if result.get("passed") else 1


if __name__ == "__main__":
    raise SystemExit(main())