ftc-match-predictor/predict_split.py at main · CycleZLab/ftc-match-predictor · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import pandas as pd
import xgboost as xgb
import os

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


def load_models():
    if not os.path.exists('model_auto.json') or not os.path.exists('model_teleop.json'):
        print("Models not found. Run train_split.py first.")
        return None, None, None

    m_auto = xgb.XGBRegressor()
    m_auto.load_model('model_auto.json')

    m_teleop = xgb.XGBRegressor()
    m_teleop.load_model('model_teleop.json')

    stats = pd.read_csv(STATS_FILE).set_index('team')
    return m_auto, m_teleop, stats


def get_prediction(m_auto, m_teleop, stats, r1, r2, b1, b2):
    feature_names = m_auto.get_booster().feature_names
    row = {'season_progress': 1.0}  # Assume end-of-season form (peak performance)

    def fill(team, prefix):
        if team in stats.index:
            for k, v in stats.loc[team].items():
                row[f"{prefix}_{k}"] = v

    fill(r1, 'red_1');
    fill(r2, 'red_2')
    fill(b1, 'blue_1');
    fill(b2, 'blue_2')

    # 1. Predict RED
    df_red = pd.DataFrame([row], columns=feature_names).fillna(0)
    red_auto = m_auto.predict(df_red)[0]
    red_teleop = m_teleop.predict(df_red)[0]

    # 2. Predict BLUE (Flip inputs)
    row_blue = {'season_progress': 1.0}
    for col in feature_names:
        if 'red_1' in col:
            row_blue[col] = row.get(col.replace('red_1', 'blue_1'), 0)
        elif 'red_2' in col:
            row_blue[col] = row.get(col.replace('red_2', 'blue_2'), 0)
        elif 'blue_1' in col:
            row_blue[col] = row.get(col.replace('blue_1', 'red_1'), 0)
        elif 'blue_2' in col:
            row_blue[col] = row.get(col.replace('blue_2', 'red_2'), 0)

    df_blue = pd.DataFrame([row_blue], columns=feature_names).fillna(0)
    blue_auto = m_auto.predict(df_blue)[0]
    blue_teleop = m_teleop.predict(df_blue)[0]

    return (red_auto, red_teleop), (blue_auto, blue_teleop)


if __name__ == "__main__":
    m_auto, m_teleop, stats = load_models()
    if m_auto:
        print("\n🤖 DUAL-CORE PREDICTOR READY")
        while True:
            i = input("\nEnter [Red1 Red2 Blue1 Blue2] (or 'q'): ")
            if i.lower() == 'q': break
            try:
                teams = list(map(int, i.replace(',', ' ').split()))
                if len(teams) == 4:
                    (ra, rt), (ba, bt) = get_prediction(m_auto, m_teleop, stats, *teams)

                    # Add Estimated Penalty Bonus (Average ~15 pts) to match real world
                    # Since we trained on Clean Score, we need to add the average "noise" back in
                    # for the user to recognize the number.
                    penalty_bonus = 15.0

                    red_total = ra + rt + penalty_bonus
                    blue_total = ba + bt + penalty_bonus

                    winner = "RED" if red_total > blue_total else "BLUE"

                    print(f"\n🔮 PREDICTION:")
                    print(f"🔴 RED:  {red_total:.0f} (Auto: {ra:.0f}, Teleop: {rt:.0f})")
                    print(f"🔵 BLUE: {blue_total:.0f} (Auto: {ba:.0f}, Teleop: {bt:.0f})")
                    print(f"🏆 Winner: {winner} (Margin: {abs(red_total - blue_total):.1f})")
            except Exception as e:
                print(f"Error: {e}")