import streamlit as st
import numpy as np
import cv2 as cv
from scipy.special import softmax
import time
import io
from threading import Lock

# Define the HairSegmentation class (same as before)
class HairSegmentation(object):
    def __init__(self):
        self.net = cv.dnn.readNet("hair_segmentation.tflite")

    def _mix_prev_mask(self, prev_mask, new_mask):
        combine_with_prev_ratio = 0.9
        eps = 1e-3
        uncertainty_alpha = 1.0 + (new_mask * np.log(new_mask + eps) + (1.0 - new_mask) * np.log(1.0 - new_mask + eps)) / np.log(2.0)
        uncertainty_alpha = np.clip(uncertainty_alpha, 0, 1)
        uncertainty_alpha *= 2.0 - uncertainty_alpha

        mixed_mask = new_mask * uncertainty_alpha + prev_mask * (1.0 - uncertainty_alpha)
        return mixed_mask * combine_with_prev_ratio + (1.0 - combine_with_prev_ratio) * new_mask

    def process_image(self, frame, color, num_runs=2):
        prev_mask = np.zeros((512, 512), dtype=np.float32)
        color = np.ones(frame.shape, dtype=np.uint8) * color

        # Prepare input
        blob = cv.dnn.blobFromImage(frame, 1.0 / 255, (512, 512), swapRB=True)
        blob = np.concatenate((blob, prev_mask.reshape(1, 1, 512, 512)), axis=1)

        for i in range(num_runs):
            # Copy previous frame mask to a new tensor
            blob[0, 3] = prev_mask

            # Run network
            self.net.setInput(blob)
            out = self.net.forward()

            out = softmax(out, axis=1)
            mask = out[0, 1]

            prev_mask = self._mix_prev_mask(prev_mask, mask)

        mask = cv.resize(prev_mask, (frame.shape[1], frame.shape[0]))
        lum = cv.cvtColor(frame, cv.COLOR_BGR2GRAY) / 255
        mask *= lum

        mask = np.repeat(np.expand_dims(mask, axis=-1), 3, axis=-1)
        result = (mask * (color.astype(np.float32) - frame) + frame).astype(np.uint8)

        return result


# Create an instance of HairSegmentation
model = HairSegmentation()

# Initialize color dictionary
colors = {}
timestamps = {}

# Initialize Streamlit application
st.title("Hair Color Segmentation")
st.write("Upload a photo, choose a color for the hair, and see the transformation!")

# Set up image upload
uploaded_image = st.file_uploader("Upload your image", type=["jpg", "jpeg", "png"])

# Color picker for hair color
color = st.color_picker("Pick a hair color", "#FF0000")

# Ensure the user inputted color hex is valid
if color.startswith('#') and len(color) == 7:
    # Remove the '#' symbol and then convert hex color to BGR format
    hex_color = color[1:]  # Remove the '#' symbol
    color_bgr = [int(hex_color[i:i+2], 16) for i in (4, 2, 0)]  # Convert hex to BGR format
else:
    st.error("Invalid color code! Please select a valid color.")

# Process image when it's uploaded
if uploaded_image is not None:
    # Read image
    img = np.frombuffer(uploaded_image.read(), dtype=np.uint8)
    img = cv.imdecode(img, cv.IMREAD_COLOR)

    # Apply hair segmentation and color transformation
    stylized_image = model.process_image(img, color_bgr)

    # Display the original and stylized images without color distortion
    st.image(cv.cvtColor(img, cv.COLOR_BGR2RGB), caption="Original Image", use_column_width=True)  # Convert to RGB for proper display
    st.image(cv.cvtColor(stylized_image, cv.COLOR_BGR2RGB), caption="Stylized Image with New Hair Color", use_column_width=True)  # Convert to RGB for proper display

    # Optionally, save the result
    save_result = st.button("Save Result")
    if save_result:
        # Convert image to bytes and allow downloading
        _, buffer = cv.imencode(".jpg", stylized_image)
        st.download_button(
            label="Download Stylized Image",
            data=buffer.tobytes(),
            file_name="stylized_image.jpg",
            mime="image/jpeg"
        )