import os
import cv2
import tempfile
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from pathlib import Path
import gradio as gr
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter

# Load the model into memory to make running multiple predictions efficient
img_colorization = pipeline(Tasks.image_colorization, model='iic/cv_ddcolor_image-colorization')

def apply_clahe(image):
    # Convert to LAB color space
    lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)
    l, a, b = cv2.split(lab)
    # Apply CLAHE to the L-channel
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    cl = clahe.apply(l)
    # Merge the CLAHE enhanced L-channel back with A and B channels
    limg = cv2.merge((cl,a,b))
    # Convert back to RGB color space
    final = cv2.cvtColor(limg, cv2.COLOR_LAB2RGB)
    return final

def enhance_image_pil(image_pil, brightness=1.0, contrast=1.0):
    # Apply brightness and contrast enhancements
    enhancer_brightness = ImageEnhance.Brightness(image_pil)
    image_pil = enhancer_brightness.enhance(brightness)
    enhancer_contrast = ImageEnhance.Contrast(image_pil)
    image_pil = enhancer_contrast.enhance(contrast)
    # Apply edge enhancement
    image_pil = image_pil.filter(ImageFilter.EDGE_ENHANCE)
    return image_pil

def process_image(img_path, brightness, contrast):
    # Read the input image
    image = cv2.imread(str(img_path))
    # Colorize the image
    output = img_colorization(image[..., ::-1])
    result = output[OutputKeys.OUTPUT_IMG].astype(np.uint8)
    # Apply CLAHE
    result_clahe = apply_clahe(result)
    # Convert to PIL Image for further enhancements
    result_pil = Image.fromarray(result_clahe)
    # Apply brightness, contrast, and edge enhancements
    enhanced_pil = enhance_image_pil(result_pil, brightness, contrast)
    # Save the enhanced image to a temporary file
    temp_dir = tempfile.mkdtemp()
    enhanced_path = os.path.join(temp_dir, 'enhanced.png')
    enhanced_pil.save(enhanced_path)
    return [img_path, enhanced_path], enhanced_path

title = "🌈 Color Restorization Model"
description = "Upload a black & white photo to restore it in color using a deep learning model."

with gr.Blocks(title=title) as demo:
    gr.Markdown(f"## {title}")
    gr.Markdown(description)

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="filepath", label="Upload B&W Image")
            brightness_slider = gr.Slider(0.5, 2.0, value=1.0, label="Brightness")
            contrast_slider = gr.Slider(0.5, 2.0, value=1.0, label="Contrast")
            submit_btn = gr.Button("Colorize")
        with gr.Column():
            comparison = gr.Gallery(label="Original vs Colorized").style(grid=[2], height="auto")
            download_btn = gr.File(label="Download Colorized Image")

    submit_btn.click(
        fn=process_image,
        inputs=[input_image, brightness_slider, contrast_slider],
        outputs=[comparison, download_btn]
    )

demo.launch(enable_queue=True)