# %%
import os
USE_HUGGINGFACE_ZEROGPU = os.getenv("USE_HUGGINGFACE_ZEROGPU", "False").lower() in ["true", "1", "yes"]
#%%
if USE_HUGGINGFACE_ZEROGPU:  # huggingface ZeroGPU, dynamic GPU allocation 
    try:
        import spaces
    except:
        USE_HUGGINGFACE_ZEROGPU = False
        
import gradio as gr

import torch
import torch.nn.functional as F
from PIL import Image
import numpy as np
import time
import threading
import os
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np

from ncut_pytorch import NCUT, eigenvector_to_rgb

from ncut_pytorch.backbone_text import MODEL_DICT as TEXT_MODEL_DICT
from ncut_pytorch.backbone_text import LAYER_DICT as TEXT_LAYER_DICT

def compute_ncut(
    features,
    num_eig=100,
    num_sample_ncut=10000,
    affinity_focal_gamma=0.3,
    knn_ncut=10,
    knn_tsne=10,
    embedding_method="UMAP",
    num_sample_tsne=300,
    perplexity=150,
    n_neighbors=150,
    min_dist=0.1,
    sampling_method="fps",
    metric="cosine",
):        
    logging_str = ""
    num_nodes = np.prod(features.shape[:-1])
    if num_nodes / 2 < num_eig:
        # raise gr.Error("Number of eigenvectors should be less than half the number of nodes.")
        gr.Warning("Number of eigenvectors should be less than half the number of nodes.\n" f"Setting num_eig to {num_nodes // 2 - 1}.")
        num_eig = num_nodes // 2 - 1
        logging_str += f"Number of eigenvectors should be less than half the number of nodes.\n" f"Setting num_eig to {num_nodes // 2 - 1}.\n"
    
    start = time.time()
    eigvecs, eigvals = NCUT(
        num_eig=num_eig,
        num_sample=num_sample_ncut,
        device="cuda" if torch.cuda.is_available() else "cpu",
        affinity_focal_gamma=affinity_focal_gamma,
        knn=knn_ncut,
        sample_method=sampling_method,
        distance=metric,
        normalize_features=False,
    ).fit_transform(features.reshape(-1, features.shape[-1]))
    # print(f"NCUT time: {time.time() - start:.2f}s")
    logging_str += f"NCUT time: {time.time() - start:.2f}s\n"
    
    start = time.time()
    _, rgb = eigenvector_to_rgb(
        eigvecs,
        method=embedding_method,
        num_sample=num_sample_tsne,
        perplexity=perplexity,
        n_neighbors=n_neighbors,
        min_distance=min_dist,
        knn=knn_tsne,
        device="cuda" if torch.cuda.is_available() else "cpu",
    )
    logging_str += f"{embedding_method} time: {time.time() - start:.2f}s\n"

    rgb = rgb.reshape(features.shape[:-1] + (3,))
    return rgb, logging_str, eigvecs


def make_plot(token_texts, rgb, num_lines=50, title=""):
    fig, ax = plt.subplots(figsize=(10, 20))
    # Define the colors
    # fill nan with 0
    rgb = np.nan_to_num(rgb)
    colors = [mcolors.rgb2hex(rgb[i]) for i in range(len(token_texts))]

    # Split the sentence into words
    words = token_texts

    y_pos = 0.96
    x_pos = 0.0
    max_word_length = max(len(word) for word in words)
    count = 0
    for word, color in zip(words, colors):
        if '\n' in word:
            word = word.replace('\n', '')
            y_pos -= 0.025
            x_pos = 0.0
            count += 1
            if count >= num_lines:
                break
        
        text_color = 'black' if sum(mcolors.hex2color(color)) > 1.3 else 'white'  # Choose text color based on background color
        # text_color = 'black'
        txt = ax.text(x_pos, y_pos, word, color=text_color, fontsize=12, bbox=dict(facecolor=color, alpha=0.8, edgecolor='none', pad=2))
        txt_width = txt.get_window_extent().width / (fig.dpi * fig.get_size_inches()[0])  # Calculate the width of the text in inches
        
        x_pos += txt_width * 1.2 + 0.01  # Adjust the spacing between words
        
        if x_pos > 0.97:
            y_pos -= 0.025
            x_pos = 0.0
            count += 1
            if count >= num_lines:
                break
        # break
            
    # Remove the axis ticks and spines
    ax.set_xticks([])
    ax.set_yticks([])
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['bottom'].set_visible(False)
    ax.spines['left'].set_visible(False)
    
    ax.set_title(title, fontsize=20)

    return fig



def ncut_run(
    model,
    text,
    model_name,
    layer=-1,
    num_eig=100,
    node_type="block",
    affinity_focal_gamma=0.3,
    num_sample_ncut=10000,
    knn_ncut=10,
    embedding_method="UMAP",
    num_sample_tsne=1000,
    knn_tsne=10,
    perplexity=500,
    n_neighbors=500,
    min_dist=0.1,
    sampling_method="fps",
):
    logging_str = ""
    if perplexity >= num_sample_tsne or n_neighbors >= num_sample_tsne:
        # raise gr.Error("Perplexity must be less than the number of samples for t-SNE.")
        gr.Warning("Perplexity/n_neighbors must be less than the number of samples.\n" f"Setting Perplexity to {num_sample_tsne-1}.")
        logging_str += f"Perplexity/n_neighbors must be less than the number of samples.\n" f"Setting Perplexity to {num_sample_tsne-1}.\n"
        perplexity = num_sample_tsne - 1
        n_neighbors = num_sample_tsne - 1

    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        
    node_type = node_type.split(":")[0].strip()
    
    model = model.to("cuda" if torch.cuda.is_available() else "cpu")
    
    start = time.time()
    out = model(text)
    features = out[node_type][layer-1].squeeze(0).detach().float()
    token_texts = out["token_texts"]
    
    if perplexity >= features.shape[0] or n_neighbors >= features.shape[0]:
        # raise gr.Error("Perplexity must be less than the number of samples.")
        gr.Warning("Perplexity/n_neighbors must be less than the number of samples.\n" f"Setting Perplexity to {features.shape[0]-1}.")
        logging_str += f"Perplexity/n_neighbors must be less than the number of samples.\n" f"Setting Perplexity to {features.shape[0]-1}.\n"
        perplexity = features.shape[0] - 1
        n_neighbors = features.shape[0] - 1
    
    # print(f"Feature extraction time (gpu): {time.time() - start:.2f}s")
    logging_str += f"Backbone time: {time.time() - start:.2f}s\n"
    
    rgb, _logging_str, _ = compute_ncut(
        features,
        num_eig=num_eig,
        num_sample_ncut=num_sample_ncut,
        affinity_focal_gamma=affinity_focal_gamma,
        knn_ncut=knn_ncut,
        knn_tsne=knn_tsne,
        num_sample_tsne=num_sample_tsne,
        embedding_method=embedding_method,
        perplexity=perplexity,
        n_neighbors=n_neighbors,
        min_dist=min_dist,
        sampling_method=sampling_method,
    )
    logging_str += _logging_str
    
    start = time.time() 
    title = f"{model_name}, Layer {layer}, {node_type}"
    fig = make_plot(token_texts, rgb, title=title)
    logging_str += f"Plotting time: {time.time() - start:.2f}s\n"
    return fig, logging_str

def _ncut_run(*args, **kwargs):
    try:
        ret = ncut_run(*args, **kwargs)
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return ret
    except Exception as e:
        gr.Error(str(e))
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return None, "Error: " + str(e)
    
if USE_HUGGINGFACE_ZEROGPU:
    @spaces.GPU(duration=30)
    def __ncut_run(*args, **kwargs):
        return _ncut_run(*args, **kwargs)
else:
    def __ncut_run(*args, **kwargs):
        return _ncut_run(*args, **kwargs)

def real_run(model_name, text, layer, node_type, num_eig, affinity_focal_gamma, num_sample_ncut, knn_ncut, embedding_method, num_sample_tsne, knn_tsne, perplexity, n_neighbors, min_dist, sampling_method):
    model = TEXT_MODEL_DICT[model_name]()
    return __ncut_run(model, text, model_name, layer, num_eig, node_type, 
                affinity_focal_gamma, num_sample_ncut, knn_ncut, embedding_method, 
                num_sample_tsne, knn_tsne, perplexity, n_neighbors, min_dist, sampling_method)

lines = \
"""1. The majestic giraffe, with its towering height and distinctive long neck, roams the savannas of Africa. These gentle giants use their elongated tongues to pluck leaves from the tallest trees, making them well-adapted to their environment. Their unique coat patterns, much like human fingerprints, are unique to each individual.
2. Penguins, the tuxedoed birds of the Antarctic, are expert swimmers and divers. These flightless seabirds rely on their dense, waterproof feathers and streamlined bodies to propel through icy waters in search of fish, krill, and other marine life. Their huddled colonies and amusing waddles make them a favorite among wildlife enthusiasts.
3. The mighty African elephant, the largest land mammal, is revered for its intelligence and strong family bonds. These gentle giants use their versatile trunks for various tasks, from drinking and feeding to communicating and greeting one another. Their massive ears and wrinkled skin make them an iconic symbol of the African wilderness.
4. The colorful and flamboyant peacock, native to Asia, is known for its stunning iridescent plumage. During mating season, the males fan out their magnificent train of feathers, adorned with intricate eye-like patterns, in an elaborate courtship display to attract potential mates, making them a true spectacle of nature.
5. The sleek and powerful cheetah, the fastest land animal, is built for speed and agility. With its distinctive black tear-like markings and slender body, this feline predator can reach top speeds of up to 70 mph during short bursts, allowing it to chase down its prey with remarkable precision.
6. The playful and intelligent dolphin, a highly social marine mammal, is known for its friendly demeanor and impressive acrobatic abilities. These aquatic creatures use echolocation to navigate and hunt, and their complex communication systems have long fascinated researchers studying their intricate social structures and cognitive abilities.
7. The majestic bald eagle, the national emblem of the United States, soars high above with its distinctive white head and tail feathers. These powerful raptors are skilled hunters, swooping down from great heights to catch fish and other prey with their sharp talons, making them an iconic symbol of strength and freedom.
8. The industrious beaver, nature's skilled engineers, are known for their remarkable ability to construct dams and lodges using their sharp incisors and webbed feet. These semiaquatic rodents play a crucial role in shaping their aquatic ecosystems, creating habitats for numerous other species while demonstrating their ingenuity and perseverance.
9. The vibrant and enchanting hummingbird, one of the smallest bird species, is a true marvel of nature. With their rapidly flapping wings and ability to hover in mid-air, these tiny feathered creatures are expert pollinators, flitting from flower to flower in search of nectar and playing a vital role in plant reproduction.
10. The majestic polar bear, the apex predator of the Arctic, is perfectly adapted to its icy environment. With its thick insulating fur and specialized paws for gripping the ice, this powerful carnivore relies on its exceptional hunting skills and keen senses to locate and capture seals, its primary prey, in the harsh Arctic landscape.
"""

def make_demo():
    with gr.Row():
        with gr.Column(scale=5, min_width=200):
            gr.Markdown("### Input Text")
            placeholder = lines
            input_text = gr.Text(value=placeholder, label="Input Text", placeholder="Type here", lines=12)
            submit_button = gr.Button("🔴 RUN", elem_id="submit_button", variant='primary')
            clear_button = gr.Button("🗑️Clear", elem_id='clear_button', variant='stop')
        with gr.Column(scale=5, min_width=200):
            gr.Markdown("### Parameters <a style='color: #0044CC;' href='https://ncut-pytorch.readthedocs.io/en/latest/how_to_get_better_segmentation/' target='_blank'>Help</a>")
            model_name = gr.Dropdown(list(TEXT_MODEL_DICT.keys()), label="Model", value="meta-llama/Meta-Llama-3.1-8B")
            layer = gr.Slider(1, 32, step=1, value=32, label="Layer")
            node_type = gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Node Type", value="block: sum of residual")
            num_eig = gr.Slider(minimum=1, maximum=1000, step=1, value=100, label="Number of Eigenvectors")
            
            with gr.Accordion("➡️ Click to expand: more parameters", open=False):
                gr.Markdown("<a href='https://ncut-pytorch.readthedocs.io/en/latest/how_to_get_better_segmentation/' target='_blank'>Docs: How to Get Better Segmentation</a>")
                affinity_focal_gamma = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.3, label="Affinity Focal Gamma")
                num_sample_ncut = gr.Slider(minimum=100, maximum=50000, step=100, value=10000, label="Number of Samples for NCUT")
                sampling_method_dropdown = gr.Dropdown(["fps", "random"], label="Sampling method", value="fps", elem_id="sampling_method")
                knn_ncut = gr.Slider(minimum=1, maximum=100, step=1, value=10, label="KNN for NCUT")
                embedding_method_dropdown = gr.Dropdown(["tsne_3d", "umap_3d", "umap_shpere", "tsne_2d", "umap_2d"], label="Coloring method", value="tsne_3d", elem_id="embedding_method")
                num_sample_tsne_slider = gr.Slider(100, 10000, step=100, label="t-SNE/UMAP: num_sample", value=300, elem_id="num_sample_tsne", info="Nyström approximation")
                knn_tsne_slider = gr.Slider(1, 100, step=1, label="t-SNE/UMAP: KNN", value=10, elem_id="knn_tsne", info="Nyström approximation")
                perplexity_slider = gr.Slider(10, 1000, step=10, label="t-SNE: Perplexity", value=150, elem_id="perplexity")
                n_neighbors_slider = gr.Slider(10, 1000, step=10, label="UMAP: n_neighbors", value=150, elem_id="n_neighbors")
                min_dist_slider = gr.Slider(0.1, 1, step=0.1, label="UMAP: min_dist", value=0.1, elem_id="min_dist")
            logging_str = gr.Textbox("", label="Logging Information", placeholder="Logging",)
            
    with gr.Row():
        gr.Markdown("### Output Embedding")
        output_image = gr.Plot(label="NCUT Output", min_width=1920)

    def change_layer_slider(model_name):
        layer_dict = TEXT_LAYER_DICT
        if model_name in layer_dict:
            value = layer_dict[model_name]
            return (gr.Slider(1, value, step=1, label="Backbone: Layer index", value=value, elem_id="layer", visible=True),
                    gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Backbone: Layer type", value="block: sum of residual", elem_id="node_type", info="which feature to take from each layer?"))
        else:
            value = 12
            return (gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Backbone: Layer type", value="block: sum of residual", elem_id="node_type", info="which feature to take from each layer?"),
                    gr.Slider(1, value, step=1, label="Backbone: Layer index", value=value, elem_id="layer", visible=True))
    model_name.change(fn=change_layer_slider, inputs=model_name, outputs=[layer, node_type])

    clear_button.click(lambda x: (None, None), outputs=[input_text, output_image])
    submit_button.click(real_run, inputs=[
        model_name, input_text, layer, node_type, num_eig, 
        affinity_focal_gamma, num_sample_ncut, knn_ncut, 
        embedding_method_dropdown, num_sample_tsne_slider, 
        knn_tsne_slider, perplexity_slider, n_neighbors_slider, 
        min_dist_slider, sampling_method_dropdown
        ], 
        outputs=[output_image, logging_str], 
    )
    


if __name__ == "__main__":
    with gr.Blocks() as demo:
        make_demo()
    demo.launch(share=True)

# %%