import gradio as gr
import numpy as np
from PIL import Image
from skimage import color as skcolor

# Import Zhang full model (download logic is inside zhang_colorizer.py)
from zhang_colorizer import zhang_colorize

# -----------------------------
# TEST MODE (subtle tint)
# -----------------------------
def test_mode_colorize(pil_img: Image.Image) -> Image.Image:
    img = pil_img.convert("RGB")
    arr = np.array(img).astype(np.float32) / 255.0
    tint = np.array([1.02, 1.0, 0.98], dtype=np.float32)
    arr = np.clip(arr * tint, 0.0, 1.0)
    arr = (arr * 255).astype(np.uint8)
    return Image.fromarray(arr)

# -----------------------------
# DEOLDIFY‑LITE (CPU SAFE)
# -----------------------------
def deoldify_lite_colorize(pil_img: Image.Image) -> Image.Image:
    img = pil_img.convert("RGB")
    arr = np.array(img).astype(np.float32) / 255.0

    lab = skcolor.rgb2lab(arr)
    L = lab[:, :, 0]
    a = lab[:, :, 1]
    b = lab[:, :, 2]

    a *= 1.35
    b *= 1.35
    a += 2.0
    b += 1.0

    lab_out = np.stack([L, a, b], axis=-1)
    rgb_out = skcolor.lab2rgb(lab_out)
    rgb_out = np.clip(rgb_out, 0.0, 1.0)
    rgb_out = (rgb_out * 255).astype(np.uint8)

    return Image.fromarray(rgb_out)

# -----------------------------
# MAIN PIPELINE
# -----------------------------
def colorize_image(input_image, mode):
    if input_image is None:
        return None

    pil_img = input_image.convert("RGB")

    if mode == "Test Mode (Very Subtle)":
        return test_mode_colorize(pil_img)

    if mode == "DeOldify‑Lite (Art Mode)":
        return deoldify_lite_colorize(pil_img)

    if mode == "Zhang ECCV (Realistic Strong Color)":
        return zhang_colorize(pil_img)

    return test_mode_colorize(pil_img)

# -----------------------------
# GRADIO UI
# -----------------------------
with gr.Blocks(title="Biker Image Colorizer – CPU (3 Modes)") as demo:
    gr.Markdown(
        """
        # Biker Image Colorizer – CPU Edition  
        **Three modes:**  
        - Test Mode (very subtle)  
        - DeOldify‑Lite (artistic strong color)  
        - Zhang ECCV (realistic strong color)  
        """
    )

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")
            mode = gr.Radio(
                choices=[
                    "Test Mode (Very Subtle)",
                    "DeOldify‑Lite (Art Mode)",
                    "Zhang ECCV (Realistic Strong Color)",
                ],
                value="Zhang ECCV (Realistic Strong Color)",
                label="Colorization Mode",
            )
            run_btn = gr.Button("Colorize", variant="primary")

        with gr.Column():
            output_image = gr.Image(type="pil", label="Output Image")

    run_btn.click(
        fn=colorize_image,
        inputs=[input_image, mode],
        outputs=[output_image],
    )

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