import gradio as gr
import joblib
import numpy as np
from fastapi import FastAPI, Request

# Load model
model = joblib.load("linear_model.pkl")

# ===== Inference function =====
def predict(x):
    X = np.array([[float(x)]])
    y_pred = model.predict(X)
    return float(y_pred[0])

# ===== Gradio UI =====
demo = gr.Interface(
    fn=predict,
    inputs=gr.Number(label="X value"),
    outputs=gr.Number(label="Predicted Y"),
    title="Simple Linear Regression",
    description="A tiny linear regression model trained with scikit-learn."
)

# ===== Add FastAPI endpoint for clean API calls =====
app = FastAPI()

@app.post("/predict")
async def predict_api(request: Request):
    body = await request.json()
    x_value = body.get("X")
    if x_value is None:
        return {"error": "Missing 'X' value"}
    return {"y_pred": predict(x_value)}

# Mount Gradio app inside FastAPI
app = gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    demo.launch()