import gradio as gr
import numpy as np
from PIL import Image


def apply_photoshop_color_balance_with_controls(
    image_pil, cyan_shift, magenta_shift, yellow_shift, brightness, contrast
):
    # Check if image is None and return early
    if image_pil is None:
        return None

    # Convert to RGB float [0–1]
    image_rgb = np.array(image_pil.convert("RGB")).astype(np.float32) / 255.0
    R0, G0, B0 = image_rgb[..., 0], image_rgb[..., 1], image_rgb[..., 2]

    # Valid area (exclude black background)
    valid_mask = np.any(image_rgb > 0.05, axis=-1)

    # Original luminance
    orig_lum = 0.299 * R0 + 0.587 * G0 + 0.114 * B0

    # Convert to CMY
    C = 1 - R0
    M = 1 - G0
    Y = 1 - B0

    # CMY Shifts scaled to [-1, 1]
    delta_C = cyan_shift / 100.0
    delta_M = magenta_shift / 100.0
    delta_Y = yellow_shift / 100.0

    # Apply CMY shifts to valid pixels
    C[valid_mask] = np.clip(C[valid_mask] + delta_C, 0, 1)
    M[valid_mask] = np.clip(M[valid_mask] + delta_M, 0, 1)
    Y[valid_mask] = np.clip(Y[valid_mask] + delta_Y, 0, 1)

    # Convert back to RGB
    R = 1 - C
    G = 1 - M
    B = 1 - Y

    # Preserve brightness
    new_lum = 0.299 * R + 0.587 * G + 0.114 * B
    gain = np.ones_like(new_lum)
    gain[valid_mask] = (orig_lum[valid_mask] + 1e-5) / (new_lum[valid_mask] + 1e-5)

    R = np.clip(R * gain, 0, 1)
    G = np.clip(G * gain, 0, 1)
    B = np.clip(B * gain, 0, 1)

    # Stack image and apply brightness/contrast
    result = np.stack([R, G, B], axis=-1)

    # Brightness [-100, 100] → shift [-0.4, +0.4]
    brightness_shift = brightness / 250.0
    result = np.clip(result + brightness_shift, 0, 1)

    # Contrast [-100, 100] → scale [0.6, 1.4]
    contrast_factor = 1.0 + (contrast / 250.0)
    result = np.clip((result - 0.5) * contrast_factor + 0.5, 0, 1)

    return Image.fromarray((result * 255).astype(np.uint8))

show_sliders = False
share = False

iface = gr.Interface(
    fn=apply_photoshop_color_balance_with_controls,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(-100, 100, value=37, label="Red–Cyan", visible=show_sliders),
        gr.Slider(-100, 100, value=-20, label="Green-Magenta", visible=show_sliders),
        gr.Slider(-100, 100, value=-10, label="Blue-Yellow", visible=show_sliders),
        gr.Slider(-100, 100, value=-25, label="Brightness", visible=show_sliders),
        gr.Slider(-100, 100, value=34, label="Contrast", visible=show_sliders),
    ],
    live=True,
    outputs=gr.Image(type="pil", label="Processed Image"),
    title="Image Color Balancer with CMY + Brightness/Contrast",
    description="Apply Photoshop-like CMY color shifts in shadows. Optionally adjust brightness and contrast.",
)


if __name__ == "__main__":
    iface.launch(share=share)