import os
from tensorflow.keras.models import load_model
from tensorflow.keras.utils import img_to_array
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

import tensorflow as tf

from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
# from tensorflow.keras.preprocessing import image
import numpy as np

import uvicorn
from PIL import Image
import io
import logging




# Suppress TensorFlow warnings
#tf.get_logger().setLevel(logging.ERROR)

app = FastAPI()

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Adjust as needed
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Replace with your actual model path
model_path = 'face_expression_detection_model3.h5'  # Update with your new model's path
model = load_model(model_path)


# def preprocess_image(img: Image.Image, target_size=(224, 224)):
#     img = img.resize(target_size)
#     x = image.img_to_array(img)
#     x = np.expand_dims(x, axis=0)
#     x = x / 255.0  # Assuming normalization was done during training
#     return x

def preprocess_image(img: Image.Image, target_size=(224, 224)):
    img = img.resize(target_size)
    x = img_to_array(img)
    x = np.expand_dims(x, axis=0)
    x = x / 255.0  # Assuming normalization was done during training
    return x

@app.get("/")
async def read_root():
    return {"message": "Welcome to the Emotion Classification API"}


@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    contents = await file.read()
    img = Image.open(io.BytesIO(contents))
    x = preprocess_image(img)
    prediction = model.predict(x)
    predicted_class = np.argmax(prediction[0])

    class_names = ['angry', 'happy', 'sad', 'surprised', 'neutral', 'disgusted', 'fearful']
    stress_levels = {
        'angry': 'high_stress',
        'happy': 'neutral',
        'sad': 'low_stress',
        'surprised': 'low_stress',
        'neutral': 'neutral',
        'disgusted': 'low_stress',
        'fearful': 'low_stress'
    }

    emotion = class_names[predicted_class]
    stress_level = stress_levels[emotion]

    result = {
        #"predicted_class": emotion,
        "predicted_class": stress_level
    }
    return JSONResponse(content=result)