import gradio as gr
import numpy as np
import torch
import re
from nltk.stem import PorterStemmer
from collections import defaultdict
from transformers import pipeline
from transformers.pipelines import PIPELINE_REGISTRY, FillMaskPipeline
from transformers import AutoModelForMaskedLM

ex_str1 = "A crustless sandwich made from two slices of baked bread. The sandwich includes first and second matching " \
          "crustless bread pieces. The bread pieces have the same general outer shape defined by an outer periphery " \
          "with central portions surrounded by an outer peripheral area, the bread pieces being at least partially " \
          "crimped together at the outer peripheral area."
ex_key1 = "sandwich bread crimped"

ex_str2 = "The present disclosure provides a DNA-targeting RNA that comprises a targeting sequence and, together with" \
          " a modifying polypeptide, provides for site-specific modification of a target DNA and/or a polypeptide" \
          " associated with the target DNA. "
ex_key2 = "DNA target modification"

ex_str3 = "The graphite plane is composed of a two-dimensional hexagonal lattice of carbon atoms and the plate has a " \
          "length and a width parallel to the graphite plane and a thickness orthogonal to the graphite plane with at " \
          "least one of the length, width, and thickness values being 100 nanometers or smaller. "
ex_key3 = "graphite lattice orthogonal "

tab_two_examples = [[ex_str1, ex_key1],
                    [ex_str2, ex_key2],
                    [ex_str3, ex_key3]]
#
# tab_one_examples = [['A crustless _ made from two slices of baked bread.'],
#                     ['The present disclosure provides a DNA-targeting RNA that comprises a targeting _.'],
#                     ['The _ plane is composed of a two-dimensional hexagonal lattice of carbon atoms.']
#                     ]

ignore_str = ['a', 'an', 'the', 'is', 'and', 'or', '!', '(', ')', '-', '[', ']', '{', '}', ';', ':', "'", '"', '\\',
              ',', '<', '>', '.', '/', '?', '@', '#', '$', '%', '^', '&', '*', '_', '~']


def add_mask(text, lower_bound=0, index=None):
    split_text = text.split()
    if index is not None:
        split_text[index] = '[MASK]'
        return ' '.join(split_text), None
    # If the user supplies a mask, don't add more
    if '_' in split_text:
        u_pos = [i for i, s in enumerate(split_text) if '_' in s][0]
        split_text[u_pos] = '[MASK]'
        return ' '.join(split_text), '[MASK]'

    idx = np.random.randint(low=lower_bound, high=len(split_text), size=1).astype(int)[0]
    # Don't mask certain words
    num_iters = 0
    while split_text[idx].lower() in ignore_str:
        num_iters += 1
        idx = np.random.randint(len(split_text), size=1).astype(int)[0]
        if num_iters > 10:
            break

    masked_string = split_text[idx]
    split_text[idx] = '[MASK]'
    masked_output = ' '.join(split_text)
    return masked_output, masked_string


class TempScalePipe(FillMaskPipeline):
    def _sanitize_parameters(self, top_k=None, targets=None, temp=None):
        postprocess_params = {}

        if targets is not None:
            target_ids = self.get_target_ids(targets, top_k)
            postprocess_params["target_ids"] = target_ids

        if top_k is not None:
            postprocess_params["top_k"] = top_k

        if temp is not None:
            postprocess_params["temp"] = temp
        return {}, {}, postprocess_params


    def __call__(self, inputs, *args, **kwargs):
        """
        Fill the masked token in the text(s) given as inputs.

        Args:
            args (`str` or `List[str]`):
                One or several texts (or one list of prompts) with masked tokens.
            targets (`str` or `List[str]`, *optional*):
                When passed, the model will limit the scores to the passed targets instead of looking up in the whole
                vocab. If the provided targets are not in the model vocab, they will be tokenized and the first
                resulting token will be used (with a warning, and that might be slower).
            top_k (`int`, *optional*):
                When passed, overrides the number of predictions to return.

        Return:
            A list or a list of list of `dict`: Each result comes as list of dictionaries with the following keys:

            - **sequence** (`str`) -- The corresponding input with the mask token prediction.
            - **score** (`float`) -- The corresponding probability.
            - **token** (`int`) -- The predicted token id (to replace the masked one).
            - **token** (`str`) -- The predicted token (to replace the masked one).
        """
        outputs = super().__call__(inputs, **kwargs)
        if isinstance(inputs, list) and len(inputs) == 1:
            return outputs[0]
        return outputs

    def postprocess(self, model_outputs, top_k=10, target_ids=None, temp=1):
        # Cap top_k if there are targets
        if target_ids is not None and target_ids.shape[0] < top_k:
            top_k = target_ids.shape[0]
        input_ids = model_outputs["input_ids"][0]
        outputs = model_outputs["logits"]

        masked_index = torch.nonzero(input_ids == self.tokenizer.mask_token_id, as_tuple=False).squeeze(-1)
        # Fill mask pipeline supports only one ${mask_token} per sample

        logits = outputs[0, masked_index, :] / temp
        probs = logits.softmax(dim=-1)
        sampling = False
        if sampling:
            predictions = torch.multinomial(probs, num_samples=3)
            values = probs[0, predictions]
        if target_ids is not None:
            probs = probs[..., target_ids]
        if not sampling:
            values, predictions = probs.topk(top_k)

        result = []
        single_mask = values.shape[0] == 1
        for i, (_values, _predictions) in enumerate(zip(values.tolist(), predictions.tolist())):
            row = []
            for v, p in zip(_values, _predictions):
                # Copy is important since we're going to modify this array in place
                tokens = input_ids.numpy().copy()
                if target_ids is not None:
                    p = target_ids[p].tolist()

                tokens[masked_index[i]] = p
                # Filter padding out:
                tokens = tokens[np.where(tokens != self.tokenizer.pad_token_id)]
                # Originally we skip special tokens to give readable output.
                # For multi masks though, the other [MASK] would be removed otherwise
                # making the output look odd, so we add them back
                sequence = self.tokenizer.decode(tokens, skip_special_tokens=single_mask)
                proposition = {"score": v, "token": p, "token_str": self.tokenizer.decode([p]), "sequence": sequence}
                row.append(proposition)
            result.append(row)
        if single_mask:
            return result[0]
        return result


PIPELINE_REGISTRY.register_pipeline(
    "temp-scale",
    pipeline_class=TempScalePipe,
    pt_model=AutoModelForMaskedLM,
)
scrambler = pipeline("temp-scale", model="anferico/bert-for-patents")

generator = pipeline('text-generation', model='gpt2')

def sample_output(out, sampling):
    score_to_str = {out[k]: k for k in out.keys()}
    score_list = list(score_to_str.keys())
    if sampling == 'multi':
        idx = np.argmax(np.random.multinomial(1, score_list, 1))
    else:
        idx = np.random.randint(0, len(score_list))
    score = score_list[idx]
    return score_to_str[score]


def unmask_single(text, temp=1):
    masked_text, _ = add_mask(text)
    res = scrambler(masked_text, temp=temp, top_k=10)
    out = {item["token_str"]: item["score"] for item in res}
    return out


def unmask(text, temp, rounds, lower_bound=0):
    sampling = 'multi'
    for _ in range(rounds):
        masked_text, masked = add_mask(text, lower_bound)
        split_text = masked_text.split()
        res = scrambler(masked_text, temp=temp, top_k=15)
        mask_pos = [i for i, t in enumerate(split_text) if 'MASK' in t][0]
        out = {item["token_str"]: item["score"] for item in res}
        new_token = sample_output(out, sampling)
        unsuccessful_iters = 0
        while masked in new_token:
            if unsuccessful_iters > 5:
                break
            print('skipped', new_token)
            new_token = sample_output(out, sampling=sampling)
            unsuccessful_iters += 1
        if masked in new_token:
            split_text[mask_pos] = new_token
        else:
            split_text[mask_pos] = '*' + new_token + '*'
        text = ' '.join(split_text)

    text = list(text)
    text[0] = text[0].upper()
    return ''.join(text)


def autocomplete(text, temp):
    output = generator(text, max_length=30, num_return_sequences=1)
    gpt_out = output[0]['generated_text']
    # diff = gpt_out.replace(text, '')
    patent_bert_out = unmask(gpt_out, temp=temp, rounds=5, lower_bound=len(text.split()))
    # Take the output from gpt-2 and randomly mask, if a mask is confident, swap it in. Iterate 5 times
    return patent_bert_out


def extract_keywords(text, queries):
    q_dict = {}
    temp = 1 # set temperature to 1
    for query in queries.split():
        # Iterate through text and mask each token
        ps = PorterStemmer()
        top_scores = defaultdict(list)
        top_k_range = 30
        text_no_punc = re.sub(r'[^\w\s]', '', text)
        indices = [i for i, t in enumerate(text_no_punc.split()) if t.lower() == query.lower()]
        for i in indices:
            masked_text, masked = add_mask(text, index=i)
            res = scrambler(masked_text, temp=temp, top_k=top_k_range)
            out = {item["token_str"]: item["score"] for item in res}
            sorted_keys = sorted(out, key=out.get)
            # If the key does not appear, floor its rank for that round
            for rank, token_str in enumerate(sorted_keys):
                if token_str in ignore_str:
                    continue
                stemmed = ps.stem(token_str)
                norm_rank = rank / top_k_range
                top_scores[stemmed].append(norm_rank)
            for key in top_scores.keys():
                if key not in out.keys():
                    top_scores[key].append(0)
        # Calc mean
        for key in top_scores.keys():
            top_scores[key] = np.mean(top_scores[key])
        # Normalize
        for key in top_scores.keys():
            top_scores[key] = top_scores[key] / np.sum(list(top_scores.values()))
        # Get top_k
        top_n = sorted(list(top_scores.values()))[-3]
        for key in list(top_scores.keys()):
            if top_scores[key] < top_n:
                del top_scores[key]
        q_dict[query] = top_scores

    keywords = ''
    for i, q in enumerate(q_dict.keys()):
        keywords += '['
        for ii, k in enumerate(q_dict[q].keys()):
            keywords += k
            if ii < len(q_dict[q].keys()) - 1:
                keywords += ' OR '
            else:
                keywords += ']'
        if i < len(q_dict.keys()) - 1:
            keywords += ' AND '
    # keywords = set([k for q in q_dict.keys() for k in q_dict[q].keys()])
    # search_str = ' OR '.join(keywords)
    output = [q_dict[q] for q in q_dict]
    output.append(keywords)
    return output
    # fig, ax = plt.subplots(nrows=1, ncols=3)
    # for q in q_dict:
    #     ax.bar(q_dict[q])
    # return fig

label0 = gr.Label(label='keyword 1', num_top_classes=3)
label01 = gr.Label(label='keyword 2', num_top_classes=3)
label02 = gr.Label(label='keyword 3', num_top_classes=3)
textbox02 = gr.Textbox(label="Input Keywords", lines=3)
textbox01 = gr.Textbox(label="Input Keywords", placeholder="Type keywords here", lines=1)
textbox0 = gr.Textbox(label="Input Sentences", placeholder="Type sentences here", lines=5)

output_textbox0 = gr.Textbox(label='Search String of Keywords', placeholder="Output will appear here", lines=4)
# temp_slider0 = gr.Slider(1.0, 3.0, value=1.0, label='Creativity')

textbox1 = gr.Textbox(label="Input Sentence", lines=5)
# output_textbox1 = gr.Textbox(placeholder="Output will appear here", lines=4)
title1 = "Patent-BERT: Context-Dependent Synonym Generator"
description1 = """<p>
Try inserting a few sentences from a patent, and pick keywords for the model to analyze. The model will analyze the 
context of the keywords in the sentences and generate the top three most likely candidates for each word. 
This can be used for more creative patent drafting or patent searches using the generated search string. The base model is 
<a href= "https://github.com/google/patents-public-data/blob/master/models/BERT%20for%20Patents.md">Patent BERT</a> created and trained by Google.

<strong>Note:</strong> Current pipeline only allows for <strong>three</strong> keyword submissions. Stemming (e.g., altering -> alter) is built into the output for 
broader search string. <br/>

Beta features (currently work-in-progress) include: (<strong>A</strong>) adjustment options for (i) the number of keywords, (ii) the number of context-dependent synonyms, 
and (iii) a 'creativity' parameter of the model; (<strong>B</strong>) analysis of where these words appear in the patent (e.g., 
claim, summary, etc.); and (<strong>C</strong>) a stemming option for input keywords. 
<br/>
<p/>"""

# textbox2 = gr.Textbox(label="Input Sentences", lines=5)
# output_textbox2 = gr.Textbox(placeholder="Output will appear here", lines=4)
# temp_slider2 = gr.Slider(1.0, 3.0, value=1.0, label='Creativity')
# edit_slider2 = gr.Slider(1, 20, step=1, value=1.0, label='Number of edits')


# title2 = "Patent-BERT Sentence Remix-er: Multiple Edits"
# description2 = """<p>
#
# Try typing in a sentence for the model to remix. Adjust the 'creativity' scale bar to change the
# the model's confidence in its likely substitutions and the 'number of edits' for the number of edits you want
# the model to attempt to make. The words substituted in the output sentence will be enclosed in asterisks (e.g., *word*).
# <br/> <p/> """

demo0 = gr.Interface(
    fn=extract_keywords,
    inputs=[textbox0, textbox01],
    outputs=[label0, label01, label02, output_textbox0],
    examples=tab_two_examples,
    allow_flagging='never',
    title=title1,
    description=description1
)

# demo1 = gr.Interface(
#     fn=unmask_single,
#     inputs=[textbox1],
#     outputs='label',
#     examples=tab_one_examples,
#     allow_flagging='never',
#     title=title1,
#     description=description1
# )

# demo2 = gr.Interface(
#     fn=unmask,
#     inputs=[textbox2, temp_slider2, edit_slider2],
#     outputs=[output_textbox2],
#     examples=tab_two_examples,
#     allow_flagging='never',
#     title=title2,
#     description=description2
# )

gr.TabbedInterface(
    [demo0], ["Keyword generator"]
).launch()