import streamlit as st
import pandas as pd
from catboost import CatBoostClassifier

# Page configuration
st.set_page_config(page_title="F1 Pit Strategy Predictor", layout="wide")

@st.cache_resource
def load_model():
    model = CatBoostClassifier()
    model.load_model("f1_pit_strategy_model.cbm", format='cbm')
    return model

model = load_model()

st.title("🏎️ F1 Pit Stop Strategy Predictor")
st.markdown("""
Predict the probability of a driver pitting on the next lap using your high-performance CatBoost model.
**Developed by: [Basak Tamer](https://huggingface.co/basaktamer)**
""")

# Create columns for organized input
col1, col2, col3 = st.columns(3)

with col1:
    driver = st.text_input("Driver Code", value="HAM")
    race = st.text_input("Race Name", value="Monaco Grand Prix")
    year = st.number_input("Year", value=2026)
    compound = st.selectbox("Tyre Compound", ["SOFT", "MEDIUM", "HARD", "INTERMEDIATE", "WET"])
    pit_stop = st.selectbox("PitStop (Current Lap)", [0, 1])

with col2:
    lap_number = st.number_input("Lap Number", min_value=1, value=20)
    stint = st.number_input("Stint Number", min_value=1, value=1)
    tyre_life = st.number_input("Tyre Life (Laps)", min_value=0, value=10)
    position = st.slider("Current Position", 1, 20, 5)
    pos_change = st.number_input("Position Change", value=0)

with col3:
    lap_time = st.number_input("LapTime (s)", value=80.0)
    lap_delta = st.number_input("LapTime_Delta", value=0.1)
    cum_deg = st.number_input("Cumulative_Degradation", value=0.5)
    race_progress = st.slider("Race Progress (%)", 0.0, 1.0, 0.3)

# Prediction button
if st.button("Predict Pit Probability"):
    # The order MUST match: ['Driver', 'Compound', 'Race', 'Year', 'PitStop', 'LapNumber', 
    # 'Stint', 'TyreLife', 'Position', 'LapTime (s)', 'LapTime_Delta', 
    # 'Cumulative_Degradation', 'RaceProgress', 'Position_Change']
    
    input_data = pd.DataFrame({
        'Driver': [driver],
        'Compound': [compound],
        'Race': [race],
        'Year': [year],
        'PitStop': [pit_stop],
        'LapNumber': [lap_number],
        'Stint': [stint],
        'TyreLife': [tyre_life],
        'Position': [position],
        'LapTime (s)': [lap_time],
        'LapTime_Delta': [lap_delta],
        'Cumulative_Degradation': [cum_deg],
        'RaceProgress': [race_progress],
        'Position_Change': [pos_change]
    })
    
    # Re-enforcing the exact order identified by your model.feature_names_
    column_order = [
        'Driver', 'Compound', 'Race', 'Year', 'PitStop', 'LapNumber', 
        'Stint', 'TyreLife', 'Position', 'LapTime (s)', 'LapTime_Delta', 
        'Cumulative_Degradation', 'RaceProgress', 'Position_Change'
    ]
    input_data = input_data[column_order]
    
    try:
        prob = model.predict_proba(input_data)[0][1]
        
        st.divider()
        st.subheader(f"Pit Probability: {prob:.2%}")
        
        if prob > 0.5:
            st.error("Strategy Alert: High probability of pitting next lap!")
        else:
            st.success("Strategy Insight: Likely to stay out.")
            
    except Exception as e:
        st.error(f"Prediction Error: {e}")

st.sidebar.markdown("### Model Details")
st.sidebar.write("Target: `PitNextLap`")
st.sidebar.write("Validation AUC: 0.9468")
st.sidebar.write("Public Score: 0.9446")