import os
import torch
import datasets
import diffusers
from .hyperparams import MOONDREAM_REVISION

auth_token = os.getenv("HF_ACCESS_TOKEN")

tokenizer = transformers.AutoTokenizer.from_pretrained("vikhyatk/moondream2")
moondream = transformers.AutoModelForCausalLM.from_pretrained(
    "vikhyatk/moondream2",
    revision=MOONDREAM_REVISION,
    trust_remote_code=True,
    attn_implementation="flash_attention_2",
    torch_dtype=torch.float16,
).to("cuda")

def collate(batch):
    images = []
    questions = []

    for sample in batch:
        images.append(sample["image"])
        questions.append("Describe this image.")

    return images, questions

flickr_dataset = datasets.load_dataset("nlphuji/flickr30k", split="test", streaming=True)\
                         .select_columns(["image"])\
                         .take(1)

wiki_art_dataset = datasets.load_dataset("huggan/wikiart", split="train", streaming=True)\
                           .select_columns(["image"])\
                           .take(1)

anime_dataset = datasets.load_dataset("animelover/danbooru2022", "1-full", trust_remote_code=True, split="train", streaming=True)\
                        .select_columns(["image"])\
                        .take(1)

coco_dataset = datasets.load_dataset("detection-datasets/coco", split="train", streaming=True)\
                       .select_columns(["image"])\
                       .take(1)

movie_poster_dataset = datasets.load_dataset("skvarre/movie_posters-100k", split="train", streaming=True)\
                               .select_columns(["image"])\
                               .take(1)

cars_dataset = datasets.load_dataset("tanganke/stanford_cars", split="train", streaming=True)\
                       .select_columns(["image"])\
                       .take(1)

website_dataset = datasets.load_dataset("silatus/1k_Website_Screenshots_and_Metadata", split="train", streaming=True)\
                          .select_columns(["image"])\
                          .take(1)

movie_scene_dataset = datasets.load_dataset("unography/movie-scenes-resized-captioned", split="train", streaming=True)\
                              .select_columns(["image"])\
                              .take(1)

ds = datasets.concatenate_datasets([
    flickr_dataset,
    wiki_art_dataset,
    anime_dataset,
    coco_dataset,
    movie_poster_dataset,
    cars_dataset,
    website_dataset,
    movie_scene_dataset,
])

data_loader = torch.utils.data.DataLoader(
    ds,
    batch_size=8,
    collate_fn=collate
)

captions = []
for batch in data_loader:
    images, questions = batch
    answers = moondream.batch_answer(
        images=images,
        prompts=questions,
        tokenizer=tokenizer
    )

    for ans in answers:
        print(ans)
        print()

    captions.extend(answers)

ds = ds.add_column("caption", captions)

del moondream

pipe = diffusers.StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-large",
    torch_dtype=torch.bfloat16,
    token=auth_token,
).to("cuda")

image = pipe(
    "A capybara holding a sign that reads Hello World",
    num_inference_steps=28,
    guidance_scale=3.5,
).images[0]
image.save("capybara.png")