import json
import os
import zipfile
from io import BytesIO
from tempfile import NamedTemporaryFile
import tempfile

import gradio as gr
import pandas as pd
from PIL import Image
import safetensors.torch
import spaces
import timm
from timm.models import VisionTransformer
import torch
from torchvision.transforms import transforms
from torchvision.transforms import InterpolationMode
import torchvision.transforms.functional as TF
from torch.utils.data import Dataset, DataLoader


torch.set_grad_enabled(False)

class Fit(torch.nn.Module):
    def __init__(
        self,
        bounds: tuple[int, int] | int,
        interpolation = InterpolationMode.LANCZOS,
        grow: bool = True,
        pad: float | None = None
    ):
        super().__init__()

        self.bounds = (bounds, bounds) if isinstance(bounds, int) else bounds
        self.interpolation = interpolation
        self.grow = grow
        self.pad = pad

    def forward(self, img: Image) -> Image:
        wimg, himg = img.size
        hbound, wbound = self.bounds

        hscale = hbound / himg
        wscale = wbound / wimg

        if not self.grow:
            hscale = min(hscale, 1.0)
            wscale = min(wscale, 1.0)

        scale = min(hscale, wscale)
        if scale == 1.0:
            return img

        hnew = min(round(himg * scale), hbound)
        wnew = min(round(wimg * scale), wbound)

        img = TF.resize(img, (hnew, wnew), self.interpolation)

        if self.pad is None:
            return img

        hpad = hbound - hnew
        wpad = wbound - wnew

        tpad = hpad // 2
        bpad = hpad - tpad

        lpad = wpad // 2
        rpad = wpad - lpad

        return TF.pad(img, (lpad, tpad, rpad, bpad), self.pad)

    def __repr__(self) -> str:
        return (
            f"{self.__class__.__name__}(" +
            f"bounds={self.bounds}, " +
            f"interpolation={self.interpolation.value}, " +
            f"grow={self.grow}, " +
            f"pad={self.pad})"
        )

class CompositeAlpha(torch.nn.Module):
    def __init__(
        self,
        background: tuple[float, float, float] | float,
    ):
        super().__init__()

        self.background = (background, background, background) if isinstance(background, float) else background
        self.background = torch.tensor(self.background).unsqueeze(1).unsqueeze(2)

    def forward(self, img: torch.Tensor) -> torch.Tensor:
        if img.shape[-3] == 3:
            return img

        alpha = img[..., 3, None, :, :]

        img[..., :3, :, :] *= alpha

        background = self.background.expand(-1, img.shape[-2], img.shape[-1])
        if background.ndim == 1:
            background = background[:, None, None]
        elif background.ndim == 2:
            background = background[None, :, :]

        img[..., :3, :, :] += (1.0 - alpha) * background
        return img[..., :3, :, :]

    def __repr__(self) -> str:
        return (
            f"{self.__class__.__name__}(" +
            f"background={self.background})"
        )

transform = transforms.Compose([
    Fit((384, 384)),
    transforms.ToTensor(),
    CompositeAlpha(0.5),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
    transforms.CenterCrop((384, 384)),
])

model = timm.create_model(
    "vit_so400m_patch14_siglip_384.webli",
    pretrained=False,
    num_classes=9083,
) # type: VisionTransformer

safetensors.torch.load_model(model, "JTP_PILOT-e4-vit_so400m_patch14_siglip_384.safetensors")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

with open("tagger_tags.json", "r") as file:
    tags = json.load(file) # type: dict
allowed_tags = list(tags.keys())

for idx, tag in enumerate(allowed_tags):
    allowed_tags[idx] = tag.replace("_", " ")

sorted_tag_score = {}

@spaces.GPU(duration=9)
def run_classifier(image, threshold):
    global sorted_tag_score
    img = image.convert('RGB')
    tensor = transform(img).unsqueeze(0).to(device)
    with torch.no_grad():
        logits = model(tensor)
        probabilities = torch.nn.functional.sigmoid(logits[0])
        indices = torch.topk(probabilities, 250).indices
        values = probabilities[indices]

    tag_score = dict()
    for i in range(indices.size(0)):
        tag_score[allowed_tags[indices[i]]] = values[i].item()
    sorted_tag_score = dict(sorted(tag_score.items(), key=lambda item: item[1], reverse=True))

    return create_tags(threshold)

def create_tags(threshold):
    global sorted_tag_score
    filtered_tag_score = {key: value for key, value in sorted_tag_score.items() if value > threshold}
    text_no_impl = ", ".join(filtered_tag_score.keys())
    return text_no_impl, filtered_tag_score

    
class ImageDataset(Dataset):
    def __init__(self, image_files, transform):
        self.image_files = image_files
        self.transform = transform
    
    def __len__(self):
        return len(self.image_files)
    
    def __getitem__(self, idx):
        img_path = self.image_files[idx]
        img = Image.open(img_path).convert('RGB')
        return self.transform(img), os.path.basename(img_path)
    
@spaces.GPU(duration=299)
def process_images(images, threshold):
    dataset = ImageDataset(images, transform)

    dataloader = DataLoader(dataset, batch_size=64, num_workers=0, pin_memory=True, drop_last=False)
    
    all_results = []

    with torch.no_grad():
        for batch, filenames in dataloader:
            
            batch = batch.to(device) 
            with torch.no_grad():
                logits = model(batch)
                probabilities = torch.nn.functional.sigmoid(logits)
            
                for i, prob in enumerate(probabilities):
                    indices = torch.where(prob > threshold)[0]
                    values = prob[indices]
                    
                    temp = []
                    tag_score = dict()
                    for j in range(indices.size(0)):
                        temp.append([allowed_tags[indices[j]], values[j].item()])
                        tag_score[allowed_tags[indices[j]]] = values[j].item()
                    
                    tags = ", ".join([t[0] for t in temp])
                    all_results.append((filenames[i], tags, tag_score))

    return all_results

def is_valid_image(file_path):
    try:
        with Image.open(file_path) as img:
            img.verify()
        return True
    except:
        return False

def process_zip(zip_file, threshold):
    if zip_file is None:
        return None, None

    with tempfile.TemporaryDirectory() as temp_dir:
        with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
            zip_ref.extractall(temp_dir)
        
        all_files = [os.path.join(temp_dir, f) for f in os.listdir(temp_dir)]
        image_files = [f for f in all_files if is_valid_image(f)]
        results = process_images(image_files, threshold)
    
        temp_file = NamedTemporaryFile(delete=False, suffix=".zip")
        with zipfile.ZipFile(temp_file, "w") as zip_ref:
            for image_name, text_no_impl, _ in results:
                with zip_ref.open(''.join(image_name.split('.')[:-1]) + ".txt", 'w') as file:
                    file.write(text_no_impl.encode())
        temp_file.seek(0)
        df = pd.DataFrame([(os.path.basename(f), t) for f, t, _ in results], columns=['Image', 'Tags'])

    return temp_file.name, df

@spaces.GPU(duration=120)  # Reduced GPU duration for less wait time...
def process_images_light(images, threshold):
    dataset = ImageDataset(images, transform)

    dataloader = DataLoader(dataset, batch_size=32, num_workers=0, pin_memory=True, drop_last=False)
    
    all_results = []

    with torch.no_grad():
        for batch, filenames in dataloader:
            
            batch = batch.to(device) 
            with torch.no_grad():
                logits = model(batch)
                probabilities = torch.nn.functional.sigmoid(logits)
            
                for i, prob in enumerate(probabilities):
                    indices = torch.where(prob > threshold)[0]
                    values = prob[indices]
                    
                    temp = []
                    tag_score = dict()
                    for j in range(indices.size(0)):
                        temp.append([allowed_tags[indices[j]], values[j].item()])
                        tag_score[allowed_tags[indices[j]]] = values[j].item()
                    
                    tags = ", ".join([t[0] for t in temp])
                    all_results.append((filenames[i], tags, tag_score))

    return all_results

def process_zip_light(zip_file, threshold):
    if zip_file is None:
        return None, None

    with tempfile.TemporaryDirectory() as temp_dir:
        with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
            zip_ref.extractall(temp_dir)
        
        all_files = [os.path.join(temp_dir, f) for f in os.listdir(temp_dir)]
        image_files = [f for f in all_files if is_valid_image(f)]
        results = process_images_light(image_files, threshold)
    
        temp_file = NamedTemporaryFile(delete=False, suffix=".zip")
        with zipfile.ZipFile(temp_file, "w") as zip_ref:
            for image_name, text_no_impl, _ in results:
                with zip_ref.open(''.join(image_name.split('.')[:-1]) + ".txt", 'w') as file:
                    file.write(text_no_impl.encode())
        temp_file.seek(0)
        df = pd.DataFrame([(os.path.basename(f), t) for f, t, _ in results], columns=['Image', 'Tags'])

    return temp_file.name, df

with gr.Blocks(css=".output-class { display: none; }") as demo:
    gr.Markdown("""
    ## Joint Tagger Project: PILOT Demo
    This tagger is designed for use on furry images (though may very well work on out-of-distribution images, potentially with funny results).  A threshold of 0.2 is recommended.  Lower thresholds often turn up more valid tags, but can also result in some amount of hallucinated tags.

    This tagger is the result of joint efforts between members of the RedRocket team.  Special thanks to Minotoro at frosting.ai for providing the compute power for this project.
    
    Usage Note for batch tagging:
    
    the normal version is limited to 300s and uses batch size 64
    
    the light version is limited to 120s with batch size 32
    
    if your image count is low use the light version for lower gpu wait time (most of the time you instantly get a gpu anyway)
    """)
    
    with gr.Tabs():
        with gr.TabItem("Single Image"):
            with gr.Row():
                with gr.Column():
                    image_input = gr.Image(label="Source", sources=['upload'], type='pil', height=512, show_label=False)
                    threshold_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")
                with gr.Column():
                    tag_string = gr.Textbox(label="Tag String")
                    label_box = gr.Label(label="Tag Predictions", num_top_classes=250, show_label=False)

            image_input.upload(
                fn=run_classifier,
                inputs=[image_input, threshold_slider],
                outputs=[tag_string, label_box]
            )

            threshold_slider.input(
                fn=create_tags,
                inputs=[threshold_slider],
                outputs=[tag_string, label_box]
            )
        
        with gr.TabItem("Multiple Images"):
            with gr.Row():
                with gr.Column():
                    zip_input = gr.File(label="Upload ZIP file", file_types=['.zip'])
                    multi_threshold_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")
                    process_button = gr.Button("Process Images")
                with gr.Column():
                    zip_output = gr.File(label="Download Tagged Text Files (ZIP)")
                    dataframe_output = gr.Dataframe(label="Image Tags Summary")
            
            process_button.click(
                fn=process_zip,
                inputs=[zip_input, multi_threshold_slider],
                outputs=[zip_output, dataframe_output]
            )
        with gr.TabItem("Multiple Images (Light)"):
            with gr.Row():
                with gr.Column():
                    zip_input_light = gr.File(label="Upload ZIP file", file_types=['.zip'])
                    multi_threshold_slider_light = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")
                    process_button_light = gr.Button("Process Images (Light)")
                with gr.Column():
                    zip_output_light = gr.File(label="Download Tagged Text Files (ZIP)")
                    dataframe_output_light = gr.Dataframe(label="Image Tags Summary")
            
            process_button_light.click(
                fn=process_zip_light,
                inputs=[zip_input_light, multi_threshold_slider_light],
                outputs=[zip_output_light, dataframe_output_light]
            )

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