from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
import joblib
import numpy as np

# Load the trained model and scaler
try:
    model = joblib.load('LR1.pkl')  # Ensure model path is correct
    scaler = joblib.load('scaler1.pkl')  # Ensure scaler is saved properly
except Exception as e:
    raise RuntimeError(f"Failed to load model or scaler: {e}")

app = FastAPI()

# Root endpoint
@app.get("/")
def read_root():
    return {"message": "Welcome to the PCOS Prediction API!"}

# Define the input data structure
class PCOSInput(BaseModel):
    Follicle_No_R: int = Field(..., ge=0, description="Number of follicles in the right ovary (>=0)")
    Follicle_No_L: int = Field(..., ge=0, description="Number of follicles in the left ovary (>=0)")
    Skin_darkening: int = Field(..., ge=0, le=1, description="Skin darkening (1: Yes, 0: No)")
    hair_growth: int = Field(..., ge=0, le=1, description="Hair growth (1: Yes, 0: No)")
    Weight_gain: int = Field(..., ge=0, le=1, description="Weight gain (1: Yes, 0: No)")
    Cycle_length: int = Field(..., ge=1, description="Menstrual cycle length in days (>=1)")
    AMH: float = Field(..., gt=0, description="Anti-Müllerian Hormone level (positive float)")
    Fast_food: int = Field(..., ge=0, le=1, description="Fast food consumption (1: Yes, 0: No)")
    Cycle_R_I: int = Field(..., ge=0, le=1, description="Regular/Irregular cycle (1: Regular, 0: Irregular)")
    FSH_LH: float = Field(..., gt=0, description="FSH/LH ratio (positive float)")
    PRL: float = Field(..., gt=0, description="Prolactin level (positive float)")
    Pimples: int = Field(..., ge=0, le=1, description="Pimples (1: Yes, 0: No)")
    Age: int = Field(..., ge=12, le=60, description="Age in years (between 12 and 60)")
    BMI: float = Field(..., gt=0, description="Body Mass Index (positive float)")

# Define the prediction response structure
class PCOSPrediction(BaseModel):
    prediction: int
    probability: float

# Define the prediction endpoint
@app.post("/predict/", response_model=PCOSPrediction)
def predict(data: PCOSInput):
    try:
        # Convert input data to array
        input_data = np.array([[data.Follicle_No_R, data.Follicle_No_L, data.Skin_darkening, data.hair_growth,
                                data.Weight_gain, data.Cycle_length, data.AMH, data.Fast_food, data.Cycle_R_I,
                                data.FSH_LH, data.PRL, data.Pimples, data.Age, data.BMI]])
        
        # Ensure scaler and model work with the same feature shape
        if input_data.shape[1] != scaler.n_features_in_:
            raise ValueError(f"Input data shape mismatch. Expected {scaler.n_features_in_}, got {input_data.shape[1]}")

        # Scale the input data using the loaded scaler
        scaled_input = scaler.transform(input_data)

        # Make prediction
        prediction = model.predict(scaled_input)
        probability = model.predict_proba(scaled_input)[0][1]  # Probability for class 1 (PCOS)

        return PCOSPrediction(prediction=int(prediction[0]), probability=probability)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction error: {e}")