import numpy as np
import streamlit as st
import torch
from PIL import Image
from torch import optim
import config
from model import EmotionModel
from utils import load_checkpoint
import os


def remove_all_from_directory(directory):
    for item in os.listdir(directory):
        path = os.path.join(directory, item)
        if os.path.isfile(path):
            os.remove(path)


def process_image(image_path):
    model = EmotionModel().to(config.DEVICE)
    opt = optim.Adam(model.parameters(), lr=config.LEARNING_RATE, betas=(0.5, 0.999), )
    if config.LOAD_MODEL:
        load_checkpoint(
            config.CHECKPOINT, model, opt, config.LEARNING_RATE,
        )
    model.eval()
    image = np.array(Image.open(image_path).convert('L'))
    image = config.transform(image=image)["image"]
    image = image.to(config.DEVICE)
    image = torch.unsqueeze(image, dim=0)
    y_pred = model(image)
    label = {0: 'anger', 1: 'disgust', 2: 'fear', 3: 'happy', 4: 'neutral', 5: 'sad', 6: 'surprise'}
    return label[torch.argmax(y_pred).item()]


def main():
    st.title("Face emotion")
    uploaded_file = st.file_uploader("Upload an image file", type=["jpg", "png", "jpeg"])
    if uploaded_file:
        image_path = f"uploads/{uploaded_file.name}"
        remove_all_from_directory("uploads")
        with open(image_path, "wb") as f:
            f.write(uploaded_file.getbuffer())
        st.success(f"Image file uploaded successfully: {uploaded_file.name}")
        st.image(image_path)
        result_string = process_image(image_path)
        st.write(f"**{result_string.upper()}**", unsafe_allow_html=True)


if __name__ == "__main__":
    main()