# app.py
# Hugging Face Space (Gradio) for color correction:
# Takes two images: (1) image to correct, (2) reference image.
# Returns a single, color-corrected image.

import gradio as gr
from PIL import Image
import numpy as np
import cv2
from skimage import exposure


def _to_rgb_np(img: Image.Image) -> np.ndarray:
    """PIL -> RGB numpy uint8"""
    if img.mode != "RGB":
        img = img.convert("RGB")
    return np.array(img)


def _lab_histogram_match(src_rgb: np.ndarray, ref_rgb: np.ndarray) -> np.ndarray:
    """
    Match histogram of src to ref in LAB space and return as RGB uint8.
    """
    src_lab = cv2.cvtColor(src_rgb, cv2.COLOR_RGB2LAB)
    ref_lab = cv2.cvtColor(ref_rgb, cv2.COLOR_RGB2LAB)

    matched_lab = exposure.match_histograms(src_lab, ref_lab, channel_axis=2)
    matched_lab = np.clip(matched_lab, 0, 255).astype(np.uint8)

    matched_rgb = cv2.cvtColor(matched_lab, cv2.COLOR_LAB2RGB)
    return matched_rgb


def _luminosity_blend(base_rgb: np.ndarray, blend_rgb: np.ndarray) -> np.ndarray:
    """
    Photoshop-like 'Luminosity' blend: keep hue/saturation from base,
    take luminance from blend. Implemented via HLS space.
    """
    base_hls = cv2.cvtColor(base_rgb, cv2.COLOR_RGB2HLS)
    blend_hls = cv2.cvtColor(blend_rgb, cv2.COLOR_RGB2HLS)

    out_hls = base_hls.copy()
    out_hls[..., 1] = blend_hls[..., 1]  # replace Lightness channel

    out_rgb = cv2.cvtColor(out_hls, cv2.COLOR_HLS2RGB)
    return out_rgb


def color_correct(img_to_correct: Image.Image, reference_img: Image.Image) -> Image.Image:
    """
    Color-correct `img_to_correct` to match the look of `reference_img`.
    Steps:
    1) Histogram match in LAB space.
    2) Luminosity blend to preserve original hue/saturation.
    """
    if img_to_correct is None or reference_img is None:
        return None

    src = _to_rgb_np(img_to_correct)
    ref = _to_rgb_np(reference_img)

    matched = _lab_histogram_match(src, ref)
    result = _luminosity_blend(src, matched)

    return Image.fromarray(result)


title = "Color Correction (Histogram Match + Luminosity Blend)"
description = (
    "Upload the image you want to correct (left) and a reference image (right). "
    "The output transfers the overall tonal feel of the reference while preserving "
    "the original image's colors via a luminosity blend."
)

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

    with gr.Row():
        with gr.Column():
            img_a = gr.Image(type="pil", label="Image to Correct")
        with gr.Column():
            img_b = gr.Image(type="pil", label="Reference Image")

    out = gr.Image(type="pil", label="Corrected Image")

    btn = gr.Button("Run Color Correction")
    btn.click(color_correct, inputs=[img_a, img_b], outputs=[out])

if __name__ == "__main__":
    demo.launch()