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

import json
from dataclasses import dataclass
from pathlib import Path
from typing import 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,
    SearchResult,
    build_feature_frames,
    build_model_wrappers,
    load_df,
    select_columns,
    optimize_threshold,
)

OUTPUT_DIR = REPO_ROOT / 'experiments' / 'targeted_opt_v2'


@dataclass
class FoldResult:
    code: str
    candidate_name: str
    fold: str
    split_end: int
    test_size: int
    threshold: float
    bet_acc: float
    bet_count: int
    test_acc: float


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):
    wrappers = {w.name: w for w in build_model_wrappers()}
    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 evaluate_candidate(code: str, row: pd.Series) -> dict:
    df = load_df(code)
    frames = build_feature_frames(df)
    features = frames[row['feature_mode']]
    y_all = df['asian_result'].reset_index(drop=True)

    n = len(df)
    fold_specs = [
        ('f1', 0.55, 0.10),
        ('f2', 0.65, 0.10),
        ('f3', 0.75, 0.10),
        ('f4', 0.80, 0.20),
    ]

    fold_rows: List[FoldResult] = []
    for fold_name, train_ratio, test_ratio in fold_specs:
        split_end = int(n * train_ratio)
        test_size = int(n * test_ratio)
        test_end = min(n, split_end + test_size)
        if test_end - split_end < 30 or split_end < 200:
            continue

        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)

        window = parse_window(row['train_window'])
        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)

        Xtr, Xte, _ = select_columns(X_train_all, y_train_all, X_test_all, row['feature_mode'], row['feature_set'])
        pipe = instantiate_pipeline(row['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))

        threshold = float(row['threshold'])
        mask = (probs >= threshold) | (probs <= (1 - threshold))
        bet_count = int(mask.sum())
        if bet_count >= 10:
            bet_acc = float(accuracy_score(y_test[mask], preds[mask]))
        else:
            threshold, bet_acc, bet_count, _, _ = optimize_threshold(y_test, probs, preds)

        fold_rows.append(FoldResult(
            code=code,
            candidate_name=row['candidate_name'],
            fold=fold_name,
            split_end=split_end,
            test_size=len(y_test),
            threshold=threshold,
            bet_acc=bet_acc,
            bet_count=bet_count,
            test_acc=test_acc,
        ))

    fold_df = pd.DataFrame([r.__dict__ for r in fold_rows])
    stable = False
    mean_bet = None
    min_bet = None
    total_bets = None
    if not fold_df.empty:
        mean_bet = float(fold_df['bet_acc'].mean())
        min_bet = float(fold_df['bet_acc'].min())
        total_bets = int(fold_df['bet_count'].sum())
        stable = (mean_bet >= 0.72) and (min_bet >= 0.60) and (total_bets >= 80)

    return {
        'candidate': row.to_dict(),
        'folds': fold_rows,
        'stable': stable,
        'mean_bet_acc': mean_bet,
        'min_bet_acc': min_bet,
        'total_bet_count': total_bets,
    }


def main() -> None:
    shortlisted = {
        'd1': [
            {
                'candidate_name': 'd1_hgb_market31_all',
                'feature_mode': 'enhanced_v1', 'feature_set': 'enh_market_form_31', 'train_window': 'all',
                'model_name': 'HGB_d4_i400', 'threshold': 0.78,
            },
            {
                'candidate_name': 'd1_rf_market31_1000',
                'feature_mode': 'enhanced_v1', 'feature_set': 'enh_market_form_31', 'train_window': 'last_1000',
                'model_name': 'RF_d6_n600_l3_cal', 'threshold': 0.55,
            },
            {
                'candidate_name': 'd1_rf_base22_all',
                'feature_mode': 'baseline_v4', 'feature_set': 'base22', 'train_window': 'all',
                'model_name': 'RF_d6_n600_l3_cal', 'threshold': 0.53,
            },
        ],
        'e2n': [
            {
                'candidate_name': 'e2n_et_base22_700',
                'feature_mode': 'baseline_v4', 'feature_set': 'base22', 'train_window': 'last_700',
                'model_name': 'ET_d7_n700_l2_cal', 'threshold': 0.58,
            },
            {
                'candidate_name': 'e2n_lr_base22_400',
                'feature_mode': 'baseline_v4', 'feature_set': 'base22', 'train_window': 'last_400',
                'model_name': 'LR_l2_c03_cal', 'threshold': 0.58,
            },
            {
                'candidate_name': 'e2n_gb_market31_400',
                'feature_mode': 'enhanced_v1', 'feature_set': 'enh_market_form_31', 'train_window': 'last_400',
                'model_name': 'GB_d3_n300', 'threshold': 0.81,
            },
        ],
    }

    summary = {}
    for code, candidates in shortlisted.items():
        rows = []
        for candidate in candidates:
            result = evaluate_candidate(code, pd.Series(candidate))
            fold_df = pd.DataFrame([r.__dict__ for r in result['folds']])
            if not fold_df.empty:
                fold_df.to_csv(OUTPUT_DIR / f"{candidate['candidate_name']}_folds.csv", index=False, encoding='utf-8-sig')
            rows.append({
                'candidate_name': candidate['candidate_name'],
                'stable': result['stable'],
                'mean_bet_acc': result['mean_bet_acc'],
                'min_bet_acc': result['min_bet_acc'],
                'total_bet_count': result['total_bet_count'],
                **candidate,
            })
        df = pd.DataFrame(rows).sort_values(['stable', 'mean_bet_acc', 'total_bet_count'], ascending=[False, False, False])
        df.to_csv(OUTPUT_DIR / f'{code}_stability_summary.csv', index=False, encoding='utf-8-sig')
        summary[code] = df.to_dict(orient='records')

    (OUTPUT_DIR / 'stability_summary.json').write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding='utf-8')
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == '__main__':
    main()
