import numpy as np
import cv2
import streamlit as st
import insightface
from insightface.app import FaceAnalysis
from PIL import Image
from nudenet import NudeDetector

# Initialize the FaceAnalysis model
app = FaceAnalysis(name='buffalo_l')
app.prepare(ctx_id=0, det_size=(640, 640))  # Ensure the model runs on the CPU if there's limited GPU

# Load the face swapping model (downloading on-demand)
swapper = insightface.model_zoo.get_model('insightface/inswapper_128.onnx', download=True)

# Initialize the NudeDetector model for nudity detection
detector = NudeDetector()

# Function to detect adult content in images
def check_for_adult_content(image):
    # Convert PIL image to a numpy array
    img = np.array(image)

    # Use NudeDetector to predict nudity
    results = detector.detect(img)  # Returns a list of dictionaries

    # Check if any of the results contain explicit labels
    explicit_labels = [
        'EXPOSED_ANUS',
        'EXPOSED_BREAST_F',
        'EXPOSED_GENITALIA_F',
        'EXPOSED_GENITALIA_M'
    ]

    # Iterate through results and check for explicit content
    for item in results:
        if 'label' in item and item['label'] in explicit_labels:
            return True  # Detected explicit content

    return False  # No explicit content detected

# Face swapping function with added sharpening
def swap_faces(destination_image, source_image):
    # Load the destination and source images from Streamlit inputs
    img = cv2.cvtColor(np.array(destination_image), cv2.COLOR_RGB2BGR)
    test = cv2.cvtColor(np.array(source_image), cv2.COLOR_RGB2BGR)

    # Detect faces in the destination and source images
    faces = app.get(img)
    test_faces = app.get(test)

    if not faces or not test_faces:
        return "No faces detected in one or both images."

    test_face = test_faces[0]

    # Perform face swapping with error handling
    res = img.copy()
    try:
        for face in faces:
            res = swapper.get(res, face, test_face, paste_back=True)
    except MemoryError:
        return "Memory error: Face swapping operation failed due to memory overload."

    # Apply sharpening for a clearer image
    kernel = np.array([[0, -1, 0],
                       [-1, 5, -1],
                       [0, -1, 0]])
    sharpened_res = cv2.filter2D(res, -1, kernel)

    # Commented smoothing (blurring) for testing purposes
    smoothed_res = cv2.GaussianBlur(res, (5, 5), 0)

    # Convert the result to RGB for display in Streamlit
    sharpened_res_rgb = cv2.cvtColor(sharpened_res, cv2.COLOR_BGR2RGB)

    return sharpened_res_rgb

# Streamlit app layout
def main():
    st.set_page_config(page_title="SnapSwap", page_icon=":camera:")

    # CSS for styling
    st.markdown("""
        <style>
            .header {
                font-size: 2em;
                font-weight: bold;
                color: #4e73df;
            }
            .subtitle {
                font-size: 1.2em;
                color: #4e73df;
            }
            .content {
                font-size: 1em;
            }
        </style>
    """, unsafe_allow_html=True)

    # Page Header (outside sidebar)
    st.header("SnapSwap :camera:")
    st.write("Welcome to the face swapping app! Upload two images and watch the magic happen!")

    # Sidebar for file uploads and instructions
    with st.sidebar:
        st.subheader("Upload Images for Face Swapping and hit Swap Faces button")
        source_image = st.file_uploader("Upload Source Image", type=["jpg", "png", "jpeg"])
        destination_image = st.file_uploader("Upload Destination Image", type=["jpg", "png", "jpeg"])

        if destination_image and source_image:
            # Check if destination image contains adult content
            destination_img = Image.open(destination_image)
            if check_for_adult_content(destination_img):
                st.error("The destination image does not fall within ethical guidelines.")
            else:
                st.write("Both images uploaded! Click below to swap faces.")
                if st.button("Swap Faces"):
                    with st.spinner("Processing..."):
                        # Read images with PIL (Streamlit works with PIL images)
                        source_img = Image.open(source_image)

                        # Call the face swapping function
                        result = swap_faces(destination_img, source_img)

                        if isinstance(result, str):
                            st.error(result)  # Display error message if no faces detected or memory error
                        else:
                            st.session_state.result = result  # Save the result in session state
                    st.success("Face swapping complete!")

        # Instructions on how to use the app (below upload buttons and swap button)
        st.subheader("How to Use the App:")
        st.markdown("""
            1. **Upload the destination and source images**: Select the images you want to use for face swapping.
            2. **Click on the \"Swap Faces\" button**: The app will process the images and swap the faces.
            3. **View the output**: The swapped face image will appear after processing.
        """)

    # Main content area for displaying the swapped face image
    if 'result' in st.session_state:
        st.image(st.session_state.result, caption="Swapped Face", use_container_width=True)

    # Footer with credits
    st.markdown("Developed with ❤ by Aklavya")

# Run the app
if __name__ == '__main__':
    main()