# Transformers and its models
#import transformers

# For Image Processing
#from transformers import ViTImageProcessor

# For Model
#from transformers import ViTModel, ViTConfig, pipeline
import insightface
from insightface.app import FaceAnalysis

# For data augmentation
from torchvision import transforms, datasets


# For Data Loaders
import datasets
from torch.utils.data import Dataset, DataLoader

# For Display
#from tqdm.notebook import tqdm

# Other Generic Libraries
import torch
from PIL import Image
import cv2
import os
import streamlit as st
import gc
from glob import glob
import shutil
import pandas as pd
import numpy as np
#import matplotlib.pyplot as plt
from io import BytesIO
import torch.nn.functional as F

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

# Initialse Globle Variables
MODEL_TRANSFORMER = 'google/vit-base-patch16-224'
BATCH_SIZE = 8

# Set Paths
data_path = 'employees'
model_path = 'vit_pytorch_GPU_1.pt'
webcam_path = 'captured_image.jpg'

IMAGE_SHAPE = 640

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

# Read images from directory
image_paths = []
image_file = glob(os.path.join(data_path, '*.jpg'))
#st.write(image_file)
image_paths.extend(image_file)
#st.write('input path size:', len(image_paths))
#st.write(image_paths)

# Initialize the app
app = FaceAnalysis(name="buffalo_l")  # buffalo_l includes ArcFace model
app.prepare(ctx_id=-1, det_size=(IMAGE_SHAPE, IMAGE_SHAPE))  # Use ctx_id=-1 if you want CPU, and ctx_id=0 for GPU



# Define the ML model - Evaluation function
def prod_function(app, prod_path, webcam_path):
    webcam_img = Image.open(webcam_path)
    np_webcam = np.array(webcam_img) # Convert to NumPy array
    cv2_webcam = cv2.cvtColor(np_webcam, cv2.COLOR_RGB2BGR) # Convert RGB (PIL) to BGR (OpenCV)
    
    webcam_emb = app.get(cv2_webcam, max_num=1)
    webcam_emb = webcam_emb[0].embedding
    webcam_emb = torch.from_numpy(np.array(webcam_emb))
    
    similarity_score = []
    for path in prod_path:
        img = cv2.imread(path)
        face_embedding = app.get(img, max_num=1)
        face_embedding = face_embedding[0].embedding
        face_embedding = torch.from_numpy(np.array(face_embedding))
        
        similarity_score.append(F.cosine_similarity(face_embedding,webcam_emb, dim=0))
        #distance = F.pairwise_distance(emb, emb_prod)
        #prod_preds.append(distance)
    similarity_score = torch.from_numpy(np.array(similarity_score))
    return similarity_score #prod_preds

about_tab, app_tab = st.tabs(["About the app", "Face Recognition"])
# About the app Tab
with about_tab:
    st.markdown(
        """
        # 👁️‍🗨️ AI-Powered Face Recognition Attendance System
        Effortless, Secure, and Accurate Attendance with Vision Transformer Technology
        
        An intelligent, facial recognition-based attendance solution that redefines how organizations manage employee presence. By leveraging cutting-edge computer vision and AI, the app automates attendance tracking with speed, precision, and reliability—no timecards, no fingerprint scans, just a glance.
        """)

# Gesture recognition Tab
with app_tab:
    # Read image from Camera
    enable = st.checkbox("Enable camera")
    picture = st.camera_input("Take a picture", disabled=not enable)
    if picture is not None:    
        with st.spinner("Wait for it...", show_time=True):
            # Run the predictions
            prediction = prod_function(app, image_paths, picture)
            #prediction = torch.cat(prediction, 0).to(device)
            match_idx = torch.argmax(prediction)
            st.write(prediction)
            st.write(image_paths)
        
            # Display the results
            if prediction[match_idx] >= 0.6:
                pname = image_paths[match_idx].split('/')[-1].split('.')[0]
                st.write('Welcome: ',pname)

                ###### upload the data to Glitch https://aimljan25f.glitch.me/addg
                # # using post method
                url = "https://sehajpreet-aiml-2025.glitch.me/addg"
                 
                data = {'rno': '15','sname': pname, 'sclass': '7'  }
                # response = requests.post(url +url1 , data=data)
    
                # Post Method is invoked if data != None
                req =  request.Request(url , method="GET", data=data)
                
                # Response
                resp = request.urlopen(req)
                if response.status_code == 200:
                    st.success("Data updated on: " + "https://sehajpreet-aiml-2025.glitch.me/")
                else:
                    st.warning("Data not updated")
                ########## end update website 

            else:
              st.write("Match not found")