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

# Load model from the local storage (ensure the model file is in the same directory)
model_path = "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))

# To run the FastAPI app locally for testing
# Uncomment the following lines
# if __name__ == "__main__":
#     import uvicorn
#     uvicorn.run(app, host="0.0.0.0", port=8501)