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

import argparse
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.metrics import accuracy_score
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import FunctionTransformer, StandardScaler

from targeted_optimize_leagues import (
    REPO_ROOT,
    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_shortlist'

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

SEARCH_PLAN = {
    'd1': {
        'architectures': [
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_market_stats_28', 'train_window': 'all', 'model_name': 'LR_l2_c03_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_market_stats_28', 'train_window': 'last_500', 'model_name': 'LR_l2_c03_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_market_stats_28', 'train_window': 'last_800', 'model_name': 'LR_l2_c03_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'mi_top15', 'train_window': 'last_800', 'model_name': 'LR_l2_c03_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_market_form_31', 'train_window': 'last_300', 'model_name': 'ET_d7_n700_l2_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_market_form_31', 'train_window': 'last_500', 'model_name': 'ET_d7_n700_l2_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_market_form_31', 'train_window': 'last_800', 'model_name': 'RF_d6_n600_l3_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_market_form_31', 'train_window': 'last_1000', 'model_name': 'RF_d6_n600_l3_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_market_form_31', 'train_window': 'last_800', 'model_name': 'RF_d8_n800_l2_cal'},
            {'feature_mode': 'baseline_v4', 'feature_set': 'base22', 'train_window': 'all', 'model_name': 'RF_d6_n600_l3_cal'},
            {'feature_mode': 'baseline_v4', 'feature_set': 'base22', 'train_window': 'last_500', 'model_name': 'RF_d6_n600_l3_cal'},
            {'feature_mode': 'baseline_v4', 'feature_set': 'base22', 'train_window': 'last_800', 'model_name': 'RF_d8_n800_l2_cal'},
            {'feature_mode': 'baseline_v4', 'feature_set': 'base22', 'train_window': 'last_800', 'model_name': 'ET_d5_n500_l2'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_baseline22', 'train_window': 'all', 'model_name': 'HGB_d4_i400'},
        ],
        'thresholds': np.arange(0.60, 0.851, 0.01),
    },
    'e2': {
        'architectures': [
            {'feature_mode': 'enhanced_v1', 'feature_set': 'mi_top15', 'train_window': 'all', 'model_name': 'ET_d7_n700_l2_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'mi_top15', 'train_window': 'last_500', 'model_name': 'ET_d7_n700_l2_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'mi_top15', 'train_window': 'last_800', 'model_name': 'ET_d7_n700_l2_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'mi_top20', 'train_window': 'all', 'model_name': 'ET_d7_n700_l2_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'mi_top20', 'train_window': 'last_800', 'model_name': 'ET_d7_n700_l2_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'mi_top15', 'train_window': 'all', 'model_name': 'RF_d6_n600_l3_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'mi_top15', 'train_window': 'all', 'model_name': 'RF_d8_n800_l2_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'mi_top20', 'train_window': 'all', 'model_name': 'RF_d8_n800_l2_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_market_form_31', 'train_window': 'last_800', 'model_name': 'RF_d4_n400_l2'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_market_form_31', 'train_window': 'last_1000', 'model_name': 'RF_d4_n400_l2'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_market_form_31', 'train_window': 'last_800', 'model_name': 'ET_d5_n500_l2_cal'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_compact_18', 'train_window': 'last_400', 'model_name': 'RF_d6_n600_l3'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_compact_18', 'train_window': 'last_500', 'model_name': 'RF_d6_n600_l3'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_compact_18', 'train_window': 'last_400', 'model_name': 'RF_d8_n800_l2_cal'},
            {'feature_mode': 'baseline_v4', 'feature_set': 'base22', 'train_window': 'all', 'model_name': 'HGB_d4_i400'},
            {'feature_mode': 'baseline_v4', 'feature_set': 'base22', 'train_window': 'last_800', 'model_name': 'HGB_d4_i400'},
            {'feature_mode': 'enhanced_v1', 'feature_set': 'enh_market_stats_28', 'train_window': 'all', 'model_name': 'RF_d8_n800_l2_cal'},
        ],
        'thresholds': np.arange(0.60, 0.821, 0.01),
    },
}

WRAPPERS = {w.name: w for w in build_model_wrappers()}


@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


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]
    estimator = CalibratedClassifierCV(clone(wrapper.estimator), method='sigmoid', cv=3) if calibrated else clone(wrapper.estimator)
    scaler = StandardScaler() if wrapper.scale else FunctionTransformer(validate=False)
    return Pipeline([('scaler', scaler), ('model', estimator)])


def fit_predict_splits(df: pd.DataFrame, frames: Dict[str, pd.DataFrame], arch: dict) -> dict | None:
    feature_mode = arch['feature_mode']
    feature_set = arch['feature_set']
    train_window = arch['train_window']
    model_name = arch['model_name']

    features = frames[feature_mode]
    y_all = df['asian_result'].reset_index(drop=True)
    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 = {}
    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 metrics_for_threshold(split: SplitPrediction, threshold: float) -> tuple[float, float, int, bool]:
    mask = (split.probs >= threshold) | (split.probs <= (1 - threshold))
    count = int(mask.sum())
    if count >= 10:
        acc = float(accuracy_score(split.y_true[mask], split.preds[mask]))
        return float(threshold), acc, count, False
    return split.threshold_opt, split.bet_acc_opt, split.bet_count_opt, True


def evaluate_code(code: str) -> pd.DataFrame:
    league = LEAGUES.get(code, code)
    df = load_df(code)
    frames = build_feature_frames(df)
    plan = SEARCH_PLAN[code]
    rows: List[dict] = []

    for idx, arch in enumerate(plan['architectures'], start=1):
        print(json.dumps({'code': code, 'architecture': idx, 'of': len(plan['architectures']), **arch}, ensure_ascii=False), flush=True)
        split_preds = fit_predict_splits(df, frames, arch)
        if split_preds is None:
            continue
        holdout = split_preds['holdout']
        for threshold in plan['thresholds']:
            _, holdout_bet_acc, holdout_bet_count, holdout_fallback = metrics_for_threshold(holdout, float(threshold))
            fold_accs = []
            total_bets = 0
            fallback_fold_count = 0
            for fold_name, _, _ in FOLD_SPECS:
                _, bet_acc, bet_count, fallback = metrics_for_threshold(split_preds[fold_name], float(threshold))
                fold_accs.append(float(bet_acc))
                total_bets += int(bet_count)
                fallback_fold_count += int(fallback)
            rows.append({
                'code': code,
                'league': league,
                **arch,
                '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),
                '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),
                'holdout_fallback': bool(holdout_fallback),
                'mean_bet_acc': float(np.mean(fold_accs)),
                'min_bet_acc': float(np.min(fold_accs)),
                'total_bet_count': int(total_bets),
                'fallback_fold_count': int(fallback_fold_count),
                'stable': bool((np.mean(fold_accs) >= 0.72) and (np.min(fold_accs) >= 0.60) and (total_bets >= 80)),
            })

    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}_short_{i+1}' for i in range(len(out))])
    return out


def main() -> None:
    warnings.filterwarnings('ignore')
    parser = argparse.ArgumentParser(description='Focused shortlist refinement for d1/e2.')
    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 = evaluate_code(code)
        df.to_csv(OUTPUT_DIR / f'{code}_results.csv', index=False, encoding='utf-8-sig')
        top20 = df.head(20)
        stable10 = df[df['stable']].head(10)
        (OUTPUT_DIR / f'{code}_top20.json').write_text(top20.to_json(orient='records', force_ascii=False, indent=2), encoding='utf-8')
        (OUTPUT_DIR / f'{code}_stable_top10.json').write_text(stable10.to_json(orient='records', force_ascii=False, indent=2), encoding='utf-8')
        summary[code] = {
            'best_overall': top20.iloc[0].to_dict() if not top20.empty else None,
            'best_stable': stable10.iloc[0].to_dict() if not stable10.empty else None,
            'stable_count': int(df['stable'].sum()),
            'candidate_count': 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()
