import pandas as pd
from sklearn.linear_model import LinearRegression
import gradio as gr

# Load dataset
teams = pd.read_csv("teams.csv")
teams = teams[['team', 'country', 'year', 'athletes', 'age', 'prev_medals', 'medals']]
teams = teams.dropna()

# Split data into training and testing sets
train = teams[teams['year'] < 2012].copy()
test = teams[teams['year'] >= 2012].copy()

# Define predictors and target
predictors = ['athletes', 'prev_medals']
target = 'medals'

# Train the Linear Regression model
reg = LinearRegression()
reg.fit(train[predictors], train['medals'])

# Define the prediction function
def predict_medals(athletes: int, prev_medals: int):
    input_data = pd.DataFrame({'athletes': [athletes], 'prev_medals': [prev_medals]})
    prediction = reg.predict(input_data)[0]
    return max(0, round(prediction))

# Create Gradio interface
interface = gr.Interface(
    fn=predict_medals,
    inputs=[
        gr.Number(label="Number of Athletes"),
        gr.Number(label="Previous Medals Won"),
    ],
    outputs="number",
    title="Olympics Medal Prediction",
    description="Predict the number of medals a team might win based on athletes and previous medals."
)

# Launch the interface
if __name__ == "__main__":
    interface.launch()