from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import numpy as np
import joblib
from tensorflow.keras.models import load_model
from typing import List

# Initialize app
app = FastAPI(title="🔥 Fire Source Classification API")

# Load model and preprocessing tools
model = load_model("fire_detection_lstm_model.keras")
scaler = joblib.load("scaler.pkl")
label_encoder = joblib.load("label_encoder.pkl")

# Request schema
class SensorInput(BaseModel):
    window: List[List[float]]  # List of 30 time steps, each with 5 sensor values

@app.get("/")
def root():
    return {"message": "Fire source classification model is ready!"}

@app.post("/predict")
def predict_fire_type(data: SensorInput):
    sensor_window = np.array(data.window)

    # Validate shape
    if sensor_window.shape != (30, 5):
        raise HTTPException(status_code=400, detail="Input must be a 30x5 list of floats (30 time steps, 5 features).")

    # Scale and reshape
    try:
        sensor_window_scaled = scaler.transform(sensor_window)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Scaling failed: {str(e)}")

    sensor_input = sensor_window_scaled.reshape(1, 30, 5)

    # Predict
    try:
        probs = model.predict(sensor_input)
        predicted_index = int(np.argmax(probs, axis=1)[0])
        predicted_class = label_encoder.inverse_transform([predicted_index])[0]
        confidence = float(np.max(probs))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")

    return {
        "predicted_class": predicted_class,
        "confidence": round(confidence, 4)
    }