from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import numpy as np
import onnxruntime as ort

app = FastAPI(
    title="Spandan ONNX API",
    version="1.0.0",
    description="REST API for Stress, Physical Activity, Focus and Cognitive Engagement prediction using ONNX models."
)

# ==========================================================
# Load Models
# ==========================================================

wearable_model = ort.InferenceSession("full_dataset_model_filewearable_model.onnx")
distraction_model = ort.InferenceSession("distraction_model.onnx")

wearable_input = wearable_model.get_inputs()[0].name
distraction_input = distraction_model.get_inputs()[0].name

CLASS_NAMES = [
    "Stress",
    "Physical Activity",
    "Normal"
]

# ==========================================================
# Request Schemas
# ==========================================================

class WearableInput(BaseModel):
    sensor_values: list[float]


class DistractionInput(BaseModel):
    sensor_values: list[float]


# ==========================================================
# Health Check
# ==========================================================

@app.get("/")
def root():
    return {
        "message": "Spandan API Running",
        "status": "healthy",
        "version": "1.0.0"
    }


# ==========================================================
# Model 1 - Stress / Physical Activity / Focus
# ==========================================================

@app.post("/predict/state")
def predict_state(data: WearableInput):

    if len(data.sensor_values) != 13:
        raise HTTPException(
            status_code=400,
            detail=f"Expected exactly 13 sensor values, got {len(data.sensor_values)}"
        )

    x = np.array([data.sensor_values], dtype=np.float32)

    try:

        logits = wearable_model.run(
            None,
            {
                wearable_input: x
            }
        )[0][0]

        prediction = int(np.argmax(logits))

        # Stable Softmax
        exp = np.exp(logits - np.max(logits))
        probabilities = exp / np.sum(exp)

        return {
            "prediction": CLASS_NAMES[prediction],
            "confidence": round(float(probabilities[prediction]), 4),
            "scores": {
                CLASS_NAMES[i]: round(float(probabilities[i]), 4)
                for i in range(len(CLASS_NAMES))
            }
        }

    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Model inference failed: {str(e)}"
        )


# ==========================================================
# Model 2 - Cognitive Engagement
# ==========================================================

@app.post("/predict/engagement")
def predict_engagement(data: DistractionInput):

    if len(data.sensor_values) != 7:
        raise HTTPException(
            status_code=400,
            detail=f"Expected exactly 7 sensor values, got {len(data.sensor_values)}"
        )

    x = np.array([data.sensor_values], dtype=np.float32)

    try:

        prediction = int(
            distraction_model.run(
                None,
                {
                    distraction_input: x
                }
            )[0][0]
        )

        status = "Neutral" if prediction == 1 else "Distracted"

        return {
            "prediction": prediction,
            "status": status
        }

    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Model inference failed: {str(e)}"
        )