import cv2
import streamlit as st
import tempfile
import torch
from torchvision import transforms
from mtcnn import MTCNN
from skimage.feature import hog
import joblib
import numpy as np

# Preprocessing for Siamese Model
def preprocess_image_siamese(img):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor()
    ])
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return transform(img)

# Preprocessing for SVM model (converting to grayscale)
def preprocess_image_svm(img):
    img = cv2.resize(img, (224, 224))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return img

# Extract HOG Features
def extract_hog_features(img):
    hog_features = hog(img, orientations=9, pixels_per_cell=(16, 16), cells_per_block=(4, 4))
    return hog_features

# Detect faces using MTCNN
def get_face(img):
    detector = MTCNN()
    faces = detector.detect_faces(img)
    if faces:
        x1, y1, w, h = faces[0]['box']
        x1, y1 = abs(x1), abs(y1)
        x2, y2 = x1 + w, y1 + h
        return img[y1:y2, x1:x2]
    return None

# Function to verify face (either HOG-SVM or Siamese model)
def verify(img1, img2, model_type, anchor_img):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_img1:
        temp_img1.write(img1.read())
        temp_img1_path = temp_img1.name
    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_img2:
        temp_img2.write(img2.read())
        temp_img2_path = temp_img2.name

    img1p = cv2.imread(temp_img1_path)
    img2p = cv2.imread(temp_img2_path)

    face1 = get_face(img1p)
    face2 = get_face(img2p)

    if face1 is not None and face2 is not None:
        st.image([face1, face2], caption=["Image 1", "Image 2"], width=200)

        if model_type == "HOG-SVM":
            with open('./svm.pkl', 'rb') as f:
                svm = joblib.load(f)
            with open('./pca.pkl', 'rb') as f:
                pca = joblib.load(f)

            face1 = preprocess_image_svm(face1)
            face2 = preprocess_image_svm(face2)

            hog1 = extract_hog_features(face1)
            hog2 = extract_hog_features(face2)

            hog1_pca = pca.transform([hog1])
            hog2_pca = pca.transform([hog2])

            pred1 = svm.predict(hog1_pca)
            pred2 = svm.predict(hog2_pca)

            if pred1 == 1 and pred2 == 1:
                st.write("Matched")
            else:
                st.write("Not Matched")
    else:
        st.write("Face not detected in one or both images")

# Main function to handle Streamlit interaction
def main():
    st.title("Real-time Face Verification App")
    
    model_type = st.selectbox("Select Model", ["Siamese", "HOG-SVM"])
    anchor_img = st.file_uploader("Select Anchor Image", type=["jpg", "png"])

    if model_type == "Siamese":
        # Implement Siamese model choice logic here if needed
        st.write("Using Siamese Network")
        
    elif model_type == "HOG-SVM":
        # Implement HOG-SVM logic here
        st.write("Using HOG-SVM")

    # Camera Input for Face Detection
    run_detection = st.checkbox("Start Camera")
    
    if run_detection:
        cap = cv2.VideoCapture(0)  # Start camera
        
        while True:
            ret, frame = cap.read()
            if not ret:
                st.write("Failed to grab frame.")
                break
            
            # Detect face in the current frame
            face = get_face(frame)
            if face is not None:
                # Draw bounding box around detected face
                x1, y1, x2, y2 = face[0], face[1], face[2], face[3]  # Update face coordinates
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                
                # Show bounding box
                st.image(frame, channels="BGR", use_column_width=True)
            
            # Stop camera when ESC is pressed
            key = cv2.waitKey(1) & 0xFF
            if key == 27:  # ESC key
                break
        
        cap.release()
        cv2.destroyAllWindows()

if __name__ == "__main__":
    main()