#!/usr/bin/env python3
"""Train goal-number models with odds features and advanced parameters."""

from __future__ import annotations

import json
import math
from collections import defaultdict
from pathlib import Path
from typing import List

import joblib
import numpy as np
import pandas as pd
from sklearn.ensemble import (
    RandomForestClassifier, 
    GradientBoostingClassifier, 
    AdaBoostClassifier,
    ExtraTreesClassifier
)
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, f1_score
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler, RobustScaler

REPO_ROOT = Path(__file__).resolve().parents[1]
HISTORICAL_DIR = REPO_ROOT / 'historical_data'
OUTPUT_DIR = REPO_ROOT / 'experiments' / 'goal_number_v2'

LEAGUES = {
    'a1': {'name': '澳超', 'min_rows': 1200},
    'br1': {'name': '巴甲', 'min_rows': 3000},
    'd2': {'name': '德乙', 'min_rows': 2400},
    'e0': {'name': '英超', 'min_rows': 2900},
    'f1': {'name': '法甲', 'min_rows': 1280},
    'f2': {'name': '法乙', 'min_rows': 2800},
    'g1': {'name': '希腊超', 'min_rows': 1870},
    'j1': {'name': '日职联', 'min_rows': 2620},
    'm1': {'name': '美职联', 'min_rows': 3500},
    'r1': {'name': '俄超', 'min_rows': 1870},
    'uz1': {'name': '乌兹别克超', 'min_rows': 780},
}


def mean_or_default(values: List[float], default: float) -> float:
    return float(np.mean(values)) if values else float(default)


def build_features_advanced(df: pd.DataFrame, has_odds: bool = False) -> pd.DataFrame:
    """Build advanced features from match history."""
    team_hist = defaultdict(lambda: {
        "gf": [], "ga": [], "tot": [], "over": [],
        "home_gf": [], "home_ga": [], "home_over": [],
        "away_gf": [], "away_ga": [], "away_over": [],
        "gf_last3": [], "ga_last3": [], "over_last3": [],
        "gf_last5": [], "ga_last5": [], "over_last5": [],
        "gf_last10": [], "ga_last10": [], "over_last10": [],
        "momentum": [],
    })
    
    rows = []
    
    for idx, match in df.iterrows():
        home = str(match["HomeTeam"])
        away = str(match["AwayTeam"])
        home_hist = team_hist[home]
        away_hist = team_hist[away]
        
        total_goals = float(match.get('FTHG', 0) or 0) + float(match.get('FTAG', 0) or 0)
        is_over = 1 if total_goals > 2.5 else 0
        
        row = {
            "home_gf_avg": mean_or_default(home_hist["gf"], 1.4),
            "home_ga_avg": mean_or_default(home_hist["ga"], 1.4),
            "away_gf_avg": mean_or_default(away_hist["gf"], 1.2),
            "away_ga_avg": mean_or_default(away_hist["ga"], 1.2),
            "home_tot_avg": mean_or_default(home_hist["tot"], 2.6),
            "away_tot_avg": mean_or_default(away_hist["tot"], 2.6),
            "home_over_rate": mean_or_default(home_hist["over"], 0.5),
            "away_over_rate": mean_or_default(away_hist["over"], 0.5),
        }
        
        # Rolling averages
        row["home_gf_last3"] = mean_or_default(home_hist["gf_last3"], 1.4)
        row["home_ga_last3"] = mean_or_default(home_hist["ga_last3"], 1.4)
        row["away_gf_last3"] = mean_or_default(away_hist["gf_last3"], 1.2)
        row["away_ga_last3"] = mean_or_default(away_hist["ga_last3"], 1.2)
        row["home_over_last3"] = mean_or_default(home_hist["over_last3"], 0.5)
        row["away_over_last3"] = mean_or_default(away_hist["over_last3"], 0.5)
        
        row["home_gf_last5"] = mean_or_default(home_hist["gf_last5"], 1.4)
        row["home_ga_last5"] = mean_or_default(home_hist["ga_last5"], 1.4)
        row["away_gf_last5"] = mean_or_default(away_hist["gf_last5"], 1.2)
        row["away_ga_last5"] = mean_or_default(away_hist["ga_last5"], 1.2)
        row["home_over_last5"] = mean_or_default(home_hist["over_last5"], 0.5)
        row["away_over_last5"] = mean_or_default(away_hist["over_last5"], 0.5)
        
        row["home_gf_last10"] = mean_or_default(home_hist["gf_last10"], 1.4)
        row["home_ga_last10"] = mean_or_default(home_hist["ga_last10"], 1.4)
        row["away_gf_last10"] = mean_or_default(away_hist["gf_last10"], 1.2)
        row["away_ga_last10"] = mean_or_default(away_hist["ga_last10"], 1.2)
        row["home_over_last10"] = mean_or_default(home_hist["over_last10"], 0.5)
        row["away_over_last10"] = mean_or_default(away_hist["over_last10"], 0.5)
        
        # Momentum features
        if len(home_hist["momentum"]) >= 5:
            row["home_momentum"] = np.mean(home_hist["momentum"][-5:])
        else:
            row["home_momentum"] = 0.5
        
        if len(away_hist["momentum"]) >= 5:
            row["away_momentum"] = np.mean(away_hist["momentum"][-5:])
        else:
            row["away_momentum"] = 0.5
        
        # Home/Away specific features
        row["home_home_gf_avg"] = mean_or_default(home_hist["home_gf"], mean_or_default(home_hist["gf"], 1.4))
        row["home_home_ga_avg"] = mean_or_default(home_hist["home_ga"], mean_or_default(home_hist["ga"], 1.4))
        row["away_away_gf_avg"] = mean_or_default(away_hist["away_gf"], mean_or_default(away_hist["gf"], 1.2))
        row["away_away_ga_avg"] = mean_or_default(away_hist["away_ga"], mean_or_default(away_hist["ga"], 1.2))
        row["home_home_over_rate"] = mean_or_default(home_hist["home_over"], mean_or_default(home_hist["over"], 0.5))
        row["away_away_over_rate"] = mean_or_default(away_hist["away_over"], mean_or_default(away_hist["over"], 0.5))
        
        # Variance features (consistency)
        if len(home_hist["gf_last10"]) >= 2:
            row["home_gf_var"] = np.var(home_hist["gf_last10"])
        else:
            row["home_gf_var"] = 0.5
        
        if len(home_hist["ga_last10"]) >= 2:
            row["home_ga_var"] = np.var(home_hist["ga_last10"])
        else:
            row["home_ga_var"] = 0.5
        
        if len(away_hist["gf_last10"]) >= 2:
            row["away_gf_var"] = np.var(away_hist["gf_last10"])
        else:
            row["away_gf_var"] = 0.5
        
        if len(away_hist["ga_last10"]) >= 2:
            row["away_ga_var"] = np.var(away_hist["ga_last10"])
        else:
            row["away_ga_var"] = 0.5
        
        # Expected goals
        expected_home = row["home_gf_avg"] + row["away_ga_avg"]
        expected_away = row["away_gf_avg"] + row["home_ga_avg"]
        row["expected_goals_proxy"] = (expected_home + expected_away) / 2
        row["combined_over_rate"] = float(np.nanmean([row["home_over_rate"], row["away_over_rate"]]))
        
        # Odds features
        if has_odds:
            for odds_prefix in ['B365', 'Avg', 'PS', 'Max']:
                over_col = odds_prefix + '>2.5'
                under_col = odds_prefix + '<2.5'
                
                if over_col in df.columns and under_col in df.columns:
                    over_odds = match.get(over_col, np.nan)
                    under_odds = match.get(under_col, np.nan)
                    
                    if pd.notna(over_odds) and pd.notna(under_odds):
                        row[f"{odds_prefix.lower()}_over_odds"] = float(over_odds)
                        row[f"{odds_prefix.lower()}_under_odds"] = float(under_odds)
                        
                        over_prob = 1.0 / float(over_odds)
                        under_prob = 1.0 / float(under_odds)
                        row[f"{odds_prefix.lower()}_over_prob"] = over_prob
                        row[f"{odds_prefix.lower()}_under_prob"] = under_prob
                        
                        fair_over = over_prob / (over_prob + under_prob)
                        row[f"{odds_prefix.lower()}_fair_over"] = fair_over
                        row[f"{odds_prefix.lower()}_margin"] = (over_prob + under_prob - 1.0)
                        break  # Use first available odds source
            
            # 1X2 odds
            for odds_prefix in ['B365', 'Avg', 'PS', 'Max']:
                home_col = odds_prefix + 'H'
                draw_col = odds_prefix + 'D'
                away_col = odds_prefix + 'A'
                
                if home_col in df.columns and draw_col in df.columns and away_col in df.columns:
                    home_odds = match.get(home_col, np.nan)
                    draw_odds = match.get(draw_col, np.nan)
                    away_odds = match.get(away_col, np.nan)
                    
                    if pd.notna(home_odds) and pd.notna(draw_odds) and pd.notna(away_odds):
                        row[f"{odds_prefix.lower()}_home_odds"] = float(home_odds)
                        row[f"{odds_prefix.lower()}_draw_odds"] = float(draw_odds)
                        row[f"{odds_prefix.lower()}_away_odds"] = float(away_odds)
                        
                        home_prob = 1.0 / float(home_odds)
                        draw_prob = 1.0 / float(draw_odds)
                        away_prob = 1.0 / float(away_odds)
                        prob_sum = home_prob + draw_prob + away_prob
                        
                        if prob_sum > 0:
                            row[f"{odds_prefix.lower()}_home_prob"] = home_prob / prob_sum
                            row[f"{odds_prefix.lower()}_draw_prob"] = draw_prob / prob_sum
                            row[f"{odds_prefix.lower()}_away_prob"] = away_prob / prob_sum
                            row[f"{odds_prefix.lower()}_home_away_gap"] = (home_prob - away_prob)
                        break
        
        row["is_over"] = is_over
        rows.append(row)
        
        # Update team history
        home_hist["gf"].append(total_goals)
        home_hist["ga"].append(total_goals)
        home_hist["tot"].append(total_goals)
        home_hist["over"].append(is_over)
        home_hist["gf_last3"].append(total_goals)
        home_hist["ga_last3"].append(total_goals)
        home_hist["gf_last5"].append(total_goals)
        home_hist["ga_last5"].append(total_goals)
        home_hist["gf_last10"].append(total_goals)
        home_hist["ga_last10"].append(total_goals)
        home_hist["over_last3"].append(is_over)
        home_hist["over_last5"].append(is_over)
        home_hist["over_last10"].append(is_over)
        
        home_result = 1 if total_goals > (match.get('FTAG', 0) or 0) else 0.5
        home_hist["momentum"].append(home_result)
        
        if pd.notna(match.get('FTHG')) and pd.notna(match.get('FTAG')):
            home_hist["home_gf"].append(float(match['FTHG']))
            home_hist["home_ga"].append(float(match['FTAG']))
            home_hist["home_over"].append(is_over)
        
        away_hist["gf"].append(total_goals)
        away_hist["ga"].append(total_goals)
        away_hist["tot"].append(total_goals)
        away_hist["over"].append(is_over)
        away_hist["gf_last3"].append(total_goals)
        away_hist["ga_last3"].append(total_goals)
        away_hist["gf_last5"].append(total_goals)
        away_hist["ga_last5"].append(total_goals)
        away_hist["gf_last10"].append(total_goals)
        away_hist["ga_last10"].append(total_goals)
        away_hist["over_last3"].append(is_over)
        away_hist["over_last5"].append(is_over)
        away_hist["over_last10"].append(is_over)
        
        away_result = 1 if total_goals > (match.get('FTHG', 0) or 0) else 0.5
        away_hist["momentum"].append(away_result)
        
        if pd.notna(match.get('FTHG')) and pd.notna(match.get('FTAG')):
            away_hist["away_gf"].append(float(match['FTHG']))
            away_hist["away_ga"].append(float(match['FTAG']))
            away_hist["away_over"].append(is_over)
    
    return pd.DataFrame(rows)


def train_league_model_v2(league_code: str, league_config: dict) -> dict | None:
    """Train an advanced goal-number model for a league."""
    print(f"\n{'='*70}")
    print(f"Training advanced model for {league_code} ({league_config['name']})")
    print(f"{'='*70}")
    
    history_path = HISTORICAL_DIR / f'{league_code}_history.csv'
    if not history_path.exists():
        print(f"  [SKIP] History file not found: {history_path}")
        return None
    
    df = pd.read_csv(history_path)
    print(f"  Loaded {len(df)} historical matches")
    
    if len(df) < league_config['min_rows']:
        print(f"  [SKIP] Not enough data: {len(df)} < {league_config['min_rows']}")
        return None
    
    has_odds = any(col in df.columns for col in ['B365>2.5', 'Avg>2.5', 'PS>2.5'])
    print(f"  Has odds data: {has_odds}")
    
    print("  Building advanced features...")
    features_df = build_features_advanced(df, has_odds=has_odds)
    print(f"  Built {len(features_df)} feature rows with {len(features_df.columns)} features")
    
    target_col = 'is_over'
    feature_cols = [col for col in features_df.columns if col != target_col]
    
    # Remove features with too many NaN values
    nan_ratio = features_df[feature_cols].isna().mean()
    valid_features = [col for col, ratio in zip(feature_cols, nan_ratio) if ratio < 0.7]
    print(f"  Features after removing high-NaN: {len(valid_features)}/{len(feature_cols)}")
    
    X = features_df[valid_features].fillna(0)
    y = features_df[target_col].astype(int)
    
    print(f"  Over rate: {y.mean():.2%}")
    
    # Time-based split
    if 'Date' in df.columns:
        df['Date_dt'] = pd.to_datetime(df['Date'], errors='coerce')
        df_sorted = df.sort_values('Date_dt').reset_index(drop=True)
        split_idx = int(len(df_sorted) * 0.8)
        
        features_sorted = features_df.iloc[df_sorted.index]
        X = features_sorted[valid_features].fillna(0)
        y = features_sorted[target_col].astype(int)
        
        X_train, X_test = X.iloc[:split_idx], X.iloc[split_idx:]
        y_train, y_test = y.iloc[:split_idx], y.iloc[split_idx:]
    else:
        from sklearn.model_selection import train_test_split
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
    
    print(f"  Train: {len(X_train)}, Test: {len(X_test)}")
    
    # Define models with hyperparameters
    models = {}
    
    # RandomForest with different parameters
    for n_est in [200, 300, 500]:
        for max_d in [10, 15, 20]:
            for min_s in [2, 5, 10]:
                model_name = f'RF_n{n_est}_d{max_d}_s{min_s}'
                models[model_name] = Pipeline([
                    ('imputer', SimpleImputer(strategy='median')),
                    ('scaler', RobustScaler()),
                    ('clf', RandomForestClassifier(
                        n_estimators=n_est,
                        max_depth=max_d,
                        min_samples_leaf=min_s,
                        random_state=42,
                        n_jobs=-1,
                        class_weight='balanced'
                    ))
                ])
    
    # ExtraTrees
    models['ET_n300_d15_s5'] = Pipeline([
        ('imputer', SimpleImputer(strategy='median')),
        ('scaler', RobustScaler()),
        ('clf', ExtraTreesClassifier(
            n_estimators=300,
            max_depth=15,
            min_samples_leaf=5,
            random_state=42,
            n_jobs=-1,
            class_weight='balanced'
        ))
    ])
    
    # GradientBoosting
    models['HGB_n200'] = Pipeline([
        ('imputer', SimpleImputer(strategy='median')),
        ('scaler', RobustScaler()),
        ('clf', GradientBoostingClassifier(
            n_estimators=200,
            max_depth=8,
            learning_rate=0.05,
            random_state=42
        ))
    ])
    
    # AdaBoost
    models['ADA_n100'] = Pipeline([
        ('imputer', SimpleImputer(strategy='median')),
        ('scaler', RobustScaler()),
        ('clf', AdaBoostClassifier(
            n_estimators=100,
            random_state=42
        ))
    ])
    
    print(f"\n  Testing {len(models)} models...")
    
    best_model = None
    best_config = None
    best_test_acc = 0.0
    
    for model_name, pipeline in models.items():
        try:
            pipeline.fit(X_train, y_train)
            
            y_pred = pipeline.predict(X_test)
            overall_acc = accuracy_score(y_test, y_pred)
            
            probs = pipeline.predict_proba(X_test)[:, 1]
            thresholds = np.arange(0.45, 0.86, 0.01)
            
            for threshold in thresholds:
                mask = (probs >= threshold) | (probs <= (1.0 - threshold))
                bet_count = int(mask.sum())
                
                if bet_count < 25:
                    continue
                
                preds = (probs >= threshold).astype(int)
                bet_acc = accuracy_score(y_test[mask], preds[mask])
                
                if bet_acc > best_test_acc or (bet_acc == best_test_acc and bet_count > best_config.get('bet_count', 0)):
                    best_test_acc = bet_acc
                    best_model = pipeline
                    best_config = {
                        'league': league_code,
                        'league_name': league_config['name'],
                        'model': model_name,
                        'threshold': float(threshold),
                        'bet_count': bet_count,
                        'bet_acc': float(bet_acc),
                        'overall_acc': float(overall_acc),
                        'f1_score': float(f1_score(y_test[mask], preds[mask])),
                        'rows': int(len(features_df)),
                        'train_rows': int(len(X_train)),
                        'test_rows': int(len(X_test)),
                        'features': valid_features,
                        'feature_count': len(valid_features),
                        'has_odds': has_odds,
                    }
        except Exception as e:
            print(f"    {model_name}: Error - {e}")
            continue
    
    if best_model is None or best_config is None:
        print(f"  [ERROR] No model met minimum bet count requirement (25)")
        return None
    
    print(f"\n  {'='*60}")
    print(f"  BEST MODEL")
    print(f"  {'='*60}")
    print(f"    Model: {best_config['model']}")
    print(f"    Threshold: {best_config['threshold']:.2f}")
    print(f"    Bet count: {best_config['bet_count']}")
    print(f"    Bet accuracy: {best_config['bet_acc']:.2%}")
    print(f"    F1 score: {best_config['f1_score']:.2%}")
    print(f"    Overall accuracy: {best_config['overall_acc']:.2%}")
    print(f"    Features: {best_config['feature_count']}")
    print(f"    Has odds: {best_config['has_odds']}")
    print(f"  {'='*60}")
    
    if best_config['bet_acc'] < 0.80:
        print(f"  [WARNING] Accuracy {best_config['bet_acc']:.2%} is below 80% target")
    else:
        print(f"  [SUCCESS] Accuracy {best_config['bet_acc']:.2%} meets 80%+ target!")
    
    # Save model
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    model_dir = OUTPUT_DIR / league_code
    model_dir.mkdir(parents=True, exist_ok=True)
    
    model_path = model_dir / 'goal_number_pipeline.pkl'
    config_path = model_dir / 'goal_number_config.json'
    
    joblib.dump(best_model, model_path)
    
    config = {
        **best_config,
        'task': 'goal-number',
        'target': 'over_under_2_5',
        'model_path': str(model_path),
    }
    
    with open(config_path, 'w', encoding='utf-8') as f:
        json.dump(config, f, ensure_ascii=False, indent=2)
    
    print(f"  Model saved to: {model_path}")
    print(f"  Config saved to: {config_path}")
    
    return best_config


def main() -> None:
    print("="*70)
    print("Training advanced goal-number models (target: 80%+ accuracy)")
    print("="*70)
    
    results = []
    
    for league_code, league_config in LEAGUES.items():
        result = train_league_model_v2(league_code, league_config)
        if result:
            results.append(result)
    
    print(f"\n{'='*70}")
    print("SUMMARY")
    print(f"{'='*70}")
    
    if results:
        print(f"\nTrained {len(results)} league models:\n")
        
        results_sorted = sorted(results, key=lambda x: x['bet_acc'], reverse=True)
        
        for idx, result in enumerate(results_sorted, 1):
            status = "✓" if result['bet_acc'] >= 0.80 else "✗"
            print(f"  {idx}. {result['league']} ({result['league_name']}): {status}")
            print(f"     Model: {result['model']}")
            print(f"     Threshold: {result['threshold']:.2f}")
            print(f"     Bet count: {result['bet_count']}")
            print(f"     Bet accuracy: {result['bet_acc']:.2%}")
            print(f"     F1 score: {result['f1_score']:.2%}")
            print(f"     Features: {result['feature_count']}")
            print(f"     Has odds: {result['has_odds']}")
            print()
        
        summary_path = OUTPUT_DIR / 'all_leagues_summary.json'
        with open(summary_path, 'w', encoding='utf-8') as f:
            json.dump(results_sorted, f, ensure_ascii=False, indent=2)
        print(f"Summary saved to: {summary_path}")
        
        high_accuracy_count = sum(1 for r in results if r['bet_acc'] >= 0.80)
        print(f"\nModels meeting 80%+ target: {high_accuracy_count}/{len(results)}")
        
        best_result = max(results, key=lambda x: x['bet_acc'])
        print(f"\nBest model: {best_result['league']} ({best_result['league_name']}) - {best_result['bet_acc']:.2%}")
        
        if best_result['bet_acc'] < 0.80:
            print(f"[WARNING] No model met 80%+ accuracy target. Best is {best_result['bet_acc']:.2%}")
        else:
            print(f"[SUCCESS] Target met! Best accuracy: {best_result['bet_acc']:.2%}")
    else:
        print("\nNo models trained!")


if __name__ == '__main__':
    main()
