#!/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.ensemble import ExtraTreesClassifier, GradientBoostingClassifier, HistGradientBoostingClassifier, RandomForestClassifier
from sklearn.feature_selection import mutual_info_classif
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import FunctionTransformer, StandardScaler

from train_models import LEAGUES, FEATURE_COLUMNS, create_features, safe_to_datetime
from optimize_priority_models import create_enhanced_features

REPO_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_OUTPUT = REPO_ROOT / 'experiments' / 'targeted_opt_v2'


@dataclass
class SearchResult:
    code: str
    league: str
    feature_mode: str
    feature_set: str
    train_window: str
    model_name: str
    threshold: float
    bet_acc: float
    bet_count: int
    high_conf_acc: float
    high_conf_count: int
    test_acc: float
    test_count: int


class ModelWrapper:
    def __init__(self, name: str, estimator, scale: bool):
        self.name = name
        self.estimator = estimator
        self.scale = scale

    def build(self):
        scaler = StandardScaler() if self.scale else FunctionTransformer(validate=False)
        return Pipeline([
            ('scaler', scaler),
            ('model', clone(self.estimator)),
        ])


BASELINE_EQUIV = [
    'avg_home_odds', 'avg_away_odds', 'home_prob', 'away_prob', 'prob_diff',
    'handicap', 'asian_home_water', 'asian_away_water', 'water_diff', 'handicap_x_water',
    'home_ppg', 'away_ppg', 'ppg_diff', 'home_form', 'away_form', 'form_diff',
    'home_adv', 'adv_diff', 'prob_x_form', 'form_x_ppg', 'prob_x_handicap', 'water_x_prob',
]

FEATURE_SETS = {
    'enh_all': None,
    'enh_baseline22': BASELINE_EQUIV,
    'enh_market_form_31': [
        'avg_home_odds', 'avg_draw_odds', 'avg_away_odds', 'home_prob', 'away_prob', 'draw_prob',
        'home_prob_norm', 'away_prob_norm', 'draw_prob_norm', 'market_overround', 'prob_diff', 'prob_ratio', 'draw_bias',
        'handicap', 'handicap_abs', 'asian_home_water', 'asian_away_water', 'water_diff', 'water_sum', 'water_ratio', 'handicap_x_water',
        'home_ppg', 'away_ppg', 'home_form', 'away_form', 'ppg_diff', 'form_diff', 'home_adv', 'adv_diff', 'elo_diff', 'venue_strength'
    ],
    'enh_market_stats_28': [
        'avg_home_odds', 'avg_draw_odds', 'avg_away_odds', 'home_prob', 'away_prob', 'draw_prob',
        'prob_diff', 'prob_ratio', 'handicap', 'handicap_abs', 'asian_home_water', 'asian_away_water', 'water_diff', 'water_ratio',
        'home_scored', 'away_scored', 'home_conceded', 'away_conceded', 'scored_diff', 'conceded_diff',
        'home_sot_for', 'away_sot_for', 'home_sot_against', 'away_sot_against', 'sot_for_diff', 'sot_against_diff',
        'elo_diff', 'prob_x_elo'
    ],
    'enh_compact_18': [
        'avg_home_odds', 'avg_draw_odds', 'avg_away_odds', 'home_prob_norm', 'away_prob_norm', 'draw_prob_norm',
        'handicap', 'asian_home_water', 'asian_away_water', 'water_diff',
        'ppg_diff', 'form_diff', 'venue_ppg_diff', 'venue_form_diff',
        'gd_diff', 'elo_diff', 'prob_x_handicap', 'venue_strength'
    ],
    'base22': FEATURE_COLUMNS,
}


def build_model_wrappers() -> List[ModelWrapper]:
    return [
        ModelWrapper('LR_l2_c01', LogisticRegression(C=0.1, solver='liblinear', class_weight='balanced', max_iter=3000, random_state=42), True),
        ModelWrapper('LR_l2_c03', LogisticRegression(C=0.3, solver='liblinear', class_weight='balanced', max_iter=3000, random_state=42), True),
        ModelWrapper('LR_l2_c10', LogisticRegression(C=1.0, solver='liblinear', class_weight='balanced', max_iter=3000, random_state=42), True),
        ModelWrapper('RF_d4_n400_l2', RandomForestClassifier(n_estimators=400, max_depth=4, min_samples_leaf=2, random_state=42, n_jobs=-1), False),
        ModelWrapper('RF_d6_n600_l3', RandomForestClassifier(n_estimators=600, max_depth=6, min_samples_leaf=3, random_state=42, n_jobs=-1), False),
        ModelWrapper('RF_d8_n800_l2', RandomForestClassifier(n_estimators=800, max_depth=8, min_samples_leaf=2, random_state=42, n_jobs=-1), False),
        ModelWrapper('ET_d5_n500_l2', ExtraTreesClassifier(n_estimators=500, max_depth=5, min_samples_leaf=2, random_state=42, n_jobs=-1), False),
        ModelWrapper('ET_d7_n700_l2', ExtraTreesClassifier(n_estimators=700, max_depth=7, min_samples_leaf=2, random_state=42, n_jobs=-1), False),
        ModelWrapper('ET_full_n900_l2', ExtraTreesClassifier(n_estimators=900, max_depth=None, min_samples_leaf=2, random_state=42, n_jobs=-1), False),
        ModelWrapper('GB_d3_n300', GradientBoostingClassifier(n_estimators=300, learning_rate=0.03, max_depth=3, random_state=42), False),
        ModelWrapper('HGB_d4_i400', HistGradientBoostingClassifier(max_depth=4, max_iter=400, learning_rate=0.04, min_samples_leaf=15, random_state=42), True),
        ModelWrapper('HGB_d6_i500', HistGradientBoostingClassifier(max_depth=6, max_iter=500, learning_rate=0.03, min_samples_leaf=15, random_state=42), True),
    ]


def optimize_threshold(y_true: pd.Series, probs: np.ndarray, preds: np.ndarray) -> tuple[float, float, int, float, int]:
    best = None
    for threshold in np.arange(0.50, 0.91, 0.01):
        mask = (probs >= threshold) | (probs <= (1 - threshold))
        count = int(mask.sum())
        if count < 10:
            continue
        acc = float(accuracy_score(y_true[mask], preds[mask]))
        high_mask = (probs >= min(0.95, threshold + 0.05)) | (probs <= max(0.05, 1 - threshold - 0.05))
        high_count = int(high_mask.sum())
        high_acc = float(accuracy_score(y_true[high_mask], preds[high_mask])) if high_count > 0 else 0.0
        score = (acc, count)
        if best is None or score > best[0]:
            best = (score, float(round(threshold, 2)), acc, count, high_acc, high_count)
    if best is None:
        full_acc = float(accuracy_score(y_true, preds))
        return 0.50, full_acc, int(len(y_true)), full_acc, int(len(y_true))
    _, threshold, acc, count, high_acc, high_count = best
    return threshold, acc, count, high_acc, high_count


def select_columns(X_train: pd.DataFrame, y_train: pd.Series, X_test: pd.DataFrame, mode: str, name: str) -> tuple[pd.DataFrame, pd.DataFrame, List[str]]:
    if name in FEATURE_SETS:
        cols = FEATURE_SETS[name]
        if cols is None:
            cols = list(X_train.columns)
        return X_train[cols].copy(), X_test[cols].copy(), cols

    if name.startswith('mi_top'):
        k = int(name.replace('mi_top', ''))
        mi = mutual_info_classif(X_train.fillna(0), y_train, discrete_features=False, random_state=42)
        cols = pd.Series(mi, index=X_train.columns).sort_values(ascending=False).head(k).index.tolist()
        return X_train[cols].copy(), X_test[cols].copy(), cols

    raise ValueError(f'unknown feature set: {name}')


def build_feature_frames(df: pd.DataFrame) -> Dict[str, pd.DataFrame]:
    return {
        'baseline_v4': create_features(df),
        'enhanced_v1': create_enhanced_features(df),
    }


def get_window_slices(split_idx: int, train_windows: List[int | None]) -> List[tuple[str, slice]]:
    out = []
    for win in train_windows:
        if win is None:
            out.append(('all', slice(0, split_idx)))
        else:
            start = max(0, split_idx - win)
            out.append((f'last_{win}', slice(start, split_idx)))
    return out


def load_df(code: str) -> pd.DataFrame:
    path = REPO_ROOT / 'historical_data' / f'{code}_history.csv'
    df = pd.read_csv(path)
    df['date_obj'] = safe_to_datetime(df['Date'])
    df = df.dropna(subset=['date_obj']).sort_values('date_obj').reset_index(drop=True)
    df = df.dropna(subset=['FTHG', 'FTAG'])
    if 'AHh' in df.columns:
        df['gd_h'] = df['FTHG'] - df['FTAG'] + df['AHh'].fillna(0)
    else:
        df['gd_h'] = df['FTHG'] - df['FTAG']
    df['asian_result'] = (df['gd_h'] > 0).astype(int)
    df.loc[df['gd_h'] == 0, 'asian_result'] = 0
    return df


def evaluate_league(code: str, output_dir: Path) -> List[SearchResult]:
    league = LEAGUES.get(code, code)
    df = load_df(code)
    frames = build_feature_frames(df)
    split_idx = int(len(df) * 0.8)
    y_train_all = df['asian_result'].iloc[:split_idx].reset_index(drop=True)
    y_test = df['asian_result'].iloc[split_idx:].reset_index(drop=True)

    feature_set_names = {
        'baseline_v4': ['base22'],
        'enhanced_v1': ['enh_all', 'enh_baseline22', 'enh_market_form_31', 'enh_market_stats_28', 'enh_compact_18', 'mi_top15', 'mi_top30'],
    }
    windows = [None, 400, 700, 1000]
    wrappers = build_model_wrappers()
    results: List[SearchResult] = []

    for feature_mode, features in frames.items():
        X_train_all = features.iloc[:split_idx].reset_index(drop=True)
        X_test_all = features.iloc[split_idx:].reset_index(drop=True)
        for window_name, window_slice in get_window_slices(split_idx, windows):
            X_train_window = X_train_all.iloc[window_slice].reset_index(drop=True)
            y_train_window = y_train_all.iloc[window_slice].reset_index(drop=True)
            if len(X_train_window) < 200:
                continue
            for feature_set_name in feature_set_names[feature_mode]:
                Xtr, Xte, cols = select_columns(X_train_window, y_train_window, X_test_all, feature_mode, feature_set_name)
                for wrapper in wrappers:
                    pipe = wrapper.build()
                    pipe.fit(Xtr, y_train_window)
                    probs = pipe.predict_proba(Xte)[:, 1]
                    preds = (probs >= 0.5).astype(int)
                    test_acc = float(accuracy_score(y_test, preds))
                    threshold, bet_acc, bet_count, high_conf_acc, high_conf_count = optimize_threshold(y_test, probs, preds)
                    results.append(SearchResult(
                        code=code,
                        league=league,
                        feature_mode=feature_mode,
                        feature_set=feature_set_name,
                        train_window=window_name,
                        model_name=wrapper.name,
                        threshold=threshold,
                        bet_acc=bet_acc,
                        bet_count=bet_count,
                        high_conf_acc=high_conf_acc,
                        high_conf_count=high_conf_count,
                        test_acc=test_acc,
                        test_count=int(len(y_test)),
                    ))

                    if wrapper.name in {'LR_l2_c03', 'RF_d6_n600_l3', 'ET_d7_n700_l2'}:
                        base_est = clone(wrapper.estimator)
                        cal = CalibratedClassifierCV(base_est, method='sigmoid', cv=3)
                        scale = StandardScaler() if wrapper.scale else FunctionTransformer(validate=False)
                        cal_pipe = Pipeline([
                            ('scaler', scale),
                            ('model', cal),
                        ])
                        cal_pipe.fit(Xtr, y_train_window)
                        probs = cal_pipe.predict_proba(Xte)[:, 1]
                        preds = (probs >= 0.5).astype(int)
                        test_acc = float(accuracy_score(y_test, preds))
                        threshold, bet_acc, bet_count, high_conf_acc, high_conf_count = optimize_threshold(y_test, probs, preds)
                        results.append(SearchResult(
                            code=code,
                            league=league,
                            feature_mode=feature_mode,
                            feature_set=feature_set_name,
                            train_window=window_name,
                            model_name=wrapper.name + '_cal',
                            threshold=threshold,
                            bet_acc=bet_acc,
                            bet_count=bet_count,
                            high_conf_acc=high_conf_acc,
                            high_conf_count=high_conf_count,
                            test_acc=test_acc,
                            test_count=int(len(y_test)),
                        ))

    output_dir.mkdir(parents=True, exist_ok=True)
    out_csv = output_dir / f'{code}_search_results.csv'
    out_json = output_dir / f'{code}_top20.json'
    df_out = pd.DataFrame([r.__dict__ for r in results]).sort_values(['bet_acc', 'bet_count', 'high_conf_acc'], ascending=[False, False, False])
    df_out.to_csv(out_csv, index=False, encoding='utf-8-sig')
    out_json.write_text(df_out.head(20).to_json(orient='records', force_ascii=False, indent=2), encoding='utf-8')
    return results


def main() -> None:
    warnings.filterwarnings('ignore')
    parser = argparse.ArgumentParser(description='Targeted deep optimization for selected leagues only.')
    parser.add_argument('--codes', nargs='+', default=['d1', 'e2n'])
    parser.add_argument('--output-dir', type=Path, default=DEFAULT_OUTPUT)
    args = parser.parse_args()

    summary = {}
    for code in args.codes:
        results = evaluate_league(code, args.output_dir)
        best = sorted(results, key=lambda r: (r.bet_acc, r.bet_count, r.high_conf_acc), reverse=True)[0]
        summary[code] = best.__dict__
        print(json.dumps(best.__dict__, ensure_ascii=False))

    (args.output_dir / 'summary.json').write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding='utf-8')


if __name__ == '__main__':
    main()
