import colorsys
import os
import gradio as gr
import matplotlib.colors as mcolors
import numpy as np
import torch
from gradio.themes.utils import sizes
from matplotlib import pyplot as plt
from matplotlib.patches import Patch
from PIL import Image, ImageOps
from torchvision import transforms

'''!pip install diffusers

import spaces
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(...)
pipe.to('cuda')

@spaces.GPU
def generate(prompt):
    return pipe(prompt).images

gr.Interface(
    fn=generate,
    inputs=gr.Text(),
    outputs=gr.Gallery(),
).launch()

@spaces.GPU(duration=120)
def generate(prompt):
   return pipe(prompt).images'''

# ----------------- HELPER FUNCTIONS ----------------- #
os.chdir(os.path.dirname(os.path.abspath(__file__)))
ASSETS_DIR = os.path.join(os.path.dirname(__file__), "assets")
os.makedirs(ASSETS_DIR, exist_ok=True)

LABELS_TO_IDS = {
    "Background": 0,
    "Apparel": 1,
    "Face Neck": 2,
    "Hair": 3,
    "Left Foot": 4,
    "Left Hand": 5,
    "Left Lower Arm": 6,
    "Left Lower Leg": 7,
    "Left Shoe": 8,
    "Left Sock": 9,
    "Left Upper Arm": 10,
    "Left Upper Leg": 11,
    "Lower Clothing": 12,
    "Right Foot": 13,
    "Right Hand": 14,
    "Right Lower Arm": 15,
    "Right Lower Leg": 16,
    "Right Shoe": 17,
    "Right Sock": 18,
    "Right Upper Arm": 19,
    "Right Upper Leg": 20,
    "Torso": 21,
    "Upper Clothing": 22,
    "Lower Lip": 23,
    "Upper Lip": 24,
    "Lower Teeth": 25,
    "Upper Teeth": 26,
    "Tongue": 27,
}

def get_palette(num_cls):
    palette = [0] * (256 * 3)
    palette[0:3] = [0, 0, 0]

    for j in range(1, num_cls):
        hue = (j - 1) / (num_cls - 1)
        saturation = 1.0
        value = 1.0 if j % 2 == 0 else 0.5
        rgb = colorsys.hsv_to_rgb(hue, saturation, value)
        r, g, b = [int(x * 255) for x in rgb]
        palette[j * 3 : j * 3 + 3] = [r, g, b]

    return palette

def create_colormap(palette):
    colormap = np.array(palette).reshape(-1, 3) / 255.0
    return mcolors.ListedColormap(colormap)

def visualize_mask_with_overlay(img: Image.Image, mask: Image.Image, labels_to_ids: dict[str, int], alpha=0.5):
    img_np = np.array(img.convert("RGB"))
    mask_np = np.array(mask)

    num_cls = len(labels_to_ids)
    palette = get_palette(num_cls)
    colormap = create_colormap(palette)

    overlay = np.zeros((*mask_np.shape, 3), dtype=np.uint8)
    for label, idx in labels_to_ids.items():
        if idx != 0:
            overlay[mask_np == idx] = np.array(colormap(idx)[:3]) * 255

    blended = Image.fromarray(np.uint8(img_np * (1 - alpha) + overlay * alpha))

    return blended

def create_legend_image(labels_to_ids: dict[str, int], filename="legend.png"):
    num_cls = len(labels_to_ids)
    palette = get_palette(num_cls)
    colormap = create_colormap(palette)

    fig, ax = plt.subplots(figsize=(4, 6), facecolor="white")

    ax.axis("off")

    legend_elements = [
        Patch(facecolor=colormap(i), edgecolor="black", label=label)
        for label, i in sorted(labels_to_ids.items(), key=lambda x: x[1])
    ]

    plt.title("Legend", fontsize=16, fontweight="bold", pad=20)

    legend = ax.legend(
        handles=legend_elements,
        loc="center",
        bbox_to_anchor=(0.5, 0.5),
        ncol=2,
        frameon=True,
        fancybox=True,
        shadow=True,
        fontsize=10,
        title_fontsize=12,
        borderpad=1,
        labelspacing=1.2,
        handletextpad=0.5,
        handlelength=1.5,
        columnspacing=1.5,
    )

    legend.get_frame().set_facecolor("#FAFAFA")
    legend.get_frame().set_edgecolor("gray")

    # Adjust layout and save
    plt.tight_layout()
    plt.savefig(filename, dpi=300, bbox_inches="tight")
    plt.close()

# ----------------- MODEL ----------------- #

URL = "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/pose/checkpoints/sapiens_1b/sapiens_1b_goliath_best_goliath_AP_640_torchscript.pt2?download=true"
CHECKPOINTS_DIR = os.path.join(ASSETS_DIR, "checkpoints")
os.makedirs(CHECKPOINTS_DIR, exist_ok=True)

model_path = os.path.join(CHECKPOINTS_DIR, "sapiens_1b_goliath_best_goliath_AP_640_torchscript.pt2")

if not os.path.exists(model_path) or os.path.getsize(model_path) == 0:
    print("Downloading model...")
    import requests

    response = requests.get(URL)
    if response.status_code == 200:
        with open(model_path, "wb") as file:
            file.write(response.content)
    else:
        raise Exception("Failed to download the model. Please check the URL.")

model = torch.jit.load(model_path)
model.eval()

@torch.no_grad()
def run_model(input_tensor, height, width):
    output = model(input_tensor)
    output = torch.nn.functional.interpolate(output, size=(height, width), mode="bilinear", align_corners=False)
    _, preds = torch.max(output, 1)
    return preds

transform_fn = transforms.Compose(
    [
        transforms.Resize((1024, 768)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)

# ----------------- CORE FUNCTION ----------------- #

def resize_and_pad(image: Image.Image, target_size=(768, 1024)):
    img_ratio = image.width / image.height
    target_ratio = target_size[0] / target_size[1]

    if img_ratio > target_ratio:
        new_width = target_size[0]
        new_height = int(target_size[0] / img_ratio)
    else:
        new_height = target_size[1]
        new_width = int(target_size[1] * img_ratio)

    resized_image = image.resize((new_width, new_height), Image.LANCZOS)

    delta_w = target_size[0] - new_width
    delta_h = target_size[1] - new_height
    padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
    padded_image = ImageOps.expand(resized_image, padding, fill="black")

    return padded_image

def segment(image: Image.Image) -> Image.Image:
    image = resize_and_pad(image, target_size=(768, 1024))
    
    input_tensor = transform_fn(image).unsqueeze(0)
    preds = run_model(input_tensor, height=image.height, width=image.width)
    mask = preds.squeeze(0).cpu().numpy()
    mask_image = Image.fromarray(mask.astype("uint8"))
    blended_image = visualize_mask_with_overlay(image, mask_image, LABELS_TO_IDS, alpha=0.5)
    return blended_image

# ----------------- GRADIO UI ----------------- #

with open("banner.html", "r") as file:
    banner = file.read()
with open("tips.html", "r") as file:
    tips = file.read()

CUSTOM_CSS = """
.image-container img {
    max-width: 512px;
    max-height: 512px;
    margin: 0 auto;
    border-radius: 0px;
}
.gradio-container {background-color: #fafafa}
"""

with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Monochrome(radius_size=sizes.radius_md)) as demo:
    gr.HTML(banner)
    gr.HTML(tips)
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input Image", type="pil", format="png")
        with gr.Column():
            result_image = gr.Image(label="Depth Output", format="png")
            run_button = gr.Button("Run")

    run_button.click(
        fn=segment,
        inputs=[input_image],
        outputs=[result_image],
    )

if __name__ == "__main__":
    demo.launch(share=False)