from pathlib import Path

import gradio as gr
import numpy as np
import pi_heif
import spaces
import torch
from huggingface_hub import hf_hub_download
from PIL import Image, ImageFilter
from refiners.foundationals.latent_diffusion import Solver, solvers

from enhancer import ESRGANUpscaler, ESRGANUpscalerCheckpoints

pi_heif.register_heif_opener()


TITLE = """
<h1>Finegrain Image Enhancer</h1>
<p>
    Transform low-resolution images into stunning high-resolution versions with intelligently generated details.
</p><p>
    This Space showcases one of Finegrain's earlier generative AI models. Today, we're building the Agentic Camera: an on-device system that produces the intended photo at capture time, eliminating image editing in most cases.
</p><p>
    ✨ Learn more at <a href="https://finegrain.ai">finegrain.ai</a> or follow <a href="https://huggingface.co/finegrain">Finegrain</a> for updates.
</p>
"""

DEFAULT_PROMPT = (
    "raw photo, candid photograph, photorealistic, natural skin texture, "
    "visible skin pores, subtle skin imperfections, "
    "natural hair strands, unretouched, film grain, dslr"
)
DEFAULT_NEGATIVE_PROMPT = (
    "cartoon, anime, illustration, painting, drawing, 3d render, cgi, "
    "airbrushed, plastic skin, smooth skin, flawless skin, waxy skin, doll, "
    "retouched, beauty filter, perfect face, makeup, glamour, de-aged, baby face, "
    "thick beard, heavy stubble, exaggerated facial hair, "
    "oversaturated, oversharpened, worst quality, low quality, normal quality"
)

CHECKPOINTS = ESRGANUpscalerCheckpoints(
    unet=Path(
        hf_hub_download(
            repo_id="refiners/juggernaut.reborn.sd1_5.unet",
            filename="model.safetensors",
            revision="347d14c3c782c4959cc4d1bb1e336d19f7dda4d2",
        )
    ),
    clip_text_encoder=Path(
        hf_hub_download(
            repo_id="refiners/juggernaut.reborn.sd1_5.text_encoder",
            filename="model.safetensors",
            revision="744ad6a5c0437ec02ad826df9f6ede102bb27481",
        )
    ),
    lda=Path(
        hf_hub_download(
            repo_id="refiners/juggernaut.reborn.sd1_5.autoencoder",
            filename="model.safetensors",
            revision="3c1aae3fc3e03e4a2b7e0fa42b62ebb64f1a4c19",
        )
    ),
    controlnet_tile=Path(
        hf_hub_download(
            repo_id="refiners/controlnet.sd1_5.tile",
            filename="model.safetensors",
            revision="48ced6ff8bfa873a8976fa467c3629a240643387",
        )
    ),
    esrgan=Path(
        hf_hub_download(
            repo_id="uwg/upscaler",
            filename="ESRGAN/4x_NMKD-Superscale-SP_178000_G.pth",
            revision="f6bace545e358eab5491f8f39b90a2dd42e8cc77",
        )
    ),
    loras={
        "more_details": Path(
            hf_hub_download(
                repo_id="philz1337x/loras",
                filename="more_details.safetensors",
                revision="a3802c0280c0d00c2ab18d37454a8744c44e474e",
            )
        ),
        "sdxl_render": Path(
            hf_hub_download(
                repo_id="philz1337x/loras",
                filename="SDXLrender_v2.0.safetensors",
                revision="a3802c0280c0d00c2ab18d37454a8744c44e474e",
            )
        ),
    },
)

# initialize the enhancer, on the cpu
DEVICE_CPU = torch.device("cpu")
DTYPE = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
enhancer = ESRGANUpscaler(checkpoints=CHECKPOINTS, device=DEVICE_CPU, dtype=DTYPE)

# "move" the enhancer to the gpu, this is handled by Zero GPU
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
enhancer.to(device=DEVICE, dtype=DTYPE)


def match_colors(image: Image.Image, reference: Image.Image) -> Image.Image:
    # Diffusion shifts colors slightly (brighter skin, washed-out veins);
    # transfer the reference's per-channel mean/std to bring them back.
    img = np.asarray(image.convert("RGB")).astype(np.float32)
    ref = np.asarray(reference.convert("RGB")).astype(np.float32)
    for c in range(3):
        img_mean, img_std = img[..., c].mean(), img[..., c].std()
        ref_mean, ref_std = ref[..., c].mean(), ref[..., c].std()
        if img_std > 1e-5:
            img[..., c] = (img[..., c] - img_mean) * (ref_std / img_std) + ref_mean
    return Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))


@spaces.GPU(duration=120)
def process(
    input_image: Image.Image,
    prompt: str = DEFAULT_PROMPT,
    negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
    seed: int = 42,
    upscale_factor: int = 2,
    controlnet_scale: float = 1.1,
    controlnet_decay: float = 1.0,
    condition_scale: int = 2,
    tile_width: int = 112,
    tile_height: int = 144,
    denoise_strength: float = 0.1,
    num_inference_steps: int = 18,
    solver: str = "DDIM",
) -> tuple[Image.Image, Image.Image]:
    solver_type: type[Solver] = getattr(solvers, solver)

    generator = torch.Generator(device=DEVICE)
    generator.manual_seed(seed)

    # Resize to avoid using too much VRAM.
    # If you have a bug GPU you can go higher.
    side_size = min(input_image.size)
    if side_size > 1024:
        scale = 1024 / side_size
        new_size = (int(input_image.width * scale), int(input_image.height * scale))
        resized_image = input_image.resize(new_size, resample=Image.Resampling.LANCZOS)
    else:
        resized_image = input_image

    enhanced_image = enhancer.upscale(
        image=resized_image,
        prompt=prompt,
        negative_prompt=negative_prompt,
        upscale_factor=upscale_factor,
        controlnet_scale=controlnet_scale,
        controlnet_scale_decay=controlnet_decay,
        condition_scale=condition_scale,
        tile_size=(tile_height, tile_width),
        denoise_strength=denoise_strength,
        num_inference_steps=num_inference_steps,
        loras_scale={"more_details": 0.55, "sdxl_render": 0.0},
        solver_type=solver_type,
        generator=generator,
    )
    enhanced_image = match_colors(enhanced_image, resized_image)
    # Deterministic texture pass: diffusion is now dialed down for face fidelity,
    # so skin/hair micro-detail is restored here instead, with zero identity bias.
    enhanced_image = enhanced_image.filter(ImageFilter.UnsharpMask(radius=3, percent=60, threshold=3))

    return (input_image, enhanced_image)


with gr.Blocks() as demo:
    gr.HTML(TITLE)

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")
            run_button = gr.ClearButton(components=None, value="Enhance Image")
        with gr.Column():
            output_slider = gr.ImageSlider(label="Before / After", max_height=1500, show_fullscreen_button=False)
            run_button.add(output_slider)

    with gr.Accordion("Advanced Options", open=False):
        prompt = gr.Textbox(
            label="Prompt",
            value=DEFAULT_PROMPT,
            placeholder=DEFAULT_PROMPT,
        )
        negative_prompt = gr.Textbox(
            label="Negative Prompt",
            value=DEFAULT_NEGATIVE_PROMPT,
            placeholder=DEFAULT_NEGATIVE_PROMPT,
        )
        seed = gr.Slider(
            minimum=0,
            maximum=10_000,
            value=42,
            step=1,
            label="Seed",
        )
        upscale_factor = gr.Slider(
            minimum=1,
            maximum=4,
            value=2,
            step=0.2,
            label="Upscale Factor",
        )
        controlnet_scale = gr.Slider(
            minimum=0,
            maximum=1.5,
            value=1.1,
            step=0.05,
            label="ControlNet Scale",
        )
        controlnet_decay = gr.Slider(
            minimum=0.5,
            maximum=1,
            value=1.0,
            step=0.025,
            label="ControlNet Scale Decay",
        )
        condition_scale = gr.Slider(
            minimum=2,
            maximum=20,
            value=2,
            step=1,
            label="Condition Scale",
        )
        tile_width = gr.Slider(
            minimum=64,
            maximum=200,
            value=112,
            step=1,
            label="Latent Tile Width",
        )
        tile_height = gr.Slider(
            minimum=64,
            maximum=200,
            value=144,
            step=1,
            label="Latent Tile Height",
        )
        denoise_strength = gr.Slider(
            minimum=0,
            maximum=1,
            value=0.1,
            step=0.05,
            label="Denoise Strength",
        )
        num_inference_steps = gr.Slider(
            minimum=1,
            maximum=30,
            value=18,
            step=1,
            label="Number of Inference Steps",
        )
        solver = gr.Radio(
            choices=["DDIM", "DPMSolver"],
            value="DDIM",
            label="Solver",
        )

    run_button.click(
        fn=process,
        inputs=[
            input_image,
            prompt,
            negative_prompt,
            seed,
            upscale_factor,
            controlnet_scale,
            controlnet_decay,
            condition_scale,
            tile_width,
            tile_height,
            denoise_strength,
            num_inference_steps,
            solver,
        ],
        outputs=output_slider,
    )

    gr.Examples(
        examples=[
            "examples/kara-eads-L7EwHkq1B2s-unsplash.jpg",
            "examples/clarity_bird.webp",
            "examples/edgar-infocus-gJH8AqpiSEU-unsplash.jpg",
            "examples/jeremy-wallace-_XjW3oN8UOE-unsplash.jpg",
            "examples/karina-vorozheeva-rW-I87aPY5Y-unsplash.jpg",
            "examples/karographix-photography-hIaOPjYCEj4-unsplash.jpg",
            "examples/melissa-walker-horn-gtDYwUIr9Vg-unsplash.jpg",
            "examples/ryoji-iwata-X53e51WfjlE-unsplash.jpg",
            "examples/tadeusz-lakota-jggQZkITXng-unsplash.jpg",
        ],
        inputs=[input_image],
        outputs=output_slider,
        fn=process,
        cache_examples=True,
        cache_mode="lazy",
        run_on_click=False,
    )

demo.launch(share=False, ssr_mode=False)