import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
import cv2
import openai
import os
import io

# Set up OpenAI API key
openai.api_key = "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"

# Load pre-trained model for face detection
face_detection_model = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

def detect_face(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    faces = face_detection_model.detectMultiScale(gray, 1.3, 5)
    if len(faces) > 0:
        (x, y, w, h) = faces[0]
        return image[y:y+h, x:x+w]
    return None

def analyze_symmetry(face):
    height, width = face.shape[:2]
    left_half = face[:, :width//2]
    right_half = cv2.flip(face[:, width//2:], 1)
    diff = cv2.absdiff(left_half, right_half)
    symmetry_score = 1 - (np.sum(diff) / (255 * height * (width//2)))
    return symmetry_score

def analyze_aesthetic_quality(image):
    # Convert the image to a byte stream
    img_byte_arr = io.BytesIO()
    Image.fromarray(image).save(img_byte_arr, format='PNG')
    img_byte_arr = img_byte_arr.getvalue()

    # Encode the image
    encoded_image = base64.b64encode(img_byte_arr).decode('ascii')

    prompt = f"Analyze the aesthetic quality of this image, focusing on lighting, composition, and overall visual appeal. Provide a brief description and a score out of 10."
    response = openai.ChatCompletion.create(
        model="gpt-4-vision-preview",
        messages=[
            {
                "role": "user", 
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{encoded_image}"}}
                ]
            }
        ]
    )
    return response.choices[0].message.content

def compare_photos(image1, image2):
    face1 = detect_face(image1)
    face2 = detect_face(image2)
    
    if face1 is None or face2 is None:
        return "Error: Face not detected in one or both images."
    
    symmetry1 = analyze_symmetry(face1)
    symmetry2 = analyze_symmetry(face2)
    
    aesthetic1 = analyze_aesthetic_quality(image1)
    aesthetic2 = analyze_aesthetic_quality(image2)
    
    return {
        "Image 1": {
            "Symmetry": symmetry1,
            "Aesthetic Quality": aesthetic1
        },
        "Image 2": {
            "Symmetry": symmetry2,
            "Aesthetic Quality": aesthetic2
        }
    }

def main():
    st.title("AttraVision")
    st.write("Compare two photos based on attractiveness parameters")
    
    uploaded_file1 = st.file_uploader("Choose the first image", type=["jpg", "jpeg", "png"])
    uploaded_file2 = st.file_uploader("Choose the second image", type=["jpg", "jpeg", "png"])
    
    if uploaded_file1 is not None and uploaded_file2 is not None:
        image1 = Image.open(uploaded_file1)
        image2 = Image.open(uploaded_file2)
        
        col1, col2 = st.columns(2)
        with col1:
            st.image(image1, caption="Image 1", use_column_width=True)
        with col2:
            st.image(image2, caption="Image 2", use_column_width=True)
        
        if st.button("Compare Photos"):
            with st.spinner("Analyzing..."):
                results = compare_photos(np.array(image1), np.array(image2))
                
                st.subheader("Comparison Results")
                col1, col2 = st.columns(2)
                with col1:
                    st.write("Image 1")
                    st.write(f"Symmetry: {results['Image 1']['Symmetry']:.2f}")
                    st.write(f"Aesthetic Quality: {results['Image 1']['Aesthetic Quality']}")
                with col2:
                    st.write("Image 2")
                    st.write(f"Symmetry: {results['Image 2']['Symmetry']:.2f}")
                    st.write(f"Aesthetic Quality: {results['Image 2']['Aesthetic Quality']}")
                
                st.write("Note: Beauty is subjective, and this tool is for entertainment purposes only.")
    
    st.sidebar.title("About AttraVision")
    st.sidebar.info("AttraVision uses machine learning to analyze and compare photos based on predefined attractiveness parameters. This tool is designed for fun and should not be used for serious judgment.")
    st.sidebar.warning("Privacy Notice: Your photos are processed securely and are not stored after analysis.")

if __name__ == "__main__":
    main()