from fastapi import FastAPI
from pydantic import BaseModel
import joblib
import numpy as np
import pandas as pd

app = FastAPI(title="EPL Predictions")

# Load model
try:
    model = joblib.load('simple_rf_model.joblib')
    scaler = joblib.load('simple_scaler.joblib')
    print("Models loaded successfully")
except:
    model = None
    scaler = None
    print("Failed to load models")

class MatchRequest(BaseModel):
    home_team: str
    away_team: str
    home_xg: float = 1.5
    away_xg: float = 1.3

class PredictionResponse(BaseModel):
    home_team: str
    away_team: str
    home_win: float
    draw: float
    away_win: float
    prediction: str
    confidence: float

@app.get("/")
def root():
    return {"status": "EPL Prediction API", "model_loaded": model is not None}

@app.get("/health")
def health():
    return {"status": "healthy", "models_loaded": model is not None}

@app.post("/predict")
def predict(match: MatchRequest):
    if model is None or scaler is None:
        return {"error": "Models not loaded"}

    # Prepare features
    features = pd.DataFrame([{
        'home_xg': match.home_xg * 0.82,  # Apply calibration
        'away_xg': match.away_xg * 0.82
    }])

    # Scale and predict
    X_scaled = scaler.transform(features)
    probs = model.predict_proba(X_scaled)[0]

    # Map probabilities (0=away, 1=draw, 2=home)
    away_prob = probs[0]
    draw_prob = probs[1] if len(probs) > 2 else 0.25
    home_prob = probs[2] if len(probs) > 2 else probs[1]

    # Get prediction
    if home_prob > draw_prob and home_prob > away_prob:
        prediction = "Home"
        confidence = home_prob
    elif away_prob > draw_prob:
        prediction = "Away"
        confidence = away_prob
    else:
        prediction = "Draw"
        confidence = draw_prob

    return PredictionResponse(
        home_team=match.home_team,
        away_team=match.away_team,
        home_win=float(home_prob),
        draw=float(draw_prob),
        away_win=float(away_prob),
        prediction=prediction,
        confidence=float(confidence)
    )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)