from fastapi import FastAPI
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from typing import Annotated
import pandas as pd
import joblib
import gradio as gr
from fastapi.middleware.cors import CORSMiddleware

# Load model
model_path = "xgb_model_reg.pkl"
model = joblib.load(model_path)

# FastAPI app
app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Change to your domain in production
    allow_methods=["GET", "POST"],
    allow_headers=["*"],
)

# Pydantic model for API validation
class UserInput(BaseModel):
    age: Annotated[int, Field(gt=0)]
    albumin_gL: Annotated[float, Field(gt=0)]
    creat_umol: Annotated[float, Field(gt=0)]
    glucose_mmol: Annotated[float, Field(gt=0)]
    lncrp: Annotated[float, Field(gt=0)]
    lymph: Annotated[float, Field(gt=0)]
    mcv: Annotated[float, Field(gt=0)]
    rdw: Annotated[float, Field(gt=0)]
    alp: Annotated[float, Field(gt=0)]
    wbc: Annotated[float, Field(gt=0)]

# FastAPI endpoint
@app.post('/predict')
def predict_api(data: UserInput):
    try:
        df = pd.DataFrame([data.dict()])
        pred = float(model.predict(df)[0])
        return JSONResponse(status_code=200, content={"Predicted Biological Age": pred})
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})

# Gradio prediction function
def predict_gradio(age, albumin_gL, creat_umol, glucose_mmol, lncrp, lymph, mcv, rdw, alp, wbc):
    df = pd.DataFrame([{
        "age": age,
        "albumin_gL": albumin_gL,
        "creat_umol": creat_umol,
        "glucose_mmol": glucose_mmol,
        "lncrp": lncrp,
        "lymph": lymph,
        "mcv": mcv,
        "rdw": rdw,
        "alp": alp,
        "wbc": wbc
    }])
    pred = float(model.predict(df)[0])
    return f"Predicted Biological Age: {pred:.2f} years"

# Gradio interface
gr_interface = gr.Interface(
    fn=predict_gradio,
    inputs=[
        gr.Number(label="Age", precision=0),
        gr.Number(label="Albumin (g/L)"),
        gr.Number(label="Creatinine (µmol/L)"),
        gr.Number(label="Glucose (mmol/L)"),
        gr.Number(label="ln(CRP)"),
        gr.Number(label="Lymph"),
        gr.Number(label="MCV"),
        gr.Number(label="RDW"),
        gr.Number(label="ALP"),
        gr.Number(label="WBC"),
    ],
    outputs="text",
    title="Biological Age Predictor",
    description="Enter patient lab values to predict biological age."
)

# Mount Gradio app on FastAPI
app = gr.mount_gradio_app(app, gr_interface, path="/gradio")

# To run locally:
# uvicorn app:app --reload