from fastapi import FastAPI
import pandas as pd
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler

app = FastAPI()

# ================= GLOBAL MODEL =================
model = None
scaler = None
features = None

# ================= TRAIN API =================
@app.post("/train")
def train_model(data: dict):

    global model, scaler, features

    df = pd.DataFrame(data["dataset"])

    target = data["target"]

    X = df.drop(columns=[target])
    y = df[target]

    features = list(X.columns)

    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    model = LogisticRegression(max_iter=5000)
    model.fit(X_scaled, y)

    return {"status": "Model trained successfully"}


# ================= PREDICT API =================
@app.post("/predict")
def predict(input_data: dict):

    global model, scaler, features

    if model is None:
        return {"error": "Model not trained yet"}

    df = pd.DataFrame([input_data])

    # Ensure correct column order
    df = df[features]

    X_scaled = scaler.transform(df)

    pred = model.predict(X_scaled)[0]
    prob = float(np.max(model.predict_proba(X_scaled)))

    return {
        "prediction": int(pred),
        "confidence": prob
    }