#!/usr/bin/env python3
from __future__ import annotations

import json
import warnings
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List

import numpy as np
import pandas as pd
from sklearn.base import clone
from sklearn.calibration import CalibratedClassifierCV
from sklearn.ensemble import ExtraTreesClassifier, GradientBoostingClassifier, HistGradientBoostingClassifier, RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import FunctionTransformer, StandardScaler

from targeted_optimize_leagues import (
    REPO_ROOT,
    ModelWrapper,
    build_feature_frames,
    build_model_wrappers,
    load_df,
    optimize_threshold,
    select_columns,
)
from train_models import LEAGUES

OUTPUT_DIR = REPO_ROOT / 'experiments' / 'd1_e2_local_refine_pass1'

FOLD_SPECS = [
    ('f1', 0.55, 0.10),
    ('f2', 0.65, 0.10),
    ('f3', 0.75, 0.10),
    ('f4', 0.80, 0.20),
]


@dataclass
class SplitPrediction:
    split_name: str
    threshold_opt: float
    bet_acc_opt: float
    bet_count_opt: int
    test_acc: float
    probs: np.ndarray
    preds: np.ndarray
    y_true: np.ndarray


@dataclass
class CandidateEval:
    code: str
    league: str
    feature_mode: str
    feature_set: str
    train_window: str
    model_name: str
    threshold: float
    holdout_bet_acc: float
    holdout_bet_count: int
    holdout_test_acc: float
    mean_bet_acc: float
    min_bet_acc: float
    total_bet_count: int
    stable: bool
    fallback_fold_count: int
    holdout_opt_threshold: float
    holdout_opt_bet_acc: float
    holdout_opt_bet_count: int


LEAGUE_SEARCH = {
    'd1': {
        'feature_configs': [
            ('enhanced_v1', 'enh_market_stats_28'),
            ('enhanced_v1', 'enh_market_form_31'),
            ('enhanced_v1', 'mi_top15'),
            ('enhanced_v1', 'enh_baseline22'),
            ('baseline_v4', 'base22'),
        ],
        'windows': ['all', 'last_300', 'last_500', 'last_800', 'last_1000'],
        'models': [
            'LR_l2_c02_cal',
            'LR_l2_c03_cal',
            'LR_l2_c05_cal',
            'RF_d6_n600_l3_cal',
            'ET_d7_n700_l2_cal',
            'HGB_d4_i400',
            'HGB_d5_i500',
        ],
        'thresholds': np.arange(0.60, 0.851, 0.01),
    },
    'e2': {
        'feature_configs': [
            ('enhanced_v1', 'mi_top10'),
            ('enhanced_v1', 'mi_top15'),
            ('enhanced_v1', 'mi_top20'),
            ('enhanced_v1', 'enh_market_form_31'),
            ('enhanced_v1', 'enh_market_stats_28'),
            ('enhanced_v1', 'enh_compact_18'),
            ('baseline_v4', 'base22'),
        ],
        'windows': ['all', 'last_300', 'last_400', 'last_500', 'last_800', 'last_1000'],
        'models': [
            'ET_d5_n500_l2_cal',
            'ET_d7_n700_l2_cal',
            'RF_d4_n400_l2',
            'RF_d6_n600_l3_cal',
            'RF_d8_n800_l2_cal',
            'HGB_d4_i400',
            'GB_d3_n300',
        ],
        'thresholds': np.arange(0.60, 0.851, 0.01),
    },
}


def build_wrapper_registry() -> Dict[str, ModelWrapper]:
    wrappers = {w.name: w for w in build_model_wrappers()}
    extras = [
        ModelWrapper('LR_l2_c02', LogisticRegression(C=0.2, solver='liblinear', class_weight='balanced', max_iter=3000, random_state=42), True),
        ModelWrapper('LR_l2_c05', LogisticRegression(C=0.5, solver='liblinear', class_weight='balanced', max_iter=3000, random_state=42), True),
        ModelWrapper('HGB_d5_i500', HistGradientBoostingClassifier(max_depth=5, max_iter=500, learning_rate=0.035, min_samples_leaf=15, random_state=42), True),
    ]
    for wrapper in extras:
        wrappers[wrapper.name] = wrapper
    return wrappers


WRAPPERS = build_wrapper_registry()


def parse_window(window_name: str) -> int | None:
    if window_name == 'all':
        return None
    return int(window_name.replace('last_', ''))


def instantiate_pipeline(model_name: str) -> Pipeline:
    calibrated = model_name.endswith('_cal')
    base_name = model_name[:-4] if calibrated else model_name
    wrapper = WRAPPERS[base_name]
    if calibrated:
        estimator = CalibratedClassifierCV(clone(wrapper.estimator), method='sigmoid', cv=3)
    else:
        estimator = clone(wrapper.estimator)
    scaler = StandardScaler() if wrapper.scale else FunctionTransformer(validate=False)
    return Pipeline([
        ('scaler', scaler),
        ('model', estimator),
    ])


def fixed_or_fallback_metrics(split: SplitPrediction, threshold: float) -> tuple[float, float, int, bool]:
    probs = split.probs
    preds = split.preds
    y_true = split.y_true
    mask = (probs >= threshold) | (probs <= (1 - threshold))
    count = int(mask.sum())
    if count >= 10:
        acc = float(accuracy_score(y_true[mask], preds[mask]))
        return float(threshold), acc, count, False
    return split.threshold_opt, split.bet_acc_opt, split.bet_count_opt, True


def precompute_architecture(df: pd.DataFrame, frames: Dict[str, pd.DataFrame], y_all: pd.Series, feature_mode: str, feature_set: str, train_window: str, model_name: str) -> dict | None:
    features = frames[feature_mode]
    n = len(df)
    split_idx = int(n * 0.8)

    specs = [('holdout', split_idx, n - split_idx)]
    for fold_name, train_ratio, test_ratio in FOLD_SPECS:
        split_end = int(n * train_ratio)
        test_size = int(n * test_ratio)
        specs.append((fold_name, split_end, min(test_size, n - split_end)))

    out: Dict[str, SplitPrediction] = {}
    window = parse_window(train_window)

    for split_name, split_end, test_size in specs:
        test_end = min(n, split_end + test_size)
        if test_end - split_end < 30 or split_end < 200:
            return None

        X_train_all = features.iloc[:split_end].reset_index(drop=True)
        y_train_all = y_all.iloc[:split_end].reset_index(drop=True)
        X_test_all = features.iloc[split_end:test_end].reset_index(drop=True)
        y_test = y_all.iloc[split_end:test_end].reset_index(drop=True)

        if window is not None and len(X_train_all) > window:
            X_train_all = X_train_all.iloc[-window:].reset_index(drop=True)
            y_train_all = y_train_all.iloc[-window:].reset_index(drop=True)

        if len(X_train_all) < 200:
            return None

        Xtr, Xte, _ = select_columns(X_train_all, y_train_all, X_test_all, feature_mode, feature_set)
        pipe = instantiate_pipeline(model_name)
        pipe.fit(Xtr, y_train_all)
        probs = pipe.predict_proba(Xte)[:, 1]
        preds = (probs >= 0.5).astype(int)
        test_acc = float(accuracy_score(y_test, preds))
        opt_threshold, bet_acc_opt, bet_count_opt, _, _ = optimize_threshold(y_test, probs, preds)

        out[split_name] = SplitPrediction(
            split_name=split_name,
            threshold_opt=float(opt_threshold),
            bet_acc_opt=float(bet_acc_opt),
            bet_count_opt=int(bet_count_opt),
            test_acc=test_acc,
            probs=np.asarray(probs),
            preds=np.asarray(preds),
            y_true=np.asarray(y_test),
        )

    return out


def evaluate_threshold(code: str, league: str, feature_mode: str, feature_set: str, train_window: str, model_name: str, threshold: float, split_preds: dict) -> CandidateEval:
    holdout = split_preds['holdout']
    _, holdout_bet_acc, holdout_bet_count, _ = fixed_or_fallback_metrics(holdout, threshold)

    fold_accs: List[float] = []
    total_bets = 0
    fallback_fold_count = 0
    for fold_name, _, _ in FOLD_SPECS:
        split = split_preds[fold_name]
        _, bet_acc, bet_count, fallback = fixed_or_fallback_metrics(split, threshold)
        fold_accs.append(float(bet_acc))
        total_bets += int(bet_count)
        fallback_fold_count += int(fallback)

    mean_bet_acc = float(np.mean(fold_accs))
    min_bet_acc = float(np.min(fold_accs))
    stable = (mean_bet_acc >= 0.72) and (min_bet_acc >= 0.60) and (total_bets >= 80)

    return CandidateEval(
        code=code,
        league=league,
        feature_mode=feature_mode,
        feature_set=feature_set,
        train_window=train_window,
        model_name=model_name,
        threshold=float(round(threshold, 2)),
        holdout_bet_acc=float(holdout_bet_acc),
        holdout_bet_count=int(holdout_bet_count),
        holdout_test_acc=float(holdout.test_acc),
        mean_bet_acc=mean_bet_acc,
        min_bet_acc=min_bet_acc,
        total_bet_count=int(total_bets),
        stable=stable,
        fallback_fold_count=int(fallback_fold_count),
        holdout_opt_threshold=float(holdout.threshold_opt),
        holdout_opt_bet_acc=float(holdout.bet_acc_opt),
        holdout_opt_bet_count=int(holdout.bet_count_opt),
    )


def search_league(code: str) -> pd.DataFrame:
    league = LEAGUES.get(code, code)
    cfg = LEAGUE_SEARCH[code]
    rows: List[dict] = []

    df = load_df(code)
    frames = build_feature_frames(df)
    y_all = df['asian_result'].reset_index(drop=True)

    total_arch = len(cfg['feature_configs']) * len(cfg['windows']) * len(cfg['models'])
    arch_idx = 0
    for feature_mode, feature_set in cfg['feature_configs']:
        for train_window in cfg['windows']:
            for model_name in cfg['models']:
                arch_idx += 1
                if arch_idx == 1 or arch_idx % 20 == 0 or arch_idx == total_arch:
                    print(json.dumps({
                        'code': code,
                        'progress': f'{arch_idx}/{total_arch}',
                        'feature_mode': feature_mode,
                        'feature_set': feature_set,
                        'train_window': train_window,
                        'model_name': model_name,
                    }, ensure_ascii=False), flush=True)
                split_preds = precompute_architecture(df, frames, y_all, feature_mode, feature_set, train_window, model_name)
                if split_preds is None:
                    continue
                for threshold in cfg['thresholds']:
                    result = evaluate_threshold(
                        code=code,
                        league=league,
                        feature_mode=feature_mode,
                        feature_set=feature_set,
                        train_window=train_window,
                        model_name=model_name,
                        threshold=float(threshold),
                        split_preds=split_preds,
                    )
                    rows.append(result.__dict__)

    out = pd.DataFrame(rows).sort_values(
        ['stable', 'mean_bet_acc', 'min_bet_acc', 'total_bet_count', 'holdout_bet_acc', 'holdout_bet_count'],
        ascending=[False, False, False, False, False, False],
    ).reset_index(drop=True)
    out.insert(0, 'candidate_name', [f'{code}_local_{i+1}' for i in range(len(out))])
    return out


def main() -> None:
    import argparse

    warnings.filterwarnings('ignore')
    parser = argparse.ArgumentParser(description='Focused local refinement search for d1/e2 without touching serving artifacts.')
    parser.add_argument('--codes', nargs='+', default=['d1', 'e2'])
    args = parser.parse_args()

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    summary = {}
    for code in args.codes:
        df = search_league(code)
        df.to_csv(OUTPUT_DIR / f'{code}_local_refine_results.csv', index=False, encoding='utf-8-sig')
        top = df.head(20)
        (OUTPUT_DIR / f'{code}_top20.json').write_text(top.to_json(orient='records', force_ascii=False, indent=2), encoding='utf-8')
        stable = df[df['stable']].head(10)
        (OUTPUT_DIR / f'{code}_stable_top10.json').write_text(stable.to_json(orient='records', force_ascii=False, indent=2), encoding='utf-8')
        summary[code] = {
            'best_overall': top.iloc[0].to_dict() if not top.empty else None,
            'best_stable': stable.iloc[0].to_dict() if not stable.empty else None,
            'stable_count': int(df['stable'].sum()),
            'total_candidates': int(len(df)),
        }
        print(json.dumps({'code': code, **summary[code]}, ensure_ascii=False), flush=True)

    summary_path = OUTPUT_DIR / 'summary.json'
    if summary_path.exists():
        try:
            old = json.loads(summary_path.read_text(encoding='utf-8'))
        except Exception:
            old = {}
        old.update(summary)
        summary = old
    summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding='utf-8')


if __name__ == '__main__':
    main()
