import streamlit as st
import random 
import os
import torch
import numpy as np
from PIL import Image
from torchvision.transforms import v2
from torchvision.models import resnet18, ResNet18_Weights

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

@st.cache_resource
def load_model():
    model = resnet18(weights=None)
    model.fc = torch.nn.Linear(512, 2)
    model_path = os.path.join(BASE_DIR, 'best_model.pth')
    model.load_state_dict(torch.load(model_path, map_location='cpu', weights_only=True))
    model.eval()
    return model

model = load_model()

transform = v2.Compose([
    v2.Resize((224, 224)),
    v2.ToImage(),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

def predict(img_path):
    img = Image.open(img_path).convert("RGB")
    tensor = transform(img).unsqueeze(0)
    with torch.no_grad():
        output = model(tensor)
        pred = torch.argmax(output, dim=1).item()
    return "Real" if pred == 1 else "Fake"

directory = os.path.join(BASE_DIR, "data", "sample")
if "filename" not in st.session_state:
    st.session_state.filename = random.choice(os.listdir(directory))

filename = st.session_state.filename

st.title("AI-Face-Detection")
st.header("Try to find which face is generated and which is real")

left, center, right = st.columns(3)
_, _, col1, col2, col3, _ = st.columns([2, 1, 1, 1, 1, 2])

with center:
    st.image(directory + "/" + filename)

with col1:
    button1 = st.button("Real")
with col2:
    button2 = st.button("Fake")
with col3:
    button3 = st.button("Next")

if button1 or button2:
    user_guess = "Real" if button1 else "Fake"
    true_label = "Real" if "real" in filename.lower() else "Fake"
    model_pred = predict(directory + "/" + filename)
    
    st.write(f"This face was **{true_label}**")
    st.write(f"The model predicted: **{model_pred}**")
    st.write("Correct!" if user_guess == true_label else "Wrong!")

if button3:
    st.session_state.filename = random.choice(os.listdir(directory))
    st.rerun()