# Transformers
from transformers import pipeline

# For Face Recognition
import insightface
from insightface.app import FaceAnalysis

# Torch and Image Processing
import torch
import torch.nn.functional as F
from torchvision import transforms, datasets

# For Display and Web Interface
import streamlit as st

# Generic Libraries
import os
import cv2
import numpy as np
from PIL import Image
from glob import glob
import requests

# Set the device (GPU or CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Global Variables
MODEL_TRANSFORMER = 'google/vit-base-patch16-224'
BATCH_SIZE = 8
IMAGE_SHAPE = 640

# Paths
data_path = 'studentsofAIML25'
webcam_path = 'captured_image.jpg'

# Title
st.title("AIML2025 Attendance System")

# Read images from employee directory
image_paths = glob(os.path.join(data_path, '*.jpg'))

# Initialize Face Recognition App
app = FaceAnalysis(name="buffalo_l")
app.prepare(ctx_id=-1, det_size=(IMAGE_SHAPE, IMAGE_SHAPE))  # CPU: -1, GPU: 0

# Initialize HuggingFace ViT Classification Pipeline
vit_pipe = pipeline("image-classification", model=MODEL_TRANSFORMER)

# Face Matching Function
def prod_function(app, prod_path, webcam_path):
    webcam_img = Image.open(webcam_path)
    np_webcam = np.array(webcam_img)
    cv2_webcam = cv2.cvtColor(np_webcam, cv2.COLOR_RGB2BGR)

    webcam_emb = app.get(cv2_webcam, max_num=1)
    if not webcam_emb:
        return torch.tensor([]), None, None

    webcam_emb_tensor = torch.from_numpy(np.array(webcam_emb[0].embedding))

    similarity_score = []
    for path in prod_path:
        img = cv2.imread(path)
        face_embedding = app.get(img, max_num=1)
        if not face_embedding:
            similarity_score.append(torch.tensor(-1.0))
            continue
        face_tensor = torch.from_numpy(np.array(face_embedding[0].embedding))
        similarity_score.append(F.cosine_similarity(face_tensor, webcam_emb_tensor, dim=0))

    return torch.stack(similarity_score), cv2_webcam, webcam_emb[0].bbox.astype(int)

# Streamlit Tabs
about_tab, app_tab = st.tabs(["About the app", "Face Recognition"])

# About Tab
with about_tab:
    st.markdown(
        """
        ## 👁️‍🗨️ AI-Powered Face Recognition Attendance System  
        **Effortless, Secure, and Accurate Attendance with Vision Transformer Technology**  
        
        This is a facial recognition-based attendance system enhanced with image classification using a Vision Transformer (ViT).  
        Just capture your image to mark your attendance automatically!
        """)

    classnames = [os.path.splitext(os.path.basename(path))[0].upper() for path in image_paths]
    st.sidebar.header("📋 Trained Students")
    for name in classnames:
        st.sidebar.write("👤", name)
# Application Tab
with app_tab:
    enable = st.checkbox("Enable Camera")
    picture = st.camera_input("Take a Picture", disabled=not enable)

    if picture is not None:
        with st.spinner("Processing..."):
            # Save captured image
            image = Image.open(picture)
            image.save(webcam_path)

            # --- Face Recognition ---
            prediction, cv2_webcam, face_bbox = prod_function(app, image_paths, webcam_path)

            if prediction is None or len(prediction) == 0:
                st.warning("No face detected in the image.")
            else:
                match_idx = torch.argmax(prediction)

                if prediction[match_idx] >= 0.6:
                    name = image_paths[match_idx].split('/')[-1].split('.')[0].upper()

                    # Draw bounding box and name
                    x1, y1, x2, y2 = face_bbox
                    cv2.rectangle(cv2_webcam, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.rectangle(cv2_webcam, (x1, y2 - 35), (x2, y2), (0, 255, 0), cv2.FILLED)
                    cv2.putText(cv2_webcam, name, (x1 + 6, y2 - 6),
                                cv2.FONT_HERSHEY_COMPLEX, 1, (255, 255, 255), 2)

                    # Show result
                    st.image(cv2.cvtColor(cv2_webcam, cv2.COLOR_BGR2RGB),
                             caption=f"✅ Matched: {name}", use_container_width=True)
                    st.success(f"Welcome, {name}!")
# The use_column_width parameter has been deprecated and will be removed in a future release. Please utilize the use_container_width parameter instead.
                    # --- Send Attendance ---
                    url = "https://skattendancesystem25.glitch.me/adds"
                    data = {'rno': 15, 'sname': name, 'sclass': 7}
                    try:
                        response = requests.post(url, data=data)
                        if response.status_code == 200:
                            st.success("Attendance marked successfully.")
                        else:
                            st.warning("Failed to update attendance.")
                    except Exception as e:
                        st.error(f"Request failed: {e}")
                else:
                    st.error("❌ Match not found. Try again.")