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

import argparse
import json
import shutil
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Tuple

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
from sklearn.preprocessing import FunctionTransformer, StandardScaler

from train_models import LEAGUES, safe_to_datetime

REPO_ROOT = Path(__file__).resolve().parents[1]
BASELINE_DIR = REPO_ROOT / 'artifacts' / 'v4_models_frozen_2026-03-20'
DEFAULT_OUTPUT_DIR = REPO_ROOT / 'experiments' / 'priority_opt_pass1'
DEFAULT_INPUT_JSON = BASELINE_DIR / 'priority_ge75.json'
SERVING_DIR = REPO_ROOT / 'artifacts' / 'v4_models_serving_current'

DEFAULT_ELO = 1500.0
HOME_ELO_BONUS = 60.0
K_FACTOR = 20.0


@dataclass
class EvalResult:
    code: str
    league: str
    baseline_model: str
    baseline_bet_acc: float
    baseline_threshold: float
    candidate_model: str
    candidate_bet_acc: float
    candidate_threshold: float
    high_conf_acc: float
    high_conf_count: int
    feature_count: int
    feature_mode: str
    improved: bool


class ModelWrapper:
    def __init__(self, name: str, estimator, needs_scaling: bool):
        self.name = name
        self.estimator = estimator
        self.needs_scaling = needs_scaling
        self.scaler = StandardScaler() if needs_scaling else FunctionTransformer(validate=False)

    def fit(self, X_train: pd.DataFrame, y_train: pd.Series):
        X = self.scaler.fit_transform(X_train)
        self.estimator.fit(X, y_train)
        return self

    def predict(self, X: pd.DataFrame) -> np.ndarray:
        data = self.scaler.transform(X)
        return self.estimator.predict(data)

    def predict_proba(self, X: pd.DataFrame) -> np.ndarray:
        data = self.scaler.transform(X)
        if hasattr(self.estimator, 'predict_proba'):
            return self.estimator.predict_proba(data)
        scores = self.estimator.decision_function(data)
        probs_1 = 1 / (1 + np.exp(-scores))
        return np.column_stack([1 - probs_1, probs_1])


def rolling_mean(values: List[float], n: int, default: float) -> float:
    if not values:
        return default
    return float(np.mean(values[-n:]))



def calc_expected(home_elo: float, away_elo: float) -> float:
    return 1.0 / (1.0 + 10 ** ((away_elo - home_elo) / 400.0))



def extract_odds(df: pd.DataFrame) -> Tuple[pd.Series, pd.Series, pd.Series]:
    home = df['AvgH'] if 'AvgH' in df.columns else df.get('B365H', pd.Series(2.0, index=df.index))
    draw = df['AvgD'] if 'AvgD' in df.columns else df.get('B365D', pd.Series(np.nan, index=df.index))
    away = df['AvgA'] if 'AvgA' in df.columns else df.get('B365A', pd.Series(3.0, index=df.index))
    return home.fillna(2.0), draw.fillna(3.2), away.fillna(3.0)



def create_enhanced_features(df: pd.DataFrame) -> pd.DataFrame:
    home_odds, draw_odds, away_odds = extract_odds(df)
    handicap = df['AHh'].fillna(0) if 'AHh' in df.columns else pd.Series(0.0, index=df.index)
    home_water = df.get('B365AHH', pd.Series(0.9, index=df.index)).fillna(0.9)
    away_water = df.get('B365AHA', pd.Series(0.9, index=df.index)).fillna(0.9)

    records: List[dict] = []
    team_stats: Dict[str, dict] = {}
    elo: Dict[str, float] = {}

    def get_team(team: str) -> dict:
        if team not in team_stats:
            team_stats[team] = {
                'pts_all': [], 'pts_home': [], 'pts_away': [],
                'gd_all': [], 'gf_all': [], 'ga_all': [],
                'gf_home': [], 'ga_home': [], 'gf_away': [], 'ga_away': [],
                'sot_for_home': [], 'sot_against_home': [], 'sot_for_away': [], 'sot_against_away': [],
                'shots_for_home': [], 'shots_against_home': [], 'shots_for_away': [], 'shots_against_away': [],
                'corners_for_home': [], 'corners_against_home': [], 'corners_for_away': [], 'corners_against_away': [],
            }
        return team_stats[team]

    for idx, row in df.iterrows():
        home = row['HomeTeam']
        away = row['AwayTeam']
        hs = get_team(home)
        aw = get_team(away)
        home_elo = elo.get(home, DEFAULT_ELO)
        away_elo = elo.get(away, DEFAULT_ELO)

        hp = 1.0 / home_odds.iloc[idx] if home_odds.iloc[idx] else 0.5
        ap = 1.0 / away_odds.iloc[idx] if away_odds.iloc[idx] else 0.5
        dp = 1.0 / draw_odds.iloc[idx] if draw_odds.iloc[idx] else 0.0
        overround = hp + ap + dp if dp else hp + ap
        hp_norm = hp / overround if overround else hp
        ap_norm = ap / overround if overround else ap
        dp_norm = dp / overround if overround else dp

        rec = {
            'avg_home_odds': home_odds.iloc[idx],
            'avg_draw_odds': draw_odds.iloc[idx],
            'avg_away_odds': away_odds.iloc[idx],
            'home_prob': hp,
            'away_prob': ap,
            'draw_prob': dp,
            'home_prob_norm': hp_norm,
            'away_prob_norm': ap_norm,
            'draw_prob_norm': dp_norm,
            'market_overround': overround,
            'prob_diff': hp - ap,
            'prob_ratio': hp / max(ap, 1e-6),
            'draw_bias': dp_norm,
            'handicap': handicap.iloc[idx],
            'handicap_abs': abs(handicap.iloc[idx]),
            'asian_home_water': home_water.iloc[idx],
            'asian_away_water': away_water.iloc[idx],
            'water_diff': home_water.iloc[idx] - away_water.iloc[idx],
            'water_sum': home_water.iloc[idx] + away_water.iloc[idx],
            'water_ratio': home_water.iloc[idx] / max(away_water.iloc[idx], 1e-6),
            'handicap_x_water': handicap.iloc[idx] * home_water.iloc[idx],
            'home_ppg': rolling_mean(hs['pts_all'], 9999, 1.5),
            'away_ppg': rolling_mean(aw['pts_all'], 9999, 1.5),
            'home_form': rolling_mean(hs['pts_all'], 5, 1.5),
            'away_form': rolling_mean(aw['pts_all'], 5, 1.5),
            'home_home_ppg': rolling_mean(hs['pts_home'], 9999, 1.5),
            'away_away_ppg': rolling_mean(aw['pts_away'], 9999, 1.5),
            'home_home_form': rolling_mean(hs['pts_home'], 5, 1.5),
            'away_away_form': rolling_mean(aw['pts_away'], 5, 1.5),
            'home_gd': rolling_mean(hs['gd_all'], 9999, 0.0),
            'away_gd': rolling_mean(aw['gd_all'], 9999, 0.0),
            'home_scored': rolling_mean(hs['gf_all'], 9999, 1.2),
            'away_scored': rolling_mean(aw['gf_all'], 9999, 1.2),
            'home_conceded': rolling_mean(hs['ga_all'], 9999, 1.2),
            'away_conceded': rolling_mean(aw['ga_all'], 9999, 1.2),
            'home_scored_home': rolling_mean(hs['gf_home'], 5, 1.2),
            'away_scored_away': rolling_mean(aw['gf_away'], 5, 1.2),
            'home_conceded_home': rolling_mean(hs['ga_home'], 5, 1.2),
            'away_conceded_away': rolling_mean(aw['ga_away'], 5, 1.2),
            'home_sot_for': rolling_mean(hs['sot_for_home'], 5, 4.0),
            'away_sot_for': rolling_mean(aw['sot_for_away'], 5, 4.0),
            'home_sot_against': rolling_mean(hs['sot_against_home'], 5, 4.0),
            'away_sot_against': rolling_mean(aw['sot_against_away'], 5, 4.0),
            'home_shots_for': rolling_mean(hs['shots_for_home'], 5, 10.0),
            'away_shots_for': rolling_mean(aw['shots_for_away'], 5, 10.0),
            'home_corners_for': rolling_mean(hs['corners_for_home'], 5, 5.0),
            'away_corners_for': rolling_mean(aw['corners_for_away'], 5, 5.0),
            'home_elo': home_elo,
            'away_elo': away_elo,
        }

        rec['ppg_diff'] = rec['home_ppg'] - rec['away_ppg']
        rec['form_diff'] = rec['home_form'] - rec['away_form']
        rec['venue_ppg_diff'] = rec['home_home_ppg'] - rec['away_away_ppg']
        rec['venue_form_diff'] = rec['home_home_form'] - rec['away_away_form']
        rec['gd_diff'] = rec['home_gd'] - rec['away_gd']
        rec['scored_diff'] = rec['home_scored'] - rec['away_scored']
        rec['conceded_diff'] = rec['away_conceded'] - rec['home_conceded']
        rec['sot_for_diff'] = rec['home_sot_for'] - rec['away_sot_for']
        rec['sot_against_diff'] = rec['away_sot_against'] - rec['home_sot_against']
        rec['shots_for_diff'] = rec['home_shots_for'] - rec['away_shots_for']
        rec['corners_for_diff'] = rec['home_corners_for'] - rec['away_corners_for']
        rec['home_adv'] = rolling_mean([1.0 if p == 3 else 0.0 for p in hs['pts_home']], 9999, 0.45)
        rec['adv_diff'] = rec['home_adv'] - 0.45
        rec['elo_diff'] = (home_elo + HOME_ELO_BONUS) - away_elo
        rec['prob_x_form'] = rec['home_prob'] * rec['home_form']
        rec['form_x_ppg'] = rec['home_form'] * rec['home_ppg']
        rec['prob_x_handicap'] = rec['home_prob'] * rec['handicap']
        rec['water_x_prob'] = rec['asian_home_water'] * rec['home_prob']
        rec['prob_x_elo'] = rec['home_prob'] * rec['elo_diff']
        rec['elo_x_handicap'] = rec['elo_diff'] * rec['handicap']
        rec['venue_strength'] = rec['venue_ppg_diff'] + rec['venue_form_diff']
        records.append(rec)

        fthg = row['FTHG']
        ftag = row['FTAG']
        if fthg > ftag:
            home_pts, away_pts = 3, 0
            score_home = 1.0
        elif fthg < ftag:
            home_pts, away_pts = 0, 3
            score_home = 0.0
        else:
            home_pts, away_pts = 1, 1
            score_home = 0.5

        hs['pts_all'].append(home_pts)
        hs['pts_home'].append(home_pts)
        aw['pts_all'].append(away_pts)
        aw['pts_away'].append(away_pts)
        hs['gd_all'].append(fthg - ftag)
        aw['gd_all'].append(ftag - fthg)
        hs['gf_all'].append(fthg)
        hs['ga_all'].append(ftag)
        aw['gf_all'].append(ftag)
        aw['ga_all'].append(fthg)
        hs['gf_home'].append(fthg)
        hs['ga_home'].append(ftag)
        aw['gf_away'].append(ftag)
        aw['ga_away'].append(fthg)

        hst = row['HST'] if 'HST' in df.columns and pd.notna(row.get('HST')) else 4.0
        ast = row['AST'] if 'AST' in df.columns and pd.notna(row.get('AST')) else 4.0
        hs_ = row['HS'] if 'HS' in df.columns and pd.notna(row.get('HS')) else 10.0
        as_ = row['AS'] if 'AS' in df.columns and pd.notna(row.get('AS')) else 10.0
        hc = row['HC'] if 'HC' in df.columns and pd.notna(row.get('HC')) else 5.0
        ac = row['AC'] if 'AC' in df.columns and pd.notna(row.get('AC')) else 5.0
        hs['sot_for_home'].append(hst)
        hs['sot_against_home'].append(ast)
        aw['sot_for_away'].append(ast)
        aw['sot_against_away'].append(hst)
        hs['shots_for_home'].append(hs_)
        hs['shots_against_home'].append(as_)
        aw['shots_for_away'].append(as_)
        aw['shots_against_away'].append(hs_)
        hs['corners_for_home'].append(hc)
        hs['corners_against_home'].append(ac)
        aw['corners_for_away'].append(ac)
        aw['corners_against_away'].append(hc)

        expected_home = calc_expected(home_elo + HOME_ELO_BONUS, away_elo)
        elo[home] = home_elo + K_FACTOR * (score_home - expected_home)
        elo[away] = away_elo + K_FACTOR * ((1.0 - score_home) - (1.0 - expected_home))

    feat_df = pd.DataFrame.from_records(records, index=df.index)
    return feat_df.fillna(0).replace([np.inf, -np.inf], 0)



def optimize_threshold(y_true: pd.Series, probs: np.ndarray, preds: np.ndarray, fallback: float) -> tuple[float, float, float, int]:
    best_threshold = 0.55
    best_bet_acc = 0.0
    best_high_conf_acc = 0.0
    best_high_conf_count = 0
    for threshold in np.arange(0.45, 0.76, 0.01):
        mask = (probs > threshold) | (probs < (1 - threshold))
        if mask.sum() < 10:
            continue
        bet_acc = accuracy_score(y_true[mask], preds[mask])
        high_mask = (probs > max(0.70, threshold + 0.05)) | (probs < min(0.30, 1 - threshold - 0.05))
        if high_mask.sum() > 0:
            high_acc = accuracy_score(y_true[high_mask], preds[high_mask])
            high_count = int(high_mask.sum())
        else:
            high_acc, high_count = 0.0, 0
        if bet_acc > best_bet_acc:
            best_threshold = float(round(threshold, 2))
            best_bet_acc = float(bet_acc)
            best_high_conf_acc = float(high_acc)
            best_high_conf_count = high_count
    if best_bet_acc < 0.5:
        best_bet_acc = float(fallback)
    return best_threshold, best_bet_acc, best_high_conf_acc, best_high_conf_count



def build_candidates() -> List[ModelWrapper]:
    return [
        ModelWrapper('LR_c03', LogisticRegression(C=0.3, max_iter=2000, class_weight='balanced', random_state=42), True),
        ModelWrapper('LR_c05', LogisticRegression(C=0.5, max_iter=2000, class_weight='balanced', random_state=42), True),
        ModelWrapper('LR_c10', LogisticRegression(C=1.0, max_iter=2000, class_weight='balanced', random_state=42), True),
        ModelWrapper('RF_d7_n500', RandomForestClassifier(n_estimators=500, max_depth=7, min_samples_leaf=2, random_state=42, n_jobs=-1), False),
        ModelWrapper('RF_d10_n700', RandomForestClassifier(n_estimators=700, max_depth=10, 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('GB_d5', GradientBoostingClassifier(n_estimators=250, learning_rate=0.05, max_depth=5, random_state=42), False),
        ModelWrapper('HGB_d6', HistGradientBoostingClassifier(max_depth=6, max_iter=350, learning_rate=0.05, min_samples_leaf=20, random_state=42), True),
    ]



def load_input(path: Path) -> List[dict]:
    return json.loads(path.read_text(encoding='utf-8'))



def save_candidate_artifacts(code: str, league: str, wrapper: ModelWrapper, result: EvalResult, 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.candidate_threshold,
        'test_acc': result.candidate_bet_acc,
        'bet_acc': result.candidate_bet_acc,
        'high_conf_acc': result.high_conf_acc,
        'high_conf_count': result.high_conf_count,
        'features': feature_names,
        'model_name': result.candidate_model,
        'data_count': None,
        'league': league,
        'feature_mode': result.feature_mode,
        'source': 'optimize_priority_models.py',
    }
    joblib.dump(config, model_dir / f'{code}_pro_v4_config.pkl')



def evaluate_league(code: str, league: str, baseline_model: str, baseline_acc: float, baseline_threshold: float, output_dir: Path) -> tuple[EvalResult, ModelWrapper, List[str]]:
    data_path = REPO_ROOT / 'historical_data' / f'{code}_history.csv'
    df = pd.read_csv(data_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

    features = create_enhanced_features(df)
    feature_names = list(features.columns)
    split_idx = int(len(df) * 0.8)
    X_train = features.iloc[:split_idx]
    X_test = features.iloc[split_idx:]
    y_train = df['asian_result'].iloc[:split_idx]
    y_test = df['asian_result'].iloc[split_idx:]

    best_payload = None
    best_wrapper = None

    for wrapper in build_candidates():
        wrapper.fit(X_train, y_train)
        probs = wrapper.predict_proba(X_test)[:, 1]
        preds = (probs > 0.5).astype(int)
        fallback_bet_acc = accuracy_score(y_test, preds)
        threshold, bet_acc, high_conf_acc, high_conf_count = optimize_threshold(y_test.reset_index(drop=True), probs, preds, fallback_bet_acc)
        payload = {
            'candidate_model': wrapper.name,
            'candidate_bet_acc': float(bet_acc),
            'candidate_threshold': float(threshold),
            'high_conf_acc': float(high_conf_acc),
            'high_conf_count': int(high_conf_count),
            'feature_count': int(features.shape[1]),
            'feature_mode': 'enhanced_v1',
        }
        if best_payload is None or payload['candidate_bet_acc'] > best_payload['candidate_bet_acc']:
            best_payload = payload
            best_wrapper = wrapper

    result = EvalResult(
        code=code,
        league=league,
        baseline_model=baseline_model,
        baseline_bet_acc=float(baseline_acc),
        baseline_threshold=float(baseline_threshold),
        candidate_model=best_payload['candidate_model'],
        candidate_bet_acc=float(best_payload['candidate_bet_acc']),
        candidate_threshold=float(best_payload['candidate_threshold']),
        high_conf_acc=float(best_payload['high_conf_acc']),
        high_conf_count=int(best_payload['high_conf_count']),
        feature_count=int(best_payload['feature_count']),
        feature_mode=best_payload['feature_mode'],
        improved=best_payload['candidate_bet_acc'] > float(baseline_acc),
    )

    if result.improved:
        save_candidate_artifacts(code, league, best_wrapper, result, feature_names, output_dir)
    return result, best_wrapper, feature_names



def build_serving_dir(output_dir: Path, summary: List[dict]) -> None:
    SERVING_DIR.mkdir(parents=True, exist_ok=True)
    for file in SERVING_DIR.glob('*'):
        if file.is_file():
            file.unlink()
    for file in BASELINE_DIR.glob('*_pro_v4_*.pkl'):
        shutil.copy2(file, SERVING_DIR / file.name)

    improved_dir = output_dir / 'improved_models'
    if improved_dir.exists():
        for row in summary:
            if not row['improved']:
                continue
            code = row['code']
            for suffix in ['config', 'model', 'scaler']:
                src = improved_dir / f'{code}_pro_v4_{suffix}.pkl'
                if src.exists():
                    shutil.copy2(src, SERVING_DIR / src.name)

    manifest = []
    for file in sorted(SERVING_DIR.glob('*_pro_v4_config.pkl')):
        code = file.name.split('_pro_v4_config.pkl')[0]
        cfg = joblib.load(file)
        manifest.append({
            'code': code,
            'league': cfg.get('league', LEAGUES.get(code, code)),
            'model_name': cfg.get('model_name'),
            'bet_acc': float(cfg.get('bet_acc', 0.0)),
            'threshold': float(cfg.get('threshold', 0.55)),
            'feature_mode': cfg.get('feature_mode', 'baseline_v4'),
            'feature_count': len(cfg.get('features', [])),
        })
    (SERVING_DIR / 'manifest.json').write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding='utf-8')



def main() -> None:
    parser = argparse.ArgumentParser(description='Optimize priority football models and deploy best versions.')
    parser.add_argument('--input-json', type=Path, default=DEFAULT_INPUT_JSON)
    parser.add_argument('--output-dir', type=Path, default=DEFAULT_OUTPUT_DIR)
    args = parser.parse_args()

    args.output_dir.mkdir(parents=True, exist_ok=True)
    priority = load_input(args.input_json)
    summary = []

    for row in priority:
        result, _, _ = evaluate_league(
            row['code'],
            row['league'],
            row['model_name'],
            row['bet_acc'],
            row['threshold'],
            args.output_dir,
        )
        summary.append(result.__dict__)
        print(json.dumps(result.__dict__, ensure_ascii=False))

    out_json = args.output_dir / 'summary.json'
    out_csv = args.output_dir / 'summary.csv'
    out_json.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding='utf-8')
    pd.DataFrame(summary).sort_values('candidate_bet_acc', ascending=False).to_csv(out_csv, index=False, encoding='utf-8-sig')
    build_serving_dir(args.output_dir, summary)
    print(f'SAVED {out_json}')
    print(f'SAVED {out_csv}')
    print(f'SERVED {SERVING_DIR}')


if __name__ == '__main__':
    main()
