from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd
import timm
import torch
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import HfHubHTTPError
from PIL import Image
from simple_parsing import field, parse_known_args
from timm.data import create_transform, resolve_data_config
from torch import Tensor, nn
from torch.nn import functional as F
import os
import time
from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
from PIL import Image, UnidentifiedImageError
from pathlib import Path
from tqdm import tqdm

@dataclass
class ScriptOptions:
    image_folder: Path = "/workspace/ds/reddit"
    model: str = field(default="vit")
    gen_threshold: float = field(default=0.7)
    char_threshold: float = field(default=0.6)

dream_model = AutoModelForCausalLM.from_pretrained(
    "moondream/moondream-2b-2025-04-14-4bit",
    trust_remote_code=True,
    device_map={"": "cuda"}
)
dream_model.model.compile()

torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

new_path = '/workspace/wdv3-timm'
os.chdir(new_path)
print(os.getcwd())

MODEL_REPO_MAP = {
    "vit": "SmilingWolf/wd-vit-tagger-v3",
    "swinv2": "SmilingWolf/wd-swinv2-tagger-v3",
    "convnext": "SmilingWolf/wd-convnext-tagger-v3",
}

def pil_ensure_rgb(image: Image.Image) -> Image.Image:
    if image.mode not in ["RGB", "RGBA"]:
        image = image.convert("RGBA") if "transparency" in image.info else image.convert("RGB")
    if image.mode == "RGBA":
        canvas = Image.new("RGBA", image.size, (255, 255, 255))
        canvas.alpha_composite(image)
        image = canvas.convert("RGB")
    return image

def pil_pad_square(image: Image.Image) -> Image.Image:
    w, h = image.size
    px = max(image.size)
    canvas = Image.new("RGB", (px, px), (255, 255, 255))
    canvas.paste(image, ((px - w) // 2, (px - h) // 2))
    return canvas

@dataclass
class LabelData:
    names: list[str]
    rating: list[np.int64]
    general: list[np.int64]
    character: list[np.int64]

def load_labels_hf(
    repo_id: str,
    revision: Optional[str] = None,
    token: Optional[str] = None,
) -> LabelData:
    try:
        csv_path = hf_hub_download(
            repo_id=repo_id, filename="selected_tags.csv", revision=revision, token=token
        )
        csv_path = Path(csv_path).resolve()
    except HfHubHTTPError as e:
        raise FileNotFoundError(f"selected_tags.csv failed to download from {repo_id}") from e

    df: pd.DataFrame = pd.read_csv(csv_path, usecols=["name", "category"])
    tag_data = LabelData(
        names=df["name"].tolist(),
        rating=list(np.where(df["category"] == 9)[0]),
        general=list(np.where(df["category"] == 0)[0]),
        character=list(np.where(df["category"] == 4)[0]),
    )

    return tag_data

def get_tags(
    probs: Tensor,
    labels: LabelData,
    gen_threshold: float,
    char_threshold: float,
):
    probs = list(zip(labels.names, probs.numpy()))

    rating_labels = dict([probs[i] for i in labels.rating])
    rating_labels = dict(sorted(rating_labels.items(), key=lambda item: item[1], reverse=True))

    gen_labels = [probs[i] for i in labels.general]
    gen_labels = dict([x for x in gen_labels if x[1] > gen_threshold])
    gen_labels = dict(sorted(gen_labels.items(), key=lambda item: item[1], reverse=True))

    char_labels = [probs[i] for i in labels.character]
    char_labels = dict([x for x in char_labels if x[1] > char_threshold])
    char_labels = dict(sorted(char_labels.items(), key=lambda item: item[1], reverse=True))

    combined_names = [x for x in gen_labels]
    combined_names.extend([x for x in char_labels])

    caption = ", ".join(combined_names)

    taglist = caption.replace("_", " ").replace("(", "\(").replace(")", "\)")

    caption = caption.replace("_", " ")
    caption += ", rating_" + next(iter(sorted(rating_labels, key=rating_labels.get, reverse=True)), '')

    return caption, taglist, rating_labels, char_labels, gen_labels

def get_all_images(folder):
    count = 0

    for path in folder.rglob('*'):
        if path.suffix.lower() in ('.jpeg', '.jpg', '.png'):
            count += 1
            yield path

def main(opts: ScriptOptions):
    repo_id = MODEL_REPO_MAP.get(opts.model)
    image_folder = Path(opts.image_folder).resolve()
    if not image_folder.is_dir():
        raise NotADirectoryError(f"Image folder not found: {image_folder}")

    print(f"Loading model '{opts.model}' from '{repo_id}'...")
    model: nn.Module = timm.create_model("hf-hub:" + repo_id).eval()
    state_dict = timm.models.load_state_dict_from_hf(repo_id)
    model.load_state_dict(state_dict)

    print("Loading tag list...")
    labels: LabelData = load_labels_hf(repo_id=repo_id)

    print("Creating data transform...")
    transform = create_transform(**resolve_data_config(model.pretrained_cfg, model=model))

    image_paths = list(get_all_images(image_folder))
    num_images = len(image_paths)

    for image_path in tqdm(image_paths, desc="Processing images"):
        txt_file = image_path.with_suffix('.txt')
        if txt_file.exists():
            continue
        try:
            img_input: Image.Image = Image.open(image_path)
            img_input = pil_ensure_rgb(img_input)
            img_input = pil_pad_square(img_input)
            inputs: Tensor = transform(img_input).unsqueeze(0)
            inputs = inputs[:, [2, 1, 0]]

            with torch.inference_mode():
                mdream_capt = dream_model.caption(img_input, length="normal")["caption"]
                mdream_capt = mdream_capt.replace("The image depicts ", "").replace("The image presents ", "").replace("The image features ", "").replace("The image portrays ", "").replace("The image is ", "").strip()

                if torch_device.type != "cpu":
                    model = model.to(torch_device)
                    inputs = inputs.to(torch_device)
                outputs = model.forward(inputs)
                outputs = F.sigmoid(outputs)
                if torch_device.type != "cpu":
                    inputs = inputs.to("cpu")
                    outputs = outputs.to("cpu")
                    model = model.to("cpu")

            caption, taglist, ratings, character, general = get_tags(
                probs=outputs.squeeze(0),
                labels=labels,
                gen_threshold=opts.gen_threshold,
                char_threshold=opts.char_threshold,
            )

            clean_name = image_path.stem
            clean_name = ' '.join(word for word in clean_name.split() if not word.startswith(('1', '2', '3', '4', '5', '6', '7', '8', '9', '0')))

            tags_filename = str(image_path.with_suffix('.tag'))
            text_filename = str(image_path.with_suffix('.txt'))

            with open(tags_filename, 'w') as file_tag:
                file_tag.write(f"{caption}")
            with open(text_filename, 'w') as file_txt:
                file_txt.write(f"{mdream_capt} {caption}. {clean_name}")

        except (OSError, UnidentifiedImageError) as e:
            print(f"Error processing {image_path}: {str(e)}")
            continue

    print("Done!")

if __name__ == "__main__":
    opts, _ = parse_known_args(ScriptOptions)
    main(opts)