import gradio as gr
import numpy as np
import torch
from datasets import load_dataset
from transformers import AutoTokenizer
import matplotlib.pyplot as plt
import seaborn as sns
from enum import Enum
from spacy import displacy


class VisType(Enum):
    SUM = 'Sum over Layers'
    

num_layers = 24
dataset = load_dataset('dar-tau/grammar-attention-maps-opt-350m')['train']
tokenizer = AutoTokenizer.from_pretrained('facebook/opt-350m', add_prefix_space=True)


def analyze_sentence(index, vis_type, vis_format):
    row = dataset[index]
    text = row['text']
    tokenized = tokenizer.batch_decode(tokenizer.encode(text, add_special_tokens=False))
    attn_map_shape = row['attention_maps_shape'][1:]
    seq_len = attn_map_shape[1]
    attn_maps = np.array(row['attention_maps']).reshape(*attn_map_shape).clip(0, 1)
    attn_maps = attn_maps[:, 1:, 1:]
    if vis_type == VisType.SUM.value:
        plot_data = attn_maps.sum(0)
    elif vis_type.startswith('Layer #'):
        layer_to_inspect = int(vis_type.split('#')[1])
        plot_data = attn_maps[layer_to_inspect]
    else:
        print(vis_type)
        0/0
    if vis_format == 'Plot':
        fig = plt.figure(figsize=(0.1 + 0.3 * len(tokenized), 0.25 * len(tokenized)))
        sns.heatmap(plot_data)
        plt.xticks(np.arange(seq_len - 1) + 0.5, tokenized[1:], rotation=90);
        plt.yticks(np.arange(seq_len - 1) + 0.5, tokenized[1:], rotation=0);
        plt.ylabel('TARGET')
        plt.xlabel('SOURCE')    
        plt.grid()
        fig_output, graph_output = fig, ""
    elif vis_format == 'Graph':
        ex = [{
            "words": [{"text": x, "tag": ""} for x in tokenized[1:]],
            "arcs": [{"start": j, "end": i, "label": "", "dir": "right"} 
                         for i in range(len(tokenized) - 1) for j in range(i) if plot_data[i, j] > 0.5 and abs(i-j) > 1
                    ]
             }]
        graph_output = displacy.render(ex, style="dep", jupyter=False, manual=True, options={"compact": True,
                                                                                             "offset_x": 20,
                                                                                             "distance": 130
                                                                                            })
        graph_output = ("<div class='displacy_container' style='max-width:100%; max-height:500px; overflow:auto'>" 
                        + graph_output + 
                        "</div>"
                       )
        fig_output = None
    else:
        fig_output = None
        plot_data[np.arange(len(plot_data)), np.arange(len(plot_data))] = 0.
        top_values, top_indices = torch.tensor(plot_data).flatten().topk(30)
        topk_data = []
        for val, ind in zip(top_values, top_indices):
            if val < 0.5:
                break
            ind = np.unravel_index(ind, plot_data.shape)
            topk_data += [str((tokenized[1+ind[0]], tokenized[1+ind[1]]))]
        graph_output = '<div><b>' + text + '</b><br/>' + '<br/>'.join(topk_data) + '</div>'
    metrics = {'Metrics': 1}
    
    metrics.update({k: v for k, v in row.items() if k not in ['text', 'attention_maps', 'attention_maps_shape']})    
    return fig_output, graph_output, metrics


demo = gr.Blocks(css=".displacy_container svg{height:500px !important; margin-top:-100px; transform: scale(0.5)}")
with demo:
    with gr.Row():
        sentence_dropdown = gr.Dropdown(label="Sentence", 
                                        choices=[x.split('</s> ')[1] for x in dataset['text']], 
                               value=0, min_width=300, type='index')
        vis_dropdown = gr.Dropdown(label="Visualization", choices=[x.value for x in VisType] + 
                                   [f'Layer #{i}' for i in range(num_layers)], 
                                   min_width=70, value=VisType.SUM, type='value')
        btn = gr.Button("Run", min_width=30)
    
    vis_format_checkbox = gr.Radio(['Plot', 'Graph', 'Text'])
        
    output = gr.Plot(label="Plot", container=True)
    graph_output = gr.HTML()
    metrics = gr.Label("Metrics")
    btn.click(analyze_sentence, 
              [sentence_dropdown, vis_dropdown, vis_format_checkbox], 
              [output, graph_output, metrics])


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