import torch
from torchvision import transforms

from config import Args
from pydantic import BaseModel, Field
from PIL import Image
from pipelines.pix2pix.pix2pix_turbo import Pix2Pix_Turbo
from pipelines.utils.canny_gpu import ScharrOperator

default_prompt = "close-up photo of the joker"
page_content = """
<h1 class="text-3xl font-bold">Real-Time pix2pix_turbo</h1>
<h3 class="text-xl font-bold">pix2pix turbo</h3>
<p class="text-sm">
    This demo showcases
    <a
    href="https://github.com/GaParmar/img2img-turbo"
    target="_blank"
    class="text-blue-500 underline hover:no-underline">One-Step Image Translation with Text-to-Image Models
    </a>
</p>
<p class="text-sm text-gray-500">
    Web app <a href="https://github.com/radames/Real-Time-Latent-Consistency-Model" target="_blank" class="text-blue-500 underline hover:no-underline">
    Real-Time Latent Consistency Models
    </a>
</p>
"""


class Pipeline:
    class Info(BaseModel):
        name: str = "img2img"
        title: str = "Image-to-Image SDXL"
        description: str = "Generates an image from a text prompt"
        input_mode: str = "image"
        page_content: str = page_content

    class InputParams(BaseModel):
        prompt: str = Field(
            default_prompt,
            title="Prompt",
            field="textarea",
            id="prompt",
        )

        width: int = Field(
            512, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
        )
        height: int = Field(
            512, min=2, max=15, title="Height", disabled=True, hide=True, id="height"
        )
        canny_low_threshold: float = Field(
            0.0,
            min=0,
            max=1.0,
            step=0.001,
            title="Canny Low Threshold",
            field="range",
            hide=True,
            id="canny_low_threshold",
        )
        canny_high_threshold: float = Field(
            1.0,
            min=0,
            max=1.0,
            step=0.001,
            title="Canny High Threshold",
            field="range",
            hide=True,
            id="canny_high_threshold",
        )
        debug_canny: bool = Field(
            False,
            title="Debug Canny",
            field="checkbox",
            hide=True,
            id="debug_canny",
        )

    def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):
        self.model = Pix2Pix_Turbo("edge_to_image")
        self.canny_torch = ScharrOperator(device=device)
        self.device = device
        self.last_time = 0.0

    def predict(self, params: "Pipeline.InputParams") -> Image.Image:
        canny_pil, canny_tensor = self.canny_torch(
            params.image,
            params.canny_low_threshold,
            params.canny_high_threshold,
            output_type="pil,tensor",
        )
        canny_tensor = torch.cat((canny_tensor, canny_tensor, canny_tensor), dim=1)
        output_image = self.model(
            canny_tensor,
            params.prompt,
        )
        output_pil = transforms.ToPILImage()(output_image[0].cpu() * 0.5 + 0.5)

        result_image = output_pil
        if params.debug_canny:
            # paste control_image on top of result_image
            w0, h0 = (200, 200)
            control_image = canny_pil.resize((w0, h0))
            w1, h1 = result_image.size
            result_image.paste(control_image, (w1 - w0, h1 - h0))
        return result_image