#!/usr/bin/env python

from __future__ import annotations

import io
import pathlib
import tarfile

import gradio as gr
import numpy as np
import PIL.Image
from huggingface_hub import hf_hub_download

TITLE = "TADNE (This Anime Does Not Exist) Image Viewer"
DESCRIPTION = """The original TADNE site is https://thisanimedoesnotexist.ai/.

You can view images generated by the TADNE model with seed 0-99999.
The original images are 512x512 in size, but they are resized to 128x128 here.

Expected execution time on Hugging Face Spaces: 4s

Related Apps:
- [TADNE](https://huggingface.co/spaces/hysts/TADNE)
- [TADNE Image Viewer](https://huggingface.co/spaces/hysts/TADNE-image-viewer)
- [TADNE Image Selector](https://huggingface.co/spaces/hysts/TADNE-image-selector)
- [TADNE Interpolation](https://huggingface.co/spaces/hysts/TADNE-interpolation)
- [TADNE Image Search with DeepDanbooru](https://huggingface.co/spaces/hysts/TADNE-image-search-with-DeepDanbooru)
"""


image_size = 128
min_seed = 0
max_seed = 99999
dirname = "0-99999"
tarball_path = hf_hub_download("hysts/TADNE-sample-images", f"{image_size}/{dirname}.tar", repo_type="dataset")


def run(
    start_seed: int,
    nrows: int,
    ncols: int,
) -> np.ndarray:
    start_seed = int(start_seed)
    num = nrows * ncols
    images = []
    dummy = np.ones((image_size, image_size, 3), dtype=np.uint8) * 255
    with tarfile.TarFile(tarball_path) as tar_file:
        for seed in range(start_seed, start_seed + num):
            if not min_seed <= seed <= max_seed:
                images.append(dummy)
                continue
            member = tar_file.getmember(f"{dirname}/{seed:07d}.jpg")
            with tar_file.extractfile(member) as f:  # type: ignore
                data = io.BytesIO(f.read())
            image = PIL.Image.open(data)
            image = np.asarray(image)
            images.append(image)
    res = (
        np.asarray(images)
        .reshape(nrows, ncols, image_size, image_size, 3)
        .transpose(0, 2, 1, 3, 4)
        .reshape(nrows * image_size, ncols * image_size, 3)
    )
    return res


demo = gr.Interface(
    fn=run,
    inputs=[
        gr.Number(label="Start Seed", value=0),
        gr.Slider(label="Number of Rows", minimum=1, maximum=10, step=1, value=2),
        gr.Slider(label="Number of Columns", minimum=1, maximum=10, step=1, value=5),
    ],
    outputs=gr.Image(label="Output"),
    title=TITLE,
    description=DESCRIPTION,
)


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