from transformers import CLIPImageProcessor, AutoModel
import torch
import json
import torch.nn as nn
from PIL import Image
import gradio as gr
import os
import faiss
import time
import requests
from huggingface_hub import login, snapshot_download

TITLE = "Danbooru Tagger"
DESCRIPTION = """
## Dataset
- Source: Cleaned Danbooru

## Metrics
- Validation Split: 10% of Dataset
- Validation Results:

### General
| Metric          | Value       |
|-----------------|-------------|
| Macro F1        | 0.4678      |
| Macro Precision | 0.4605      |
| Macro Recall    | 0.5229      |
| Micro F1        | 0.6661      |
| Micro Precision | 0.6049      |
| Micro Recall    | 0.7411      |

### Character
| Metric          | Value       |
|-----------------|-------------|
| Macro F1        | 0.8925      |
| Macro Precision | 0.9099      |
| Macro Recall    | 0.8935      |
| Micro F1        | 0.9232      |
| Micro Precision | 0.9264      |
| Micro Recall    | 0.9199      |

### Artist
| Metric          | Value       |
|-----------------|-------------|
| Macro F1        | 0.7904      |
| Macro Precision | 0.8286      |
| Macro Recall    | 0.7904      |
| Micro F1        | 0.5989      |
| Micro Precision | 0.5975      |
| Micro Recall    | 0.6004      |
"""

kaomojis = [
    "0_0",
    "(o)_(o)",
    "+_+",
    "+_-",
    "._.",
    "<o>_<o>",
    "<|>_<|>",
    "=_=",
    ">_<",
    "3_3",
    "6_9",
    ">_o",
    "@_@",
    "^_^",
    "o_o",
    "u_u",
    "x_x",
    "|_|",
    "||_||",
]

device = torch.device('cpu')
dtype = torch.float32

hf_token = os.getenv("HF_TOKEN")
if hf_token:
    login(token=hf_token)
else:
    raise ValueError("environment variable HF_TOKEN not found.")

repo = snapshot_download('Johnny-Z/vit-e4')
model = AutoModel.from_pretrained(repo, dtype=dtype, trust_remote_code=True, device_map=device)

index_dir = snapshot_download('Johnny-Z/dan_index', repo_type='dataset')

processor = CLIPImageProcessor.from_pretrained(repo)

class MultiheadAttentionPoolingHead(nn.Module):
    def __init__(self, input_size):
        super().__init__()

        self.map_probe = nn.Parameter(torch.randn(1, 1, input_size))
        self.map_layernorm0 = nn.LayerNorm(input_size, eps=1e-08)
        self.map_attention = torch.nn.MultiheadAttention(input_size, input_size // 64, batch_first=True)
        self.map_layernorm1 = nn.LayerNorm(input_size, eps=1e-08)
        self.map_ffn = nn.Sequential(
            nn.Linear(input_size, input_size * 4),
            nn.SiLU(),
            nn.Linear(input_size * 4, input_size)
        )

    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
        batch_size = hidden_state.shape[0]
        probe = self.map_probe.repeat(batch_size, 1, 1)

        hidden_state = self.map_layernorm0(hidden_state)
        hidden_state = self.map_attention(probe, hidden_state, hidden_state)[0]
        hidden_state = self.map_layernorm1(hidden_state)

        residual = hidden_state
        hidden_state = residual + self.map_ffn(hidden_state)
        return hidden_state[:, 0]

class MLP(nn.Module):
    def __init__(self, input_size, class_num):
        super().__init__()
        self.mlp_layer0 = nn.Sequential(
            nn.LayerNorm(input_size, eps=1e-08),
            nn.Linear(input_size, input_size // 2),
            nn.SiLU()
        )
        self.mlp_layer1 = nn.Linear(input_size // 2, class_num)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.mlp_layer0(x)
        x = self.mlp_layer1(x)
        x = self.sigmoid(x)
        return x

class MLP_Retrieval(nn.Module):
    def __init__(self, input_size, class_num):
        super().__init__()
        self.mlp_layer0 = nn.Sequential(
            nn.Linear(input_size, input_size // 2),
            nn.SiLU()
        )
        self.mlp_layer1 = nn.Linear(input_size // 2, class_num)

    def forward(self, x):
        x = self.mlp_layer0(x)
        x = self.mlp_layer1(x)
        x1, x2 = x[:, :15], x[:, 15:]
        x1 = torch.softmax(x1, dim=1)
        x2 = torch.softmax(x2, dim=1)
        x = torch.cat([x1, x2], dim=1)

        return x
    
class MLP_R(nn.Module):
    def __init__(self, input_size):
        super().__init__()
        self.mlp_layer0 = nn.Sequential(
            nn.Linear(input_size, 256),
        )

    def forward(self, x):
        x = self.mlp_layer0(x)
        return x

with open(os.path.join(repo, 'general_tag_dict.json'), 'r', encoding='utf-8') as f:
    general_dict = json.load(f)

with open(os.path.join(repo, 'character_tag_dict.json'), 'r', encoding='utf-8') as f:
    character_dict = json.load(f)
    
with open(os.path.join(repo, 'artist_tag_dict.json'), 'r', encoding='utf-8') as f:
    artist_dict = json.load(f)

with open(os.path.join(repo, 'implications_list.json'), 'r', encoding='utf-8') as f:
    implications_list = json.load(f)

with open(os.path.join(repo, 'artist_threshold.json'), 'r', encoding='utf-8') as f:
    artist_thresholds = json.load(f)

with open(os.path.join(repo, 'character_threshold.json'), 'r', encoding='utf-8') as f:
    character_thresholds = json.load(f)

with open(os.path.join(repo, 'general_threshold.json'), 'r', encoding='utf-8') as f:
    general_thresholds = json.load(f)

model_map = MultiheadAttentionPoolingHead(2048)
model_map.load_state_dict(torch.load(os.path.join(repo, "map_head.pth"), map_location=device, weights_only=True))
model_map.to(device).to(dtype).eval()

general_class = 9775
mlp_general = MLP(2048, general_class)
mlp_general.load_state_dict(torch.load(os.path.join(repo, "cls_predictor_general.pth"), map_location=device, weights_only=True))
mlp_general.to(device).to(dtype).eval()

character_class = 7568
mlp_character = MLP(2048, character_class)
mlp_character.load_state_dict(torch.load(os.path.join(repo, "cls_predictor_character.pth"), map_location=device, weights_only=True))
mlp_character.to(device).to(dtype).eval()

artist_class = 13957
mlp_artist = MLP(2048, artist_class)
mlp_artist.load_state_dict(torch.load(os.path.join(repo, "cls_predictor_artist.pth"), map_location=device, weights_only=True))
mlp_artist.to(device).to(dtype).eval()

mlp_artist_retrieval = MLP_Retrieval(2048, artist_class)
mlp_artist_retrieval.load_state_dict(torch.load(os.path.join(repo, "cls_predictor_artist_retrieval.pth"), map_location=device, weights_only=True))
mlp_artist_retrieval.to(device).to(dtype).eval()

mlp_r = MLP_R(2048)
mlp_r.load_state_dict(torch.load(os.path.join(repo, "retrieval_head.pth"), map_location=device, weights_only=True))
mlp_r.to(device).to(dtype).eval()

def prediction_to_tag(prediction, tag_dict, class_num):
    prediction = prediction.view(class_num)
    predicted_ids = (prediction >= 0.2).nonzero(as_tuple=True)[0].cpu().numpy() + 1
    
    general = {}
    character = {}
    artist = {}
    date = {}
    rating = {}
    
    for tag, value in tag_dict.items():
        if value[2] in predicted_ids:
            tag_value = round(prediction[value[2] - 1].item(), 6)
            if value[1] == "general" and tag_value >= general_thresholds.get(tag, {}).get("Threshold", 0.75):
                general[tag] = tag_value
            elif value[1] == "character" and tag_value >= character_thresholds.get(tag, {}).get("Threshold", 0.75):
                character[tag] = tag_value
            elif value[1] == "artist" and tag_value >= artist_thresholds.get(tag, {}).get("Threshold", 0.75):
                artist[tag] = tag_value
            elif value[1] == "rating":
                rating[tag] = tag_value
            elif value[1] == "date":
                date[tag] = tag_value

    general = dict(sorted(general.items(), key=lambda item: item[1], reverse=True))
    character = dict(sorted(character.items(), key=lambda item: item[1], reverse=True))
    artist = dict(sorted(artist.items(), key=lambda item: item[1], reverse=True))

    if date:
        date = {max(date, key=date.get): date[max(date, key=date.get)]}
    if rating:
        rating = {max(rating, key=rating.get): rating[max(rating, key=rating.get)]}
    
    return general, character, artist, date, rating

def prediction_to_retrieval(prediction, tag_dict, class_num, top_k):
    prediction = prediction.view(class_num)
    predicted_ids = (prediction>=0.005).nonzero(as_tuple=True)[0].cpu().numpy() + 1

    artist = {}
    date = {}
    
    for tag, value in tag_dict.items():
        if value[2] in predicted_ids:
            tag_value = round(prediction[value[2] - 1].item(), 6)
            if value[1] == "artist":
                artist[tag] = tag_value
            elif value[1] == "date":
                date[tag] = tag_value

    artist = dict(sorted(artist.items(), key=lambda item: item[1], reverse=True))
    artist = dict(list(artist.items())[:top_k])

    if date:
        date = {max(date, key=date.get): date[max(date, key=date.get)]}
    
    return artist, date

def load_id_map(id_map_path):
    with open(id_map_path, "r") as f:
        id_map = json.load(f)

    id_map = {int(k): int(v) for k, v in id_map.items()}

    inv_map = {v: k for k, v in id_map.items()}
    return id_map, inv_map

def search_index(query_vector, k=32, distance_threshold_min=0, distance_threshold_max=64, nprobe=4):
    global index_dir
    index_path = os.path.join(index_dir, 'danbooru_retrieval.index')
    id_map_path = os.path.join(index_dir, 'danbooru_retrieval_id_map.json')
    distance_threshold_min = distance_threshold_min**2
    distance_threshold_max = distance_threshold_max**2

    index = faiss.read_index(index_path)

    if nprobe is not None and hasattr(index, "nprobe"):
        index.nprobe = nprobe  
    _, inv_map = load_id_map(id_map_path)

    qv = query_vector.detach().to(torch.float32).cpu().numpy()

    distances, internal_ids = index.search(qv, k)  
    distances = distances[0]
    internal_ids = internal_ids[0]

    results = []
    for dist, internal_id in zip(distances, internal_ids):
        if internal_id == -1:
            continue  
        if dist < distance_threshold_min or dist > distance_threshold_max:
            continue
        original_id = inv_map.get(int(internal_id))
        if original_id is None:
            continue
        results.append({"original_id": original_id, "l2_distance": float(dist**0.5)})
    results.sort(key=lambda x: x["l2_distance"])

    return results

def fetch_retrieval_image_urls(retrieval_results, sleep_sec=0.25, timeout=4.0):
    pairs = []
    for item in retrieval_results:
        oid = item.get("original_id")
        if oid is None:
            continue
        api_url = f"https://danbooru.donmai.us/posts/{oid}.json"
        try:
            resp = requests.get(api_url, timeout=timeout)
            if resp.status_code != 200:

                time.sleep(sleep_sec)
                continue
            data = resp.json()
            url = data.get("large_file_url") or data.get("file_url") or data.get("preview_file_url")
            if not url:
                time.sleep(sleep_sec)
                continue

            if url.startswith("//"):
                url = "https:" + url
            elif url.startswith("/"):
                url = "https://danbooru.donmai.us" + url
            pairs.append((url, oid))
        except Exception:

            pass
        finally:

            time.sleep(sleep_sec)
    return pairs

def process_image(image, k, distance_threshold_min, distance_threshold_max):
    try:
        image = image.convert('RGBA')
        background = Image.new('RGBA', image.size, (255, 255, 255, 255))
        image = Image.alpha_composite(background, image).convert('RGB')

        image_inputs = processor(images=[image], return_tensors="pt").to(device).to(dtype)

    except (OSError, IOError) as e:
        print(f"Error opening image: {e}")
        return
    with torch.no_grad():
        embedding = model(image_inputs.pixel_values)

        embedding = model_map(embedding)

        embedding_r = mlp_r(embedding)

        retrieval_results = search_index(embedding_r, k, distance_threshold_min, distance_threshold_max)

        url_id_pairs = fetch_retrieval_image_urls(retrieval_results)

        retrieval_gallery_items = [(url, f"https://danbooru.donmai.us/posts/{oid}") for url, oid in url_id_pairs]

        general_prediction = mlp_general(embedding)
        general_ = prediction_to_tag(general_prediction, general_dict, general_class)
        general_tags = general_[0]
        rating = general_[4]

        character_prediction = mlp_character(embedding)
        character_ = prediction_to_tag(character_prediction, character_dict, character_class)
        character_tags = character_[1]

        artist_retrieval_prediction = mlp_artist_retrieval(embedding)
        artist_retrieval_ = prediction_to_retrieval(artist_retrieval_prediction, artist_dict, artist_class, 10)
        artist_tags = artist_retrieval_[0]
        date = artist_retrieval_[1]

    combined_tags = {**general_tags}

    tags_list = [tag for tag in combined_tags]
    remove_list = []
    for tag in tags_list:
        if tag in implications_list:
            for implication in implications_list[tag]:
                remove_list.append(implication)
    tags_list = [tag for tag in tags_list if tag not in remove_list]
    tags_list = [tag.replace("_", " ") if tag not in kaomojis else tag for tag in tags_list]

    tags_str = ", ".join(tags_list).replace("(", r"\(").replace(")", r"\)")

    return (
        tags_str,
        artist_tags,
        character_tags,
        general_tags,
        rating,
        date,
        retrieval_gallery_items,
    )

def main():
    with gr.Blocks(title=TITLE) as demo:
        with gr.Column():
            gr.Markdown(
                value=f"<h1 style='text-align: center; margin-bottom: 1rem'>{TITLE}</h1>"
            )
            with gr.Row():
                with gr.Column(variant="panel"):
                    submit = gr.Button(value="Submit", variant="primary", size="lg")
                    image = gr.Image(type="pil", image_mode="RGBA", label="Input")
                    k_slider = gr.Slider(1, 100, value=32, step=1, label="Top K Results")
                    distance_min_slider = gr.Slider(0, 128, value=0, step=1, label="Min Distance Threshold")
                    distance_max_slider = gr.Slider(0, 128, value=80, step=1, label="Max Distance Threshold")
                    with gr.Row():
                        clear = gr.ClearButton(
                            components=[
                                image,
                                k_slider,
                                distance_min_slider,
                                distance_max_slider,
                            ],
                            variant="secondary",
                            size="lg",
                        )
                    gr.Markdown(value=DESCRIPTION)
                with gr.Column(variant="panel"):
                    tags_str = gr.Textbox(label="Output", lines=4)
                    with gr.Row():
                        rating = gr.Label(label="Rating")
                        date = gr.Label(label="Year")
                    artist_tags = gr.Label(label="Artist")
                    character_tags = gr.Label(label="Character")
                    general_tags = gr.Label(label="General")
            with gr.Row():
                retrieval_gallery = gr.Gallery(
                    label="Retrieval Preview",
                    columns=5,
                )
            clear.add(
                [
                    tags_str,
                    artist_tags,
                    general_tags,
                    character_tags,
                    rating,
                    date,
                    retrieval_gallery,
                ]
            )

        submit.click(
            process_image,
            inputs=[image, k_slider, distance_min_slider, distance_max_slider],
            outputs=[
                tags_str,
                artist_tags,
                character_tags,
                general_tags,
                rating,
                date,
                retrieval_gallery,
            ],
        )

    demo.queue(max_size=10)
    demo.launch()

if __name__ == "__main__":
    main()