import cv2
import numpy as np
import io
import streamlit as st
from PIL import Image

def color_transfer(source, target):
    # Convert images to LAB color space
    source = cv2.cvtColor(source, cv2.COLOR_BGR2LAB).astype("float32")
    target = cv2.cvtColor(target, cv2.COLOR_BGR2LAB).astype("float32")

    # Compute mean and std for each channel
    mean_src, std_src = cv2.meanStdDev(source)
    mean_tgt, std_tgt = cv2.meanStdDev(target)

    # Reshape for broadcasting
    mean_src, std_src = mean_src.reshape((1, 1, 3)), std_src.reshape((1, 1, 3))
    mean_tgt, std_tgt = mean_tgt.reshape((1, 1, 3)), std_tgt.reshape((1, 1, 3))

    # Apply color transfer
    adjusted = ((target - mean_tgt) * (std_src / std_tgt)) + mean_src
    adjusted = np.clip(adjusted, 0, 255).astype("uint8")

    return cv2.cvtColor(adjusted, cv2.COLOR_LAB2BGR)

def main():
    st.set_page_config(layout="wide")
    st.sidebar.header("Upload Images")
    
    # Upload two images in the sidebar
    uploaded_file1 = st.sidebar.file_uploader("Upload first image", type=["jpg", "png", "jpeg"])
    uploaded_file2 = st.sidebar.file_uploader("Upload second image", type=["jpg", "png", "jpeg"])
    
    st.sidebar.image("test.png", width=700)
    
    # Display the images in the center
    col1, col2, col3 = st.columns([1, 1, 1])
    
    with col1:
        st.header("Uploaded Images")
        
        if uploaded_file1 and uploaded_file2:
            ref_img = Image.open(uploaded_file1)
            input_img = Image.open(uploaded_file2)
            
            # Display images side by side
            st.image([ref_img, input_img], caption=["Reference Image", "Input Image"], width=300)
        elif uploaded_file1:
            ref_img = Image.open(uploaded_file1)
            st.image(ref_img, caption="Reference Image", width=300)
        elif uploaded_file2:
            input_img = Image.open(uploaded_file2)
            st.image(input_img, caption="Input Image", width=300)
        else:
            st.write("Upload images to display them here.")
    
    with col2:
        st.header(" ")
        if uploaded_file1 and uploaded_file2:
            img1 = Image.open(uploaded_file1)
            img2 = Image.open(uploaded_file2)
            
            ref_img = cv2.cvtColor(np.array(img1), cv2.COLOR_RGB2BGR)
            input_img = cv2.cvtColor(np.array(img2), cv2.COLOR_RGB2BGR)
            
            orig_h, orig_w = input_img.shape[:2]
            input_img = cv2.resize(input_img, (ref_img.shape[1], ref_img.shape[0]))
            # Apply color transfer
            output_img = color_transfer(ref_img, input_img)
            output_img_resize = cv2.resize(output_img, (orig_w, orig_h))
            st.image(output_img_resize, caption="Result", width=500, channels="BGR")
            
            _, buffer = cv2.imencode(".jpg", output_img_resize)
            img_bytes = io.BytesIO(buffer)

            # Add Download Button
            st.download_button(
                label="Download Image",
                data=img_bytes,
                file_name="image.jpg",
                mime="image/jpeg"
            )

if __name__ == "__main__":
    main()