import gradio as gr
import numpy as np
from PIL import Image
import torch
import torchvision.transforms as T
from skimage import color as skcolor

# -----------------------------
# DEVICE
# -----------------------------
DEVICE = "cpu"


# -----------------------------
# 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)


# -----------------------------
# ZHANG ECCV16 COLORIZER
# -----------------------------
try:
    import colorizers
    _ZHANG_AVAILABLE = True
except Exception:
    _ZHANG_AVAILABLE = False
    colorizers = None

_zhang_model = None
_zhang_transform = T.Compose([
    T.Resize(256),
    T.CenterCrop(256),
    T.ToTensor(),
])


def load_zhang_model():
    global _zhang_model
    if not _ZHANG_AVAILABLE:
        return None
    if _zhang_model is None:
        _zhang_model = colorizers.eccv16().eval().to(DEVICE)
    return _zhang_model


def zhang_colorize(pil_img: Image.Image) -> Image.Image:
    model = load_zhang_model()
    if model is None:
        return test_mode_colorize(pil_img)

    img = pil_img.convert("RGB")
    img_resized = _zhang_transform(img).unsqueeze(0).to(DEVICE)

    np_img = img_resized[0].permute(1, 2, 0).cpu().numpy()
    lab = skcolor.rgb2lab(np_img)
    L = lab[:, :, 0]

    tens_l = torch.from_numpy(L).unsqueeze(0).unsqueeze(0).float().to(DEVICE)

    with torch.no_grad():
        out_ab = model(tens_l).cpu()

    out_ab = out_ab[0].permute(1, 2, 0).numpy()

    H_orig, W_orig = img.size[1], img.size[0]
    out_ab_resized = np.array(
        Image.fromarray((out_ab * 255).astype(np.uint8)).resize((W_orig, H_orig), Image.BILINEAR),
        dtype=np.float32
    ) / 255.0

    img_np = np.array(img).astype(np.float32) / 255.0
    lab_orig = skcolor.rgb2lab(img_np)
    L_orig = lab_orig[:, :, 0]

    lab_out = np.zeros((H_orig, W_orig, 3), dtype=np.float32)
    lab_out[:, :, 0] = L_orig
    lab_out[:, :, 1:] = out_ab_resized * 128.0

    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)


# -----------------------------
# 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 == "Zhang ECCV16 (Deep Colorizer)":
        return zhang_colorize(pil_img)

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

    return test_mode_colorize(pil_img)


# -----------------------------
# GRADIO UI
# -----------------------------
with gr.Blocks(title="Biker Image Colorizer – CPU (Test + Zhang + DeOldify‑Lite)") as demo:
    gr.Markdown(
        """
        # Biker Image Colorizer – CPU Edition  
        **Three modes:**  
        - Test Mode (very subtle)  
        - Zhang ECCV16 (deep neural colorizer)  
        - DeOldify‑Lite (artistic 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)",
                    "Zhang ECCV16 (Deep Colorizer)",
                    "DeOldify‑Lite (Art Mode)",
                ],
                value="Zhang ECCV16 (Deep Colorizer)",
                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()