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

import argparse
import json
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Dict, List

import joblib
import numpy as np
import pandas as pd
from sklearn.ensemble import (
    ExtraTreesClassifier,
    GradientBoostingClassifier,
    HistGradientBoostingClassifier,
    RandomForestClassifier,
)
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

REPO_ROOT = Path(__file__).resolve().parents[1]
SCRIPTS_DIR = REPO_ROOT / 'scripts'
if str(SCRIPTS_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPTS_DIR))

from optimize_priority_models import ModelWrapper, create_enhanced_features  # noqa: E402
from train_models import LEAGUES, create_features, safe_to_datetime  # noqa: E402

BASELINE_DIR = REPO_ROOT / 'artifacts' / 'v4_models_frozen_2026-03-20'
DEFAULT_OUTPUT_DIR = REPO_ROOT / 'experiments' / 'd1_e2n_opt_pass2'
DEFAULT_CODES = ['d1', 'e2n']


@dataclass
class SearchResult:
    code: str
    league: str
    baseline_model: str
    baseline_bet_acc: float
    baseline_threshold: float
    feature_mode: str
    feature_count: int
    model_name: str
    threshold: float
    bet_acc: float
    bet_count: int
    full_test_acc: float
    high_conf_acc: float
    high_conf_count: int
    improved_vs_baseline: bool


@dataclass
class Bundle:
    name: str
    frame: pd.DataFrame


THRESHOLDS = np.arange(0.50, 0.91, 0.01)
MIN_BET_COUNT = 10
MAX_FEATURE_MODES = {
    'baseline_v4',
    'enhanced_v1',
    'enhanced_compact',
    'enhanced_market_form',
    'enhanced_no_shot',
    'enhanced_no_elo',
}


def load_baseline(code: str) -> dict:
    return joblib.load(BASELINE_DIR / f'{code}_pro_v4_config.pkl')


def load_league_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']).reset_index(drop=True)
    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 build_feature_bundles(df: pd.DataFrame) -> list[Bundle]:
    base = create_features(df).copy()
    enhanced = create_enhanced_features(df).copy()

    shot_cols = [c for c in enhanced.columns if any(k in c for k in ['sot', 'shots', 'corners'])]
    elo_cols = [c for c in enhanced.columns if 'elo' in c]

    market_cols = [
        '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',
    ]
    form_cols = [
        'home_ppg', 'away_ppg', 'home_form', 'away_form',
        'home_home_ppg', 'away_away_ppg', 'home_home_form', 'away_away_form',
        'home_gd', 'away_gd', 'home_scored', 'away_scored', 'home_conceded', 'away_conceded',
        'home_scored_home', 'away_scored_away', 'home_conceded_home', 'away_conceded_away',
        'ppg_diff', 'form_diff', 'venue_ppg_diff', 'venue_form_diff',
        'gd_diff', 'scored_diff', 'conceded_diff', 'venue_strength',
        'home_adv', 'adv_diff',
    ]
    interaction_cols = [
        'prob_x_form', 'form_x_ppg', 'prob_x_handicap', 'water_x_prob', 'prob_x_elo', 'elo_x_handicap',
    ]
    compact_cols = [
        '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',
        'home_home_ppg', 'away_away_ppg', 'ppg_diff', 'form_diff', 'venue_ppg_diff', 'venue_form_diff',
        'home_adv', 'adv_diff',
        'prob_x_form', 'form_x_ppg', 'prob_x_handicap', 'water_x_prob',
    ]

    def pick(cols: list[str]) -> pd.DataFrame:
        return enhanced[[c for c in cols if c in enhanced.columns]].copy()

    bundles = [
        Bundle('baseline_v4', base),
        Bundle('enhanced_v1', enhanced),
        Bundle('enhanced_compact', pick(compact_cols)),
        Bundle('enhanced_market_form', pick(market_cols + form_cols + interaction_cols + elo_cols)),
        Bundle('enhanced_no_shot', enhanced[[c for c in enhanced.columns if c not in shot_cols]].copy()),
        Bundle('enhanced_no_elo', enhanced[[c for c in enhanced.columns if c not in elo_cols]].copy()),
        Bundle('enhanced_market_only', pick(market_cols + interaction_cols)),
        Bundle('enhanced_form_only', pick(form_cols + interaction_cols + elo_cols)),
    ]
    return bundles


def build_candidates() -> list[ModelWrapper]:
    return [
        ModelWrapper('LR_c003_bal', LogisticRegression(C=0.03, max_iter=4000, class_weight='balanced', random_state=42), True),
        ModelWrapper('LR_c01_bal', LogisticRegression(C=0.10, max_iter=4000, class_weight='balanced', random_state=42), True),
        ModelWrapper('LR_c03_bal', LogisticRegression(C=0.30, max_iter=4000, class_weight='balanced', random_state=42), True),
        ModelWrapper('LR_c05_bal', LogisticRegression(C=0.50, max_iter=4000, class_weight='balanced', random_state=42), True),
        ModelWrapper('LR_c10_bal', LogisticRegression(C=1.00, max_iter=4000, class_weight='balanced', random_state=42), True),
        ModelWrapper('LR_c03_plain', LogisticRegression(C=0.30, max_iter=4000, class_weight=None, random_state=42), True),
        ModelWrapper('RF_d5_n500', RandomForestClassifier(n_estimators=500, max_depth=5, min_samples_leaf=2, random_state=42, n_jobs=-1), False),
        ModelWrapper('RF_d7_n700', RandomForestClassifier(n_estimators=700, max_depth=7, min_samples_leaf=2, random_state=42, n_jobs=-1), False),
        ModelWrapper('RF_d10_n1000', RandomForestClassifier(n_estimators=1000, max_depth=10, min_samples_leaf=2, random_state=42, n_jobs=-1), False),
        ModelWrapper('RF_full_n1000', RandomForestClassifier(n_estimators=1000, max_depth=None, min_samples_leaf=2, random_state=42, n_jobs=-1), False),
        ModelWrapper('ET_d8_n700', ExtraTreesClassifier(n_estimators=700, max_depth=8, min_samples_leaf=2, random_state=42, n_jobs=-1), False),
        ModelWrapper('ET_n700', ExtraTreesClassifier(n_estimators=700, max_depth=None, min_samples_leaf=2, random_state=42, n_jobs=-1), False),
        ModelWrapper('ET_n1200', ExtraTreesClassifier(n_estimators=1200, max_depth=None, min_samples_leaf=2, random_state=42, n_jobs=-1), False),
        ModelWrapper('GB_d3_lr003', GradientBoostingClassifier(n_estimators=500, learning_rate=0.03, max_depth=3, random_state=42), False),
        ModelWrapper('GB_d5_lr005', GradientBoostingClassifier(n_estimators=250, learning_rate=0.05, max_depth=5, random_state=42), False),
        ModelWrapper('HGB_d4', HistGradientBoostingClassifier(max_depth=4, max_iter=450, learning_rate=0.03, min_samples_leaf=12, random_state=42), True),
        ModelWrapper('HGB_d6', HistGradientBoostingClassifier(max_depth=6, max_iter=350, learning_rate=0.05, min_samples_leaf=20, random_state=42), True),
    ]


def optimize_threshold(y_true: np.ndarray, probs: np.ndarray, preds: np.ndarray) -> tuple[float, float, int, float, int]:
    best: tuple[float, float, int, float, int] | None = None
    for threshold in THRESHOLDS:
        mask = (probs >= threshold) | (probs <= (1.0 - threshold))
        bet_count = int(mask.sum())
        if bet_count < MIN_BET_COUNT:
            continue
        bet_acc = float(accuracy_score(y_true[mask], preds[mask]))

        high_t = min(0.95, threshold + 0.05)
        low_t = max(0.05, 1.0 - threshold - 0.05)
        high_mask = (probs >= high_t) | (probs <= low_t)
        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

        candidate = (round(float(threshold), 2), bet_acc, bet_count, high_acc, high_count)
        if best is None:
            best = candidate
            continue
        _, best_acc, best_count, _, _ = best
        if bet_acc > best_acc or (bet_acc == best_acc and bet_count > best_count):
            best = candidate
    if best is None:
        full_acc = float(accuracy_score(y_true, preds))
        return 0.55, full_acc, int(len(y_true)), full_acc, int(len(y_true))
    return best


def save_artifacts(code: str, league: str, wrapper: ModelWrapper, result: SearchResult, feature_names: list[str], output_dir: Path) -> None:
    model_dir = output_dir / 'improved_models'
    model_dir.mkdir(parents=True, exist_ok=True)
    joblib.dump(wrapper.scaler, model_dir / f'{code}_pro_v4_scaler.pkl')
    joblib.dump(wrapper.estimator, model_dir / f'{code}_pro_v4_model.pkl')
    config = {
        'threshold': result.threshold,
        'test_acc': result.bet_acc,
        'bet_acc': result.bet_acc,
        'bet_count': result.bet_count,
        'high_conf_acc': result.high_conf_acc,
        'high_conf_count': result.high_conf_count,
        'features': feature_names,
        'model_name': result.model_name,
        'data_count': None,
        'league': league,
        'feature_mode': result.feature_mode,
        'source': 'targeted_optimize_d1_e2n.py',
        'search_scope': 'd1_e2n_only',
    }
    joblib.dump(config, model_dir / f'{code}_pro_v4_config.pkl')


def evaluate_league(code: str, output_dir: Path) -> dict:
    baseline = load_baseline(code)
    league = LEAGUES.get(code, code)
    df = load_league_df(code)
    split_idx = int(len(df) * 0.8)
    y_train = df['asian_result'].iloc[:split_idx]
    y_test = df['asian_result'].iloc[split_idx:].to_numpy()

    rows: list[SearchResult] = []
    best_result: SearchResult | None = None
    best_wrapper: ModelWrapper | None = None
    best_features: list[str] = []

    for bundle in build_feature_bundles(df):
        X_train = bundle.frame.iloc[:split_idx]
        X_test = bundle.frame.iloc[split_idx:]
        feature_names = list(bundle.frame.columns)

        for template in build_candidates():
            wrapper = ModelWrapper(template.name, template.estimator, template.needs_scaling)
            wrapper.fit(X_train, y_train)
            probs = wrapper.predict_proba(X_test)[:, 1]
            preds = (probs >= 0.5).astype(int)
            threshold, bet_acc, bet_count, high_conf_acc, high_conf_count = optimize_threshold(y_test, probs, preds)
            full_test_acc = float(accuracy_score(y_test, preds))
            row = SearchResult(
                code=code,
                league=league,
                baseline_model=str(baseline.get('model_name')),
                baseline_bet_acc=float(baseline.get('bet_acc', 0.0)),
                baseline_threshold=float(baseline.get('threshold', 0.55)),
                feature_mode=bundle.name,
                feature_count=int(bundle.frame.shape[1]),
                model_name=wrapper.name,
                threshold=float(threshold),
                bet_acc=float(bet_acc),
                bet_count=int(bet_count),
                full_test_acc=full_test_acc,
                high_conf_acc=float(high_conf_acc),
                high_conf_count=int(high_conf_count),
                improved_vs_baseline=float(bet_acc) > float(baseline.get('bet_acc', 0.0)),
            )
            rows.append(row)
            if best_result is None or row.bet_acc > best_result.bet_acc or (row.bet_acc == best_result.bet_acc and row.bet_count > best_result.bet_count):
                best_result = row
                best_wrapper = wrapper
                best_features = feature_names

    assert best_result is not None and best_wrapper is not None
    save_artifacts(code, league, best_wrapper, best_result, best_features, output_dir)

    rows_sorted = sorted(rows, key=lambda x: (x.bet_acc, x.bet_count, -abs(x.threshold - 0.70)), reverse=True)
    frontier = {}
    for minimum in [10, 15, 20, 30, 40, 50, 60, 80, 100]:
        feasible = [r for r in rows_sorted if r.bet_count >= minimum]
        if feasible:
            pick = feasible[0]
            frontier[str(minimum)] = {
                'model_name': pick.model_name,
                'feature_mode': pick.feature_mode,
                'feature_count': pick.feature_count,
                'threshold': pick.threshold,
                'bet_acc': pick.bet_acc,
                'bet_count': pick.bet_count,
                'high_conf_acc': pick.high_conf_acc,
                'high_conf_count': pick.high_conf_count,
            }

    return {
        'code': code,
        'league': league,
        'baseline': baseline,
        'best': best_result.__dict__,
        'frontier': frontier,
        'top20': [r.__dict__ for r in rows_sorted[:20]],
        'total_candidates': len(rows_sorted),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description='Targeted wider offline search for d1 and e2n only.')
    parser.add_argument('--codes', nargs='+', default=DEFAULT_CODES)
    parser.add_argument('--output-dir', type=Path, default=DEFAULT_OUTPUT_DIR)
    args = parser.parse_args()

    args.output_dir.mkdir(parents=True, exist_ok=True)
    summaries = []
    for code in args.codes:
        summary = evaluate_league(code, args.output_dir)
        summaries.append(summary)
        print(json.dumps(summary['best'], ensure_ascii=False))

    (args.output_dir / 'summary.json').write_text(json.dumps(summaries, ensure_ascii=False, indent=2), encoding='utf-8')
    top_rows = []
    for summary in summaries:
        for row in summary['top20']:
            top_rows.append(row)
    pd.DataFrame(top_rows).to_csv(args.output_dir / 'top20.csv', index=False, encoding='utf-8-sig')
    print(f'SAVED {args.output_dir / "summary.json"}')
    print(f'SAVED {args.output_dir / "top20.csv"}')


if __name__ == '__main__':
    main()
