import spaces

import json
import yaml
import os
from PIL import Image

import numpy as np

import torch
import torchvision.utils as vutils


import gradio as gr

from huggingface_hub import hf_hub_download

from model.pipeline import UJiTModel, JiTConfig
from model.config import ClassContextConfig


MODEL_REPO = os.environ.get("MODEL_REPO", "p1atdev/JiT-AnimeFace-experiment")
MODEL_PATH = os.environ.get(
    "MODEL_PATH",
    "ujit-b512-p32-cls-derf/jit-anime_00020e_311840s.safetensors",
)
LABEL2ID_PATH = os.environ.get("LABEL2ID_PATH", "ujit-b512-p32-cls-derf/label2id.json")
CONFIG_PATH = os.environ.get(
    "CONFIG_PATH",
    "ujit-b512-p32-cls-derf/config.yml",
)

DEVICE = (
    torch.device("cuda")
    if torch.cuda.is_available()
    else torch.device("mps")
    if torch.backends.mps.is_available()
    else torch.device("cpu")
)
DTYPE = torch.bfloat16 if DEVICE.type in ["cuda"] else torch.float16
MAX_TOKEN_LENGTH = 32

MAX_IMAGE_SIZE = 1024
MIN_IMAGE_SIZE = 256

COMMON_NEGATIVE_PROMPT = (
    "retro artstyle, speech bubble, doujinshi, comic, title, cover, "
    "logo, watermark, signature, bad anatomy, sketch, abstract, unfinished, multiple views, "
)

model_map: dict[str, UJiTModel] = {}  # {model_path: model}
label2id_map: dict[str, dict] = {}  # {label2id_path: label2id}


def images_to_tensor(
    images: list[Image.Image],
    dtype: torch.dtype,
    device: torch.device,
) -> torch.Tensor:
    # 0~255 -> -1~1
    return torch.stack(
        [
            torch.tensor(np.array(image), dtype=dtype, device=device).permute(2, 0, 1)
            / 127.5
            - 1.0
            for image in images
        ]
    )


@torch.no_grad()
def images_to_grid_image(
    images: list[Image.Image] | torch.Tensor,
    padding: int = 2,
) -> Image.Image:
    if isinstance(images, list):
        tensor_images = images_to_tensor(
            images,
            dtype=torch.float16,
            device=torch.device("cpu"),
        )

    grid = vutils.make_grid(
        tensor_images,
        nrow=int(len(tensor_images) ** 0.5),
        padding=padding,
        normalize=True,
    )

    # TensorをPIL画像に変換して保存
    # (C, H, W) -> (H, W, C) への変換なども自動化できる
    image = Image.fromarray(
        grid.mul(255)
        .add_(0.5)
        .clamp_(0, 255)
        .permute(1, 2, 0)
        .to("cpu", torch.uint8)
        .numpy()
    )

    return image


def get_file_path(repo: str, path: str) -> str:
    """Hugging Face Hub からファイルを取得"""

    return hf_hub_download(repo, path)


def load_label2id(label2id_path: str) -> dict:
    """label2id.json を読み込む"""
    with open(label2id_path, "r") as f:
        return json.load(f)


def load_config(config_path: str) -> JiTConfig:
    """設定ファイルを読み込む"""
    with open(config_path, "r") as f:
        if config_path.endswith(".json"):
            config_dict = json.load(f)
        elif config_path.endswith((".yaml", ".yml")):
            config_dict = yaml.safe_load(f)
        else:
            raise ValueError("Unsupported config file format. Use .json or .yaml/.yml")

    return JiTConfig.model_validate(config_dict)


def load_model(
    model_path: str,
    label2id_path: str,
    config_path: str,
    device: torch.device,
    dtype: torch.dtype = DTYPE,
) -> tuple[UJiTModel, dict]:
    """モデルを読み込む"""

    if model_path in model_map:  # use cache
        model = model_map[model_path]
        label2id = label2id_map[label2id_path]
        return model, label2id

    config = load_config(get_file_path(MODEL_REPO, config_path))
    if isinstance(config.context_encoder, ClassContextConfig):
        config.context_encoder.label2id_map_path = get_file_path(
            MODEL_REPO, label2id_path
        )

    model = UJiTModel.from_pretrained(
        config=config,
        checkpoint_path=get_file_path(MODEL_REPO, model_path),
    )
    model.eval()
    model.requires_grad_(False)
    model.to(device=device, dtype=dtype)
    model_map[model_path] = model  # cache

    label2id = load_label2id(get_file_path(MODEL_REPO, label2id_path))
    label2id_map[label2id_path] = label2id  # cache

    return model, label2id


@spaces.GPU(duration=15)
def generate_images(
    prompt: str,
    negative_prompt: str,
    num_steps: int,
    cfg_scale: float,
    batch_size: int,
    height: int,
    width: int,
    seed: int,
    #
    model_path: str = MODEL_PATH,
    label2id_path: str = LABEL2ID_PATH,
    config_path: str = CONFIG_PATH,
    progress=gr.Progress(track_tqdm=True),
):
    model, _label2id = load_model(
        model_path=model_path,
        label2id_path=label2id_path,
        config_path=config_path,
        device=DEVICE,
        dtype=DTYPE,
    )

    with torch.inference_mode(), torch.autocast(device_type=DEVICE.type, dtype=DTYPE):
        images = model.generate(
            prompt=[prompt] * batch_size,
            negative_prompt=negative_prompt,
            num_inference_steps=num_steps,
            cfg_scale=cfg_scale,
            height=height,
            width=width,
            max_token_length=MAX_TOKEN_LENGTH,
            cfg_time_range=[0.1, 1.0],
            seed=seed if seed >= 0 else None,
            device=DEVICE,
            execution_dtype=DTYPE,
        )

    return [images, images_to_grid_image(images, padding=4)]


LABEL2ID_URL = f"https://huggingface.co/{MODEL_REPO}/blob/main/{LABEL2ID_PATH}"


def demo():
    with gr.Blocks() as ui:
        gr.Markdown(f"""
# JiT-Anime Demo
Pixel-space x-prediction flow-matching 380M parameter model for anime image generation, trained from scratch.

- See full supported tags: [label2id.json]({LABEL2ID_URL}). 対応しているタグ一覧は [こちら]({LABEL2ID_URL}) から確認できます。ここに載っていないタグは反応しません。
- Current model: [{MODEL_PATH}](https://huggingface.co/{MODEL_REPO}/blob/main/{MODEL_PATH})
""")

        with gr.Row():
            with gr.Column():
                prompt = gr.TextArea(
                    label="Prompt",
                    info=f"Comma-separated tags. Not all of danbooru tags are supported. See [the full supported tags]({LABEL2ID_URL}). カンマで区切ってください。",
                    value="1girl, solo, cowboy shot, looking at viewer, blue hair, short hair, colored inner hair, hair intakes, cat ears, animal ears, red eyes, white background, collared shirt, long sleeves",
                    placeholder="e.g.: general, 1girl, solo, portrait, looking at viewer",
                )
                negative_prompt = gr.TextArea(
                    label="Negative Prompt",
                    info="Comma-separated negative tags to avoid in generation. カンマで区切ってください。",
                    value=COMMON_NEGATIVE_PROMPT,
                    lines=2,
                    placeholder="e.g.: retro artstyle, 1990s (style), sketch",
                )
                num_steps = gr.Slider(
                    minimum=1,
                    maximum=100,
                    value=32,
                    step=1,
                    label="Number of Steps",
                    info="Recommended: more than 20 steps for better quality.",
                )
                cfg_scale = gr.Slider(
                    minimum=1.0,
                    maximum=15.0,
                    value=7.5,
                    step=0.25,
                    label="CFG Scale",
                    info="Recommended: more than 2.0 for better adherence to the prompt.",
                )
                batch_size = gr.Slider(
                    minimum=1,
                    maximum=36,
                    value=16,
                    step=1,
                    label="Batch Size",
                    info="Number of images to generate in one batch.",
                )
                with gr.Row():
                    height = gr.Slider(
                        minimum=MIN_IMAGE_SIZE,
                        maximum=MAX_IMAGE_SIZE,
                        value=640,
                        step=32,
                        label="Image Height",
                    )
                    width = gr.Slider(
                        minimum=MIN_IMAGE_SIZE,
                        maximum=MAX_IMAGE_SIZE,
                        value=448,
                        step=32,
                        label="Image Width",
                    )
                seed = gr.Number(
                    value=-1,
                    label="Seed (-1 for random)",
                )

            with gr.Column(scale=2):
                generate_button = gr.Button("Generate Images", variant="primary")
                output_gallery = gr.Gallery(
                    label="Generated Images",
                    columns=5,
                    height="768px",
                    object_fit="contain",
                    preview=False,
                    show_label=False,
                )

                with gr.Accordion("Grid", open=False):
                    grid_image = gr.Image(
                        label="Grid Image",
                        type="pil",
                        show_label=False,
                    )

        gr.Examples(
            examples=[
                [
                    "1girl, solo, cowboy shot, looking at viewer, blue hair, short hair, colored inner hair, hair intakes, cat ears, animal ears, red eyes, white background, collared shirt, long sleeves",
                    COMMON_NEGATIVE_PROMPT,
                ],
                [
                    "1girl, solo, original, portrait, upper body, looking at viewer, long hair, blue ribbon, hair ornament, hairclip, depth of field, head tilt, collared shirt, white shirt, simple background",
                    COMMON_NEGATIVE_PROMPT,
                ],
                [
                    "1girl, aqua eyes, baseball cap, blonde hair, closed mouth, earrings, green background, hat, jewelry, looking at viewer, shirt, short hair, simple background, solo, portrait, yellow shirt",
                    COMMON_NEGATIVE_PROMPT,
                ],
            ],
            inputs=[prompt, negative_prompt],
            label="Examples",
            examples_per_page=20,
        )

        gr.on(
            triggers=[generate_button.click, prompt.submit],
            fn=generate_images,
            inputs=[
                prompt,
                negative_prompt,
                num_steps,
                cfg_scale,
                batch_size,
                height,
                width,
                seed,
            ],
            outputs=[output_gallery, grid_image],
        )

    return ui


if __name__ == "__main__":
    load_model(
        model_path=MODEL_PATH,
        label2id_path=LABEL2ID_PATH,
        config_path=CONFIG_PATH,
        device=DEVICE,
        dtype=DTYPE,
    )

    demo().launch()