from fastapi import FastAPI
from pydantic import BaseModel
import pandas as pd
import pickle
from fastapi.middleware.cors import CORSMiddleware


# Load the saved Random Forest model
with open("random_forest_pkl.pkl", "rb") as f:
    model = pickle.load(f)

# Initialize FastAPI app
app = FastAPI(
    title="Soil Fertility Prediction API",
    description="Predict soil fertility level (0=Low, 1=Medium, 2=High) using a trained Random Forest model.",
    version="1.0.0"
)

# Enable CORS (for browser and external app access)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # or replace * with your website URL for security
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Define input data model
class SoilInput(BaseModel):
    N: float
    P: float
    K: float
    pH: float
    EC: float
    OC: float
    S: float
    Zn: float
    Fe: float
    Cu: float
    Mn: float
    B: float

# Root endpoint
@app.get("/")
def root():
    return {"message": "Welcome to the Soil Fertility Prediction API"}

# Prediction endpoint
@app.post("/predict")
def predict_fertility(data: SoilInput):
    df = pd.DataFrame([data.model_dump()])
    pred = model.predict(df)[0]
    labels = {0: "Low Fertility", 1: "Medium Fertility", 2: "High Fertility"}
    return {
        "prediction": int(pred),
        "class_label": labels[int(pred)]
    }