import tensorflow as tf
from fastapi import FastAPI
from pydantic import BaseModel
import numpy as np
import os

app = FastAPI()

# Load your specific model

MODEL_PATH = os.path.join(os.path.dirname(__file__), "driving_behavior_model.keras")

model = tf.keras.models.load_model(MODEL_PATH)

# Define the labels used in your notebook
CLASSES = ["Safe", "Moderate", "Dangerous"]

class ModelInput(BaseModel):
    # Expecting a list of 10 lists, each containing 6 sensor values
    # Example: [[ax, ay, az, gx, gy, gz], [...], ... (10 times)]
    data: list 

@app.post("/predict")
async def predict(input_params: ModelInput):
    # 1. Convert input to numpy array
    input_data = np.array(input_params.data) # Shape: (10, 6)
    
    # 2. Add the batch dimension to make it (1, 10, 6)
    input_data = np.expand_dims(input_data, axis=0)
    
    # 3. Run prediction
    prediction_probs = model.predict(input_data)
    
    # 4. Get the index of the highest probability
    predicted_class_index = np.argmax(prediction_probs[0])
    
    # 5. Return readable result
    return {
        "prediction_index": int(predicted_class_index),
        "label": CLASSES[predicted_class_index],
        "confidence": float(np.max(prediction_probs[0]))
    }