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

# Load model from the local storage (ensure the model file is in the same directory)
model_path = "rf_model.pkl"
gb_model_loaded = joblib.load(model_path)

# Create FastAPI app
app = FastAPI()

# Define class labels
class_names = [
    'Emergency & Accident Unit', 'Heart Clinic',
    'Neuro Med Center', 'OPD:EYE', 'Dental',
    'OPD:MED', 'OPD:ENT', 'OPD:OBG',
    'OPD:Surgery + Uro.', 'Orthopedic Surgery',
    'GI Clinic', 'Breast Clinic', 'Skin & Dermatology'
]

# Define the input format for FastAPI using Pydantic BaseModel
class InputData(BaseModel):
    features: list[float]  # List of 32 feature inputs

@app.post("/predict")
def predict(data: InputData):
    try:
        # Validate input length
        if len(data.features) != 32:
            raise HTTPException(status_code=400, detail=f"Expected 32 features, but got {len(data.features)}")

        # Convert list to numpy array and reshape
        input_array = np.array(data.features).reshape(1, -1)

        # Get predictions
        prediction = gb_model_loaded.predict_proba(input_array)

        # Convert probabilities to percentage and format
        probabilities = (prediction[0] * 100).round(2)
        result_pro = {class_name: f"{prob:.2f}%" for class_name, prob in zip(class_names, probabilities)}

        # Return result as JSON
        return {'result': result_pro}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Run the application with the following command if needed
# if __name__ == "__main__":
#     uvicorn.run(app, host="0.0.0.0", port=8501)