import nltk
import spacy
nltk.download('wordnet')
spacy.cli.download('en_core_web_sm')

import torch
import joblib, json
import numpy as np
import pandas as pd
import gradio as gr
from const import used_indices, name_map
from model import get_model
from options import parse_args
from transformers import T5Tokenizer
from compute_lng import compute_lng
from sklearn.experimental import enable_iterative_imputer
from sklearn.impute import IterativeImputer
from sklearn.linear_model import Ridge


def process_examples(samples, full_names):
    processed = []
    for sample in samples:
        processed.append([
            sample['sentence1'],
            pd.DataFrame({'Index': full_names, 'Source': sample['sentence1_ling'], 'Target': sample['sentence2_ling']})
                    ])
    return processed

args, args_list, lng_names = parse_args(ckpt='./ckpt/model.pt')

tokenizer = T5Tokenizer.from_pretrained(args.model_name)
device = 'cuda' if torch.cuda.is_available() else 'cpu'

lng_names = [name_map[x] for x in lng_names]
examples = json.load(open('assets/examples.json'))
examples = process_examples(examples, lng_names)

stats = json.load(open('assets/stats.json'))

scaler = joblib.load('assets/scaler.bin')
scale_ratio = np.load('assets/ratios.npy')

ling_collection = np.load('assets/ling_collection.npy')
ling_collection_scaled = scaler.transform(ling_collection)

model, ling_disc, sem_emb = get_model(args, tokenizer, device)

state = torch.load(args.ckpt, map_location=torch.device('cpu'))
model.load_state_dict(state['model'], strict=True)
model.eval()
ling_disc.eval()

state = torch.load(args.sem_ckpt, map_location=torch.device('cpu'))
sem_emb.load_state_dict(state['model'], strict=True)
sem_emb.eval()

device = model.backbone.device

############# Start demo code
def round_ling(x):
    is_int = stats['is_int']
    mins = stats['min']
    maxs = stats['max']
    for i in range(len(x)):
        # if is_int[i]:
        #     x[i] = round(x[i])
        # else:
        #     x[i] = round(x[i], 3)
        x[i] = round(x[i], 3)
    return np.clip(x, mins, maxs)

def visibility(mode):
    if mode == 0:
        vis_group = group1
    elif mode == 1:
        vis_group = group2
    elif mode == 2:
        vis_group = group3

    output = [gr.update(value=''), gr.update(value='')]
    for component in components:
        if component in vis_group:
            output.append(gr.update(visible=True))
        else:
            output.append(gr.update(visible=False))
    return output

def generate(sent1, ling):
    input_ids = tokenizer.encode(sent1, return_tensors='pt').to(device)
    ling1 = scaler.transform([ling['Source']])
    ling2 = scaler.transform([ling['Target']])
    inputs = {'sentence1_input_ids': input_ids,
            'sentence1_ling': torch.tensor(ling1).float().to(device),
            'sentence2_ling': torch.tensor(ling2).float().to(device),
            'sentence1_attention_mask': torch.ones_like(input_ids)}
    preds = []
    with torch.no_grad():
        pred = model.infer(inputs).cpu().numpy()
    pred = tokenizer.batch_decode(pred,
            skip_special_tokens=True)[0]

    return pred

def generate_with_feedback(sent1, ling, approx):
    if sent1 == '':
        return ['Please input a source text.', '']


    input_ids = tokenizer.encode(sent1, return_tensors='pt').to(device)
    ling2 = torch.tensor(scaler.transform([ling['Target']])).float().to(device)
    inputs = {
            'sentence1_input_ids': input_ids,
            'sentence2_ling': ling2,
            'sentence1_attention_mask': torch.ones_like(input_ids)
              }

    pred, (pred_text, interpolations) = model.infer_with_feedback_BP(ling_disc, sem_emb, inputs, tokenizer)

    interpolation = '-- ' + '\n-- '.join(interpolations)
    return [pred_text, interpolation]

def generate_random(sent1, ling, count, approx):
    preds, interpolations = [], []
    for c in range(count):
        idx = np.random.randint(0, len(ling_collection))
        ling_ex = ling_collection[idx]
        ling['Target'] = ling_ex
        pred, interpolation =  generate_with_feedback(sent1, ling, approx)
        preds.append(pred)
        interpolations.append(interpolation)
    return '\n***\n'.join(preds), '\n***\n'.join(interpolations), ling

def estimate_gen(sent1, sent2, ling, approx):
    if 'approximate' in approx:
        input_ids = tokenizer.encode(sent2, return_tensors='pt').to(device)
        with torch.no_grad():
            ling_pred = ling_disc(input_ids=input_ids).cpu().numpy()
        ling_pred = scaler.inverse_transform(ling_pred)[0]
    elif 'exact' in approx:
        ling_pred = np.array(compute_lng(sent2))[used_indices]
    else:
        raise ValueError()

    ling_pred = round_ling(ling_pred)
    ling['Target'] = ling_pred
    gen = generate_with_feedback(sent1, ling, approx)
    results = gen + [ling]

    return results

def estimate_tgt(sent2, ling, approx):
    if 'approximate' in approx:
        input_ids = tokenizer.encode(sent2, return_tensors='pt').to(device)
        with torch.no_grad():
            ling_pred = ling_disc(input_ids=input_ids).cpu().numpy()
        ling_pred = scaler.inverse_transform(ling_pred)[0]
    elif 'exact' in approx:
        ling_pred = np.array(compute_lng(sent2))[used_indices]
    else:
        raise ValueError()

    ling_pred = round_ling(ling_pred)
    ling['Target'] = ling_pred
    return ling

def estimate_src(sent1, ling, approx):
    if 'approximate' in approx:
        input_ids = tokenizer.encode(sent1, return_tensors='pt').to(device)
        with torch.no_grad():
            ling_pred = ling_disc(input_ids=input_ids).cpu().numpy()
        ling_pred = scaler.inverse_transform(ling_pred)[0]
    elif 'exact' in approx:
        ling_pred = np.array(compute_lng(sent1))[used_indices]
    else:
        raise ValueError()

    ling['Source'] = ling_pred
    return ling

def rand_target(ling):
    ling['Target'] = scaler.inverse_transform([np.random.randn(*ling['Target'].shape)])[0]
    return ling

def rand_ex_target(ling):
    idx = np.random.randint(0, len(ling_collection))
    ling_ex = ling_collection[idx]
    ling['Target'] = ling_ex
    return ling

def copy(ling):
    ling['Target'] = ling['Source']
    return ling

def add(ling):
    scale_stepsize = np.random.uniform(1.0, 5.0)
    x = ling['Target'] + scale_stepsize * scale_ratio
    x = round_ling(x)
    ling['Target'] =  x
    return ling

def sub(ling):
    scale_stepsize = np.random.uniform(1.0, 5.0)
    x = ling['Target'] - scale_stepsize * scale_ratio
    x = round_ling(x)
    ling['Target'] =  x
    return ling

def impute(ling):
    ling['Target'] = ling['Target'].replace("", np.nan)
    ling['Target'] = scaler.transform([ling['Target']])[0]
    estimator = Ridge(alpha=1e3, fit_intercept=False)
    imputer = IterativeImputer(estimator=estimator, imputation_order='random', max_iter=100)

    combined_matrix = np.vstack([ling_collection, ling['Target']])
    interpolated_matrix = imputer.fit_transform(combined_matrix)
    interpolated_vector = interpolated_matrix[-1]

    interp_raw = scaler.inverse_transform([interpolated_vector])[0]

    ling['Target'] = round_ling(interp_raw)
    return ling

title = """
<h1 style="text-align: center;">Controlled Paraphrase Generation with Linguistic Feature Control</h1>

<p style="font-size:1.2em;">This system utilizes an encoder-decoder model to generate text with controlled complexity, guided by 40 linguistic complexity indices.
The model can generate diverse paraphrases of a given sentence, each adjusted to maintain consistent meaning while varying
in linguistic complexity according to the desired level.</p>
<p style="font-size:1.2em;">It is important to note that not all index combinations are feasible (e.g., a sentence of "length" 5 with 10 "unique words").
To ensure high-quality outputs, our approach iteratively adjusts the generated text to match the closest, yet coherent
achievable set of indices for the given target.</p>
"""

guide = """
You may use the system in on of the following ways:

**Randomized Paraphrase Generation**: Select this option to produce multiple paraphrases with a range
of linguistic complexity. You need to provide a source text, specify the number of paraphrases you want,
and click "Generate." The linguistic complexity of the paraphrases will be determined randomly.

**Complexity-Matched Paraphrasing**: Select this option to generate a paraphrase of the given source
sentence that closely mirrors the linguistic complexity of another given sentence. Input your source
sentence along with another sentence (which will serve only to extract linguistic indices for the
paraphrase generation). Then, click "Generate."

**Manual Linguistic Control**: Select this option to manually control the linguistic complexity of the
generated text. We provided a set of tools for manual adjustments of the desired linguistic complexity of
the target sentence. These tools enable the user to extract linguistic indices from a given sentence,
generate a random (yet coherent) set of linguistic indices, and add or subtract to them.
These tools are designed for experimental use and require the user to possess linguistic expertise for
effective input of linguistic indices. To use these tools, select "Tools to assist in setting linguistic
indices." Once indices are entered, click "Generate."


Second, you may select to use exact or approximate computation of linguistic indices. Approximate computation is significantly faster.

Third, you may view the intermediate sentences of the quality control process by selecting the checkbox under "Advanced Options".

Fourth, you may try out some examples by clicking on "Examples...". Examples consist of a source sentences
and a sample set of target linguistic indices.

Please make your choice below.

"""

sent1 = gr.Textbox(label='Source text')
ling = gr.Dataframe(value = [[x, 0, 0] for x in lng_names],
        headers=['Index', 'Source', 'Target'],
        datatype=['str', 'number', 'number'], visible=False)
css = """
#guide span.svelte-1w6vloh {font-size: 22px !important; font-weight: 600 !important}
#mode span.svelte-1gfkn6j {font-size: 18px !important; font-weight: 600 !important}
#mode {border: 0px; box-shadow: none}
#mode .block {padding: 0px}

#estimate textarea {border: 1px solid; border-radius: 7px}

div.gradio-container {color: black}
div.form {background: inherit}

body {
  --text-sm: 12px;
  --text-md: 16px;
  --text-lg: 18px;
  --input-text-size: 16px;
  --section-text-size: 16px;
  --input-background: --neutral-50;
  }

.top-separator {
  width: 100%;
  height: 4px; /* Adjust the height for boldness */
  background-color: #000; /* Adjust the color as needed */
  margin-top: 20px; /* Adjust the margin as needed */
}
.bottom-separator {
  width: 100%;
  height: 4px; /* Adjust the height for boldness */
  background-color: #000; /* Adjust the color as needed */
  margin-bottom: 20px; /* Adjust the margin as needed */
}
"""

with gr.Blocks(
        theme=gr.themes.Default(
            spacing_size=gr.themes.sizes.spacing_md,
            text_size=gr.themes.sizes.text_md,
            ),
        css=css) as demo:
    gr.Image('assets/logo.png', height=100, container=False, show_download_button=False)
    gr.Markdown(title)
    with gr.Accordion("🚀 Quick Start Guide", open=False, elem_id='guide'):
        gr.Markdown(guide)

    with gr.Group(elem_classes='top-separator'):
        pass
    with gr.Group(elem_id='mode'):
        mode = gr.Radio(
                value='Randomized Paraphrase Generation',
                label='Operation Modes',
                type="index",
                choices=['🔄 Randomized Paraphrase Generation',
                         '⚖️ Complexity-Matched Paraphrasing',
                         '🎛️ Manual Linguistic Control'],
                        )
    with gr.Accordion("⚙️ Advanced Options", open=False):
        approx = gr.Radio(value='Use approximate computation of linguistic indices (faster)',
                choices=['Use approximate computation of linguistic indices (faster)',
                    'Use exact computation of linguistic indices'], container=False, show_label=False)
        control_interpolation = gr.Checkbox(label='View the intermediate sentences in the interpolation of linguistic indices')

    with gr.Accordion("📑 Examples...", open=False):
        gr.Examples(examples, [sent1, ling], examples_per_page=4, label=None)

    with gr.Row():
        sent1.render()
        with gr.Column():
            sent2 = gr.Textbox(label='Generated text')
    interpolation = gr.Textbox(label='Quality control interpolation', visible=False, lines=5)
    with gr.Group(elem_classes='bottom-separator'):
        pass
    #####################
    with gr.Row():
        generate_random_btn = gr.Button("Generate",
                variant='primary', scale=1, visible=True)
        count = gr.Number(label='Number of generated sentences', value=3, precision=0, scale=1, visible=True)
    # generate_fb_btn = gr.Button("Generate with auto-adjust (towards pred)")
    # generate_fb_s_btn = gr.Button("Generate with auto-adjust (moving s)")
    #####################
    with gr.Row():
        estimate_gen_btn = gr.Button("Generate", 
                variant='primary',
                scale=1, visible=False)
        sent_ling_gen = gr.Textbox(label='Text to estimate linguistic indices', scale=1, visible=False)
    #####################
    generate_btn = gr.Button("Generate", variant='primary', visible=False)
    with gr.Accordion("Tools to assist in the setting of linguistic indices...", open=False, visible=False) as ling_tools:
        rand_ex_btn = gr.Button("Random target", size='lg', visible=False)
        impute_btn = gr.Button("Impute Missing Values", size='lg', visible=False)
        with gr.Row():
            estimate_src_btn = gr.Button("Estimate linguistic indices of source sentence", visible=False)
            copy_btn = gr.Button("Copy linguistic indices of source to target", size='lg', visible=False)
        with gr.Row():
            sub_btn = gr.Button('Subtract \u03B5 from target linguistic indices', visible=False)
            add_btn = gr.Button('Add \u03B5 to target linguistic indices', visible=False)
        with gr.Row():
            estimate_tgt_btn = gr.Button("Estimate linguistic indices of this sentence →", visible=False)
            sent_ling_est = gr.Textbox(label='Text to estimate linguistic indices', scale=2, visible=False, container=False, elem_id='estimate')
    ling.render()
    #####################

    estimate_src_btn.click(estimate_src, inputs=[sent1, ling, approx], outputs=[ling])
    estimate_tgt_btn.click(estimate_tgt, inputs=[sent_ling_est, ling, approx], outputs=[ling])
    estimate_gen_btn.click(estimate_gen, inputs=[sent1, sent_ling_gen, ling, approx], outputs=[sent2, interpolation, ling])
    rand_ex_btn.click(rand_ex_target, inputs=[ling], outputs=[ling])
    impute_btn.click(impute, inputs=[ling], outputs=[ling])
    copy_btn.click(copy, inputs=[ling], outputs=[ling])
    generate_btn.click(generate_with_feedback, inputs=[sent1, ling, approx], outputs=[sent2, interpolation])
    generate_random_btn.click(generate_random, inputs=[sent1, ling, count, approx],
            outputs=[sent2, interpolation, ling])
    add_btn.click(add, inputs=[ling], outputs=[ling])
    sub_btn.click(sub, inputs=[ling], outputs=[ling])

    group1 = [generate_random_btn, count]
    group2 = [estimate_gen_btn, sent_ling_gen]
    group3 = [generate_btn, estimate_src_btn, impute_btn, estimate_tgt_btn, sent_ling_est, rand_ex_btn, copy_btn, add_btn, sub_btn, ling, ling_tools]
    components = group1 + group2 + group3
    mode.change(visibility, inputs=[mode], outputs=[sent2, interpolation] + components)
    control_interpolation.change(lambda v: gr.update(visible=v), inputs=[control_interpolation],
            outputs=[interpolation])

print('Finished loading')
demo.launch(share=True)