train_split.py

import pandas as pd
import numpy as np
import xgboost as xgb
from scipy.sparse import coo_matrix
from scipy.sparse.linalg import lsqr
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error
import warnings
import os

warnings.simplefilter(action='ignore', category=FutureWarning)

# --- CONFIGURATION ---
DATA_FILE = 'ftc_decode_2025_data_threaded.csv'
STATS_FILE = 'team_stats_split.csv'


# --- 1. FEATURE ENGINEERING ---
def calculate_opr_and_stats(df):
    print("Stats: Calculating Component OPRs...")
    all_teams = pd.unique(df[['red_1', 'red_2', 'blue_1', 'blue_2']].values.ravel('K'))
    all_teams.sort()
    team_to_idx = {team: i for i, team in enumerate(all_teams)}
    n_matches = len(df) * 2

    # Participation Matrix
    red_rows = np.arange(len(df))
    blue_rows = np.arange(len(df)) + len(df)
    cols = np.concatenate([
        df['red_1'].map(team_to_idx).values, df['red_2'].map(team_to_idx).values,
        df['blue_1'].map(team_to_idx).values, df['blue_2'].map(team_to_idx).values
    ])
    rows = np.concatenate([red_rows, red_rows, blue_rows, blue_rows])
    A = coo_matrix((np.ones(len(rows)), (rows, cols)), shape=(n_matches, len(all_teams)))

    team_stats = pd.DataFrame({'team': all_teams})

    # Calculate OPRs for detailed metrics
    # We focus heavily on the specific scoring components now
    metrics = [
        'auto_score', 'dc_score',  # The big two
        'auto_artifacts', 'teleop_artifacts',
        'auto_patterns', 'teleop_patterns'
    ]

    for metric in metrics:
        if f'red_{metric}' in df.columns:
            b = np.concatenate([df[f'red_{metric}'].values, df[f'blue_{metric}'].values])
            team_stats[f'opr_{metric}'] = lsqr(A, b, damp=2.0)[0]  # Lower damp for more aggressive fitting

    return team_stats


def prepare_data(df):
    stats = calculate_opr_and_stats(df)
    stats.to_csv(STATS_FILE, index=False)
    print(f"Stats saved to {STATS_FILE}")

    # Merge stats back
    def merge_slot(match_df, slot_col, prefix):
        s = stats.copy()
        s.columns = [f"{prefix}_{c}" if c != 'team' else slot_col for c in s.columns]
        return pd.merge(match_df, s, on=slot_col, how='left')

    train_df = df.copy()

    # Add Season Progress Feature (0.0 to 1.0)
    # This helps the model know that late-season matches have higher scores
    train_df['season_progress'] = train_df['match_id'] / train_df['match_id'].max()

    for slot in ['red_1', 'red_2', 'blue_1', 'blue_2']:
        train_df = merge_slot(train_df, slot, slot)

    return train_df


# --- 2. TRAINING ENGINE ---
def train_component_model(X, y, name):
    print(f"\nTraining {name} Model...")
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.15, random_state=42)

    model = xgb.XGBRegressor(
        objective='reg:squarederror',
        n_estimators=10000,
        learning_rate=0.003,  # Very precise
        max_depth=5,  # Moderate depth to prevent overfitting
        subsample=0.6,
        colsample_bytree=0.7,
        reg_alpha=2.0,  # L1 Regularization
        early_stopping_rounds=300,
        n_jobs=-1
    )

    model.fit(
        X_train, y_train,
        eval_set=[(X_test, y_test)],
        verbose=500
    )

    model.save_model(f'model_{name}.json')
    preds = model.predict(X_test)
    error = mean_absolute_error(y_test, preds)
    print(f"✅ {name.upper()} MAE: +/- {error:.2f} pts")
    return model


# --- MAIN ---
if not os.path.exists(DATA_FILE):
    print(f"Error: {DATA_FILE} not found.");
    exit()

df = pd.read_csv(DATA_FILE)
train_df = prepare_data(df)

# Feature Selection
# We use only numeric OPR columns + season_progress
X = train_df.select_dtypes(include=[np.number])
cols_to_keep = [c for c in X.columns if 'opr_' in c or c == 'season_progress']
X = X[cols_to_keep]

# 3. SPLIT TRAINING
# Model A: Predicts ONLY Auto Score
y_auto = train_df['red_auto_score']
train_component_model(X, y_auto, "auto")

# Model B: Predicts ONLY Teleop (DC) Score
y_teleop = train_df['red_dc_score']
train_component_model(X, y_teleop, "teleop")

print("\n🚀 DONE! Models saved as 'model_auto.json' and 'model_teleop.json'")
print("Use predict_split.py to combine them.")