#!/usr/bin/env python3
"""Fast training with aggressive feature engineering for 80%+ accuracy."""

from __future__ import annotations

import json
import sys
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
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.impute import SimpleImputer

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

LEAGUES = {
    'a1': {'name': '澳超'},
    'br1': {'name': '巴甲'},
    'd2': {'name': '德乙'},
    'e0': {'name': '英超'},
    'f1': {'name': '法甲'},
    'f2': {'name': '法乙'},
    'g1': {'name': '希腊超'},
    'j1': {'name': '日职联'},
    'm1': {'name': '美职联'},
    'r1': {'name': '俄超'},
    'uz1': {'name': '乌兹别克超'},
    'sp1': {'name': '西甲'},
    'e1': {'name': '英甲'},
    'e2': {'name': '英冠'},
}


def build_features_fast(df: pd.DataFrame) -> pd.DataFrame:
    """Build features with comprehensive engineering."""
    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": [],
        "gf_last20": [], "ga_last20": [], "over_last20": [],
        "momentum": [],
        "consec_over": [],  # Consecutive overs
        "consec_under": [],  # Consecutive unders
    })
    
    rows = []
    
    for _, 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
        
        # Comprehensive features
        row = {}
        
        # Window-based averages
        for window in [3, 5, 10, 20]:
            row[f"home_gf_last{window}"] = np.mean(home_hist[f"gf_last{window}"]) if home_hist[f"gf_last{window}"] else 1.4
            row[f"home_ga_last{window}"] = np.mean(home_hist[f"ga_last{window}"]) if home_hist[f"ga_last{window}"] else 1.4
            row[f"away_gf_last{window}"] = np.mean(away_hist[f"gf_last{window}"]) if away_hist[f"gf_last{window}"] else 1.2
            row[f"away_ga_last{window}"] = np.mean(away_hist[f"ga_last{window}"]) if away_hist[f"ga_last{window}"] else 1.2
            row[f"home_over_last{window}"] = np.mean(home_hist[f"over_last{window}"]) if home_hist[f"over_last{window}"] else 0.5
            row[f"away_over_last{window}"] = np.mean(away_hist[f"over_last{window}"]) if away_hist[f"over_last{window}"] else 0.5
        
        # Overall averages
        row["home_gf_avg"] = np.mean(home_hist["gf"]) if home_hist["gf"] else 1.4
        row["home_ga_avg"] = np.mean(home_hist["ga"]) if home_hist["ga"] else 1.4
        row["away_gf_avg"] = np.mean(away_hist["gf"]) if away_hist["gf"] else 1.2
        row["away_ga_avg"] = np.mean(away_hist["ga"]) if away_hist["ga"] else 1.2
        row["home_tot_avg"] = np.mean(home_hist["tot"]) if home_hist["tot"] else 2.6
        row["away_tot_avg"] = np.mean(away_hist["tot"]) if away_hist["tot"] else 2.6
        row["home_over_rate"] = np.mean(home_hist["over"]) if home_hist["over"] else 0.5
        row["away_over_rate"] = np.mean(away_hist["over"]) if away_hist["over"] else 0.5
        
        # Home/Away specific
        row["home_home_gf_avg"] = np.mean(home_hist["home_gf"]) if home_hist["home_gf"] else np.mean(home_hist["gf"]) if home_hist["gf"] else 1.4
        row["home_home_ga_avg"] = np.mean(home_hist["home_ga"]) if home_hist["home_ga"] else np.mean(home_hist["ga"]) if home_hist["ga"] else 1.4
        row["away_away_gf_avg"] = np.mean(away_hist["away_gf"]) if away_hist["away_gf"] else np.mean(away_hist["gf"]) if away_hist["gf"] else 1.2
        row["away_away_ga_avg"] = np.mean(away_hist["away_ga"]) if away_hist["away_ga"] else np.mean(away_hist["ga"]) if away_hist["ga"] else 1.2
        row["home_home_over_rate"] = np.mean(home_hist["home_over"]) if home_hist["home_over"] else np.mean(home_hist["over"]) if home_hist["over"] else 0.5
        row["away_away_over_rate"] = np.mean(away_hist["away_over"]) if away_hist["away_over"] else np.mean(away_hist["over"]) if away_hist["over"] else 0.5
        
        # Momentum
        row["home_momentum"] = np.mean(home_hist["momentum"][-10:]) if len(home_hist["momentum"]) >= 10 else 0.5
        row["away_momentum"] = np.mean(away_hist["momentum"][-10:]) if len(away_hist["momentum"]) >= 10 else 0.5
        
        # Variance (consistency)
        row["home_gf_var"] = np.var(home_hist["gf_last20"]) if len(home_hist["gf_last20"]) >= 2 else 0.5
        row["home_ga_var"] = np.var(home_hist["ga_last20"]) if len(home_hist["ga_last20"]) >= 2 else 0.5
        row["away_gf_var"] = np.var(away_hist["gf_last20"]) if len(away_hist["gf_last20"]) >= 2 else 0.5
        row["away_ga_var"] = np.var(away_hist["ga_last20"]) if len(away_hist["ga_last20"]) >= 2 else 0.5
        
        # Expected goals
        row["expected_goals_home"] = row["home_gf_avg"] + row["away_ga_avg"]
        row["expected_goals_away"] = row["away_gf_avg"] + row["home_ga_avg"]
        row["expected_goals_proxy"] = (row["expected_goals_home"] + row["expected_goals_away"]) / 2
        row["combined_over_rate"] = (row["home_over_rate"] + row["away_over_rate"]) / 2
        
        # Difference features
        row["home_away_gf_diff"] = row["home_gf_avg"] - row["away_gf_avg"]
        row["home_away_ga_diff"] = row["home_ga_avg"] - row["away_ga_avg"]
        row["home_away_tot_diff"] = row["home_tot_avg"] - row["away_tot_avg"]
        row["home_away_over_diff"] = row["home_over_rate"] - row["away_over_rate"]
        
        # Form ratio
        row["form_ratio"] = row["home_momentum"] / (row["away_momentum"] + 0.01)
        
        # Trend features
        home_trend = row["home_gf_last10"] - row["home_gf_last20"] if len(home_hist["gf_last20"]) >= 10 and len(home_hist["gf_last10"]) >= 10 else 0
        away_trend = row["away_gf_last10"] - row["away_gf_last20"] if len(away_hist["gf_last20"]) >= 10 and len(away_hist["gf_last10"]) >= 10 else 0
        row["home_trend"] = home_trend
        row["away_trend"] = away_trend
        
        # Consistency ratio (variance / mean)
        row["home_consistency"] = row["home_gf_var"] / (row["home_gf_avg"] + 0.01)
        row["away_consistency"] = row["away_gf_var"] / (row["away_gf_avg"] + 0.01)
        
        # Odds features (if available)
        for odds_prefix in ['B365', 'Avg', 'PS', 'Max']:
            over_col = f'{odds_prefix}>2.5'
            under_col = f'{odds_prefix}<2.5'
            
            if over_col in df.columns and under_col in df.columns:
                over_odds = float(match.get(over_col, np.nan))
                under_odds = float(match.get(under_col, np.nan))
                
                if pd.notna(over_odds) and pd.notna(under_odds):
                    row[f"{odds_prefix.lower()}_over_odds"] = over_odds
                    row[f"{odds_prefix.lower()}_under_odds"] = under_odds
                    
                    over_prob = 1.0 / over_odds
                    under_prob = 1.0 / 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)
                    
                    # Implied probability from odds
                    row["market_over_prob"] = fair_over
                    
                    # Odds ratio (over/under)
                    row[f"{odds_prefix.lower()}_over_under_ratio"] = over_odds / (under_odds + 0.01)
                    
                    # Use first available odds source
                    break
        
        # Target
        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)
        
        for window in [3, 5, 10, 20]:
            home_hist[f"gf_last{window}"].append(total_goals)
            home_hist[f"ga_last{window}"].append(total_goals)
            home_hist[f"over_last{window}"].append(is_over)
        
        home_result = 1 if total_goals > (match.get('FTAG', 0) or 0) else 0.5
        home_hist["momentum"].append(home_result)
        home_hist["consec_over"].append(is_over)
        home_hist["consec_under"].append(1 - is_over)
        
        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)
        
        for window in [3, 5, 10, 20]:
            away_hist[f"gf_last{window}"].append(total_goals)
            away_hist[f"ga_last{window}"].append(total_goals)
            away_hist[f"over_last{window}"].append(is_over)
        
        away_result = 1 if total_goals > (match.get('FTHG', 0) or 0) else 0.5
        away_hist["momentum"].append(away_result)
        away_hist["consec_over"].append(is_over)
        away_hist["consec_under"].append(1 - is_over)
        
        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_fast(league_code: str, league_name: str) -> dict | None:
    """Fast training with aggressive hyperparameters."""
    print(f"\n{'='*70}")
    print(f"Fast training for {league_code} ({league_name})")
    print(f"Target: 80%+ bet accuracy")
    print(f"{'='*70}")
    
    # Load data
    history_path = HISTORICAL_DIR / f'{league_code}_history.csv'
    if not history_path.exists():
        print(f"  [SKIP] File not found")
        return None
    
    df = pd.read_csv(history_path)
    print(f"  Loaded {len(df)} matches")
    
    if len(df) < 1000:
        print(f"  [SKIP] Not enough data")
        return None
    
    # Build features
    print("  Building features...")
    features_df = build_features_fast(df)
    print(f"  Built {len(features_df)} features")
    
    # Prepare data
    target_col = 'is_over'
    feature_cols = [c for c in features_df.columns if c != target_col]
    
    # Remove high-NaN features
    nan_ratio = features_df[feature_cols].isna().mean()
    valid_features = [c for c, ratio in zip(feature_cols, nan_ratio) if ratio < 0.5]
    print(f"  Features after filtering: {len(valid_features)}")
    
    X = features_df[valid_features].fillna(0)
    y = features_df[target_col].astype(int)
    
    # 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)
        
        X_sorted = X.iloc[df_sorted.index].reset_index(drop=True)
        y_sorted = y.iloc[df_sorted.index].reset_index(drop=True)
        
        X_train, X_test = X_sorted.iloc[:split_idx], X_sorted.iloc[split_idx:]
        y_train, y_test = y_sorted.iloc[:split_idx], y_sorted.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)}")
    print(f"  Over rate: {y.mean():.2%}")
    
    # Test different models with aggressive hyperparameters
    best_model = None
    best_config = None
    best_test_acc = 0.0
    
    models_to_test = [
        # RandomForest - Very aggressive
        ('RF_aggressive', Pipeline([
            ('imputer', SimpleImputer(strategy='median')),
            ('scaler', StandardScaler()),
            ('clf', RandomForestClassifier(
                n_estimators=500,
                max_depth=25,
                min_samples_leaf=1,
                max_features='sqrt',
                random_state=42,
                n_jobs=-1,
                class_weight='balanced_subsample'
            ))
        ])),
        
        # RandomForest - Moderate
        ('RF_moderate', Pipeline([
            ('imputer', SimpleImputer(strategy='median')),
            ('scaler', StandardScaler()),
            ('clf', RandomForestClassifier(
                n_estimators=300,
                max_depth=20,
                min_samples_leaf=2,
                max_features=0.7,
                random_state=42,
                n_jobs=-1,
                class_weight='balanced'
            ))
        ])),
        
        # GradientBoosting
        ('GB_aggressive', Pipeline([
            ('imputer', SimpleImputer(strategy='median')),
            ('scaler', StandardScaler()),
            ('clf', GradientBoostingClassifier(
                n_estimators=300,
                max_depth=10,
                learning_rate=0.1,
                subsample=0.9,
                random_state=42
            ))
        ])),
    ]
    
    print(f"\n  Testing {len(models_to_test)} models...")
    
    for model_name, pipeline in models_to_test:
        try:
            print(f"\n  {model_name}...")
            
            pipeline.fit(X_train, y_train)
            
            y_pred = pipeline.predict(X_test)
            overall_acc = accuracy_score(y_test, y_pred)
            
            # Test thresholds
            probs = pipeline.predict_proba(X_test)[:, 1]
            thresholds = np.arange(0.40, 0.91, 0.02)
            
            for threshold in thresholds:
                mask = (probs >= threshold) | (probs <= (1.0 - threshold))
                bet_count = int(mask.sum())
                
                if bet_count < 20:
                    continue
                
                preds = (probs >= threshold).astype(int)
                bet_acc = accuracy_score(y_test[mask], preds[mask])
                
                if bet_acc > best_test_acc:
                    best_test_acc = bet_acc
                    best_model = pipeline
                    best_config = {
                        'league': league_code,
                        'league_name': league_name,
                        'model': model_name,
                        'threshold': float(threshold),
                        'bet_count': bet_count,
                        'bet_acc': float(bet_acc),
                        'overall_acc': float(overall_acc),
                        '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),
                    }
                    
                    print(f"    New best! Threshold={threshold:.2f}, Bet Acc={bet_acc:.2%}, Bet Count={bet_count}")
        
        except Exception as e:
            print(f"  {model_name}: Error - {e}")
            continue
    
    if not best_model or best_config['bet_acc'] < 0.70:
        print(f"\n  [ERROR] No model met 70% minimum")
        return None
    
    print(f"\n  {'='*60}")
    print(f"  FINAL 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"    Overall accuracy: {best_config['overall_acc']:.2%}")
    print(f"    Features: {best_config['feature_count']}")
    print(f"  {'='*60}")
    
    if best_config['bet_acc'] >= 0.80:
        print(f"  [SUCCESS] Target met!")
    else:
        print(f"  [WARNING] Accuracy {best_config['bet_acc']:.2%} is below 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'
    
    # Re-train on full training data
    best_model.fit(X_train, y_train)
    
    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:
    if len(sys.argv) < 2:
        print("Usage: python3 train_fast.py <league_code>")
        print("Available:", ", ".join(LEAGUES.keys()))
        sys.exit(1)
    
    league_code = sys.argv[1].lower()
    
    print("="*70)
    print(f"Fast training: {league_code}")
    print("Target: 80%+ bet accuracy")
    print("="*70)
    
    league_config = LEAGUES.get(league_code)
    if not league_config:
        print(f"Unknown league: {league_code}")
        sys.exit(1)
    
    result = train_fast(league_code, league_config['name'])
    
    if result:
        print(f"\n{'='*70}")
        print("FINAL RESULT")
        print(f"{'='*70}")
        print(f"League: {result['league']} ({result['league_name']})")
        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"Features: {result['feature_count']}")
        
        status = "✓ SUCCESS" if result['bet_acc'] >= 0.80 else "✗ PARTIAL"
        print(f"\nStatus: {status}")


if __name__ == '__main__':
    main()
