import torch
import numpy as np
from helper import *
from config.GlobalVariables import *
from SynthesisNetwork import SynthesisNetwork
from DataLoader import DataLoader
import convenience
import gradio as gr


def update_chosen_writers(writer1, writer2, weight, words, all_loaded_data):
    net.clamp_mdn = 0
    chosen_writers = [int(writer1.split(" ")[1]), int(writer2.split(" ")[1])]

    all_loaded_data = []
    for writer_id in chosen_writers:
        loaded_data = dl.next_batch(TYPE='TRAIN', uid=writer_id, tids=list(range(num_samples)))
        all_loaded_data.append(loaded_data)

    writer_mean_Ws = []
    for loaded_data in all_loaded_data:
        mean_global_W = convenience.get_mean_global_W(net, loaded_data, device)
        writer_mean_Ws.append(mean_global_W.detach())

    return gr.Slider.update(label=f"{writer1} vs. {writer2}"), chosen_writers, writer_mean_Ws, *update_writer_word(" ".join(words), writer_mean_Ws, all_loaded_data, weight)

def update_writer_word(target_word, writer_mean_Ws, all_loaded_data, writer_weight, device="cpu"):
    words = []
    for word in target_word.split(" "):
        if len(word) > 0:
            words.append(word)
        
    word_Ws = []
    word_Cs = []
    for word in words:
        writer_Ws, writer_Cs = convenience.get_DSD(net, word, writer_mean_Ws, all_loaded_data, device)
        word_Ws.append(writer_Ws)
        word_Cs.append(writer_Cs)

    if len(words) == 0:
        word_Ws.append(torch.tensor([]))
        word_Cs.append(torch.tensor([]))

    return words, word_Ws, word_Cs, *update_writer_slider(writer_weight, words, word_Ws, word_Cs)

def update_writer_slider(weight, words, all_word_Ws, all_word_Cs):
    weights = [1 - weight, weight]
    net.clamp_mdn = 0
    svg = convenience.draw_words_svg(words, all_word_Ws, all_word_Cs, weights, net)
    return gr.HTML.update(value=svg.tostring()), gr.File.update(visible=False), gr.Button.update(visible=True), weight, svg


def update_writer_download(writer_svg):
    writer_svg.saveas("./DSD_writer_interpolation.svg")
    return gr.File.update(value="./DSD_writer_interpolation.svg", visible=True), gr.Button.update(visible=False)

# for character blend
def update_blend_chars(c1, c2, weight, char_Ws):
    blend_chars = [c1, c2]
    char_Cs = torch.zeros(1, 2, convenience.L, convenience.L)
    for i in range(2):  # get corners of grid
        _, char_matrix = convenience.get_DSD(net, blend_chars[i], default_mean_global_W, [default_loaded_data], device)
        char_Cs[:, i, :, :] = char_matrix

    return gr.Slider.update(label=f"'{c1}' vs. '{c2}'"), char_Cs.detach(), blend_chars, *update_char_slider(weight, char_Ws, char_Cs, blend_chars)

def update_char_slider(weight, char_Ws, char_Cs, blend_chars):
    """Generates an image of handwritten text based on target_sentence"""
    net.clamp_mdn = 0
    character_weights = [1 - weight, weight]

    all_W_c = convenience.get_character_blend_W_c(character_weights, char_Ws, char_Cs)
    all_commands = convenience.get_commands(net, blend_chars[0], all_W_c)
    svg = convenience.commands_to_svg(all_commands, 750, 160, 375)
    return gr.HTML.update(value=svg.tostring()), gr.File.update(visible=False), gr.Button.update(visible=True), weight, svg

def update_char_download(char_svg):
    char_svg.saveas("./DSD_char_interpolation.svg")
    return gr.File.update(value="./DSD_char_interpolation.svg", visible=True), gr.Button.update(visible=False)

# for MDN
def update_mdn_word(target_word, scale_sd, clamp_mdn):
    mdn_words = []
    for word in target_word.split(" "):
        mdn_words.append(word)

    all_word_mdn_Ws = []
    all_word_mdn_Cs = []
    for word in mdn_words:
        all_writer_Ws, all_writer_Cs = convenience.get_DSD(net, word, default_mean_global_W, [default_loaded_data], device)
        all_word_mdn_Ws.append(all_writer_Ws)
        all_word_mdn_Cs.append(all_writer_Cs)

    return mdn_words, all_word_mdn_Ws, all_word_mdn_Cs, *sample_mdn(scale_sd, clamp_mdn, mdn_words, all_word_mdn_Ws, all_word_mdn_Cs)


def sample_mdn(maxs, maxr, mdn_words, all_word_mdn_Ws, all_word_mdn_Cs):
    net.clamp_mdn = maxr
    net.scale_sd = maxs
    svg = convenience.draw_words_svg(mdn_words, all_word_mdn_Ws, all_word_mdn_Cs, [1], net)
    return gr.HTML.update(value=svg.tostring()), gr.File.update(visible=False), gr.Button.update(visible=True), maxr, maxs, svg

def update_mdn_download(mdn_svg):
    mdn_svg.saveas("./DSD_add_randomness.svg")
    return gr.File.update(value="./DSD_add_randomness.svg", visible=True), gr.Button.update(visible=False)

device = 'cpu'
num_samples = 10

net = SynthesisNetwork(weight_dim=256, num_layers=3).to(device)

if not torch.cuda.is_available():
    net.load_state_dict(torch.load('./model/250000.pt', map_location=torch.device(device))["model_state_dict"])


dl = DataLoader(num_writer=1, num_samples=10, divider=5.0, datadir='./data/writers')

writer_options = [5, 14, 15, 16, 17, 22, 25, 80, 120, 137, 147, 151]
all_loaded_data_DEFAULT = []
chosen_writers_DEFAULT = [120, 80]
avail_char = "0 1 2 3 4 5 6 7 8 9 a b c d e f g h i j k l m n o p q r s t u v w x y z A B C D E F G H I J K L M N O P Q R S T U V W X Y Z ! ? \" ' * + - = : ; , . < > \ / [ ] ( ) # $ % &"
avail_char_list = avail_char.split(" ")
for writer_id in chosen_writers_DEFAULT:
    loaded_data = dl.next_batch(TYPE='TRAIN', uid=writer_id, tids=list(range(num_samples)))
    all_loaded_data_DEFAULT.append(loaded_data)

default_loaded_data = all_loaded_data_DEFAULT[-1]
default_mean_global_W = convenience.get_mean_global_W(net, default_loaded_data, device)

# data for writer interpolation
writer_words_DEFAULT = ["hello", "world"]
writer_mean_Ws_DEFAULT = []
writer_all_word_Ws_DEFAULT = []
writer_all_word_Cs_DEFAULT = []
writer_weight_DEFAULT = 0.7
writer_svg_DEFAULT = None

# data for char interpolation
char_chosen_DEFAULT = ["y", "s"]
char_mean_global_W_DEFAULT = None
char_weight_DEFAULT = 0.7
char_Ws_DEFAULT = default_mean_global_W.reshape(1, 1, convenience.L)
char_Cs_DEFAULT = None
char_svg_DEFAULT = None

# # data for MDN
mdn_words_DEFAULT = ["hello", "world"]
all_word_mdn_Ws_DEFAULT = None
all_word_mdn_Cs_DEFAULT = None
clamp_mdn_DEFAULT = 0.5
scale_sd_DEFAULT = 1
mdn_svg_DEFAULT = None

_wrds, writer_all_word_Ws_DEFAULT, writer_all_word_Cs_DEFAULT, _html, _file, _btn, _wt, _svg = update_writer_word(" ".join(writer_words_DEFAULT), writer_mean_Ws_DEFAULT, all_loaded_data_DEFAULT, writer_weight_DEFAULT)
_sldr, _wrtrs, writer_mean_Ws_DEFAULT, _wrds, _waww, _wawc, _html, _file, _btn, _wt, writer_svg_DEFAULT = update_chosen_writers(f"Writer {chosen_writers_DEFAULT[0]}", f"Writer {chosen_writers_DEFAULT[1]}", writer_weight_DEFAULT, writer_words_DEFAULT, all_loaded_data_DEFAULT)

_wrds, all_word_mdn_Ws_DEFAULT, all_word_mdn_Cs_DEFAULT, _html, _file, _btn, _maxr, _maxs, mdn_svg_DEFAULT = update_mdn_word(" ".join(mdn_words_DEFAULT), scale_sd_DEFAULT, clamp_mdn_DEFAULT)
_sldr, char_Cs_DEFAULT, _chrs, _html, _file, _btn, _wght, char_svg_DEFAULT = update_blend_chars(*char_chosen_DEFAULT, char_weight_DEFAULT, char_Ws_DEFAULT)

with gr.Blocks() as demo:
    all_loaded_data_var = gr.State(all_loaded_data_DEFAULT)
    chosen_writers_var = gr.State(chosen_writers_DEFAULT)
    # data for writer interpolation
    writer_words_var = gr.State(writer_words_DEFAULT)
    writer_mean_Ws_var = gr.State(writer_mean_Ws_DEFAULT)
    writer_all_word_Ws_var = gr.State([e.detach() for e in writer_all_word_Ws_DEFAULT])
    writer_all_word_Cs_var = gr.State([e.detach() for e in writer_all_word_Cs_DEFAULT])
    writer_weight_var = gr.State(writer_weight_DEFAULT)
    writer_svg_var = gr.State(writer_svg_DEFAULT)
    # data for char interpolation
    char_chosen_var = gr.State(char_chosen_DEFAULT)
    char_mean_global_W_var = gr.State(char_mean_global_W_DEFAULT)
    char_weight_var = gr.State(char_weight_DEFAULT)
    char_Ws_var = gr.State(char_Ws_DEFAULT.detach())
    char_Cs_var = gr.State(char_Cs_DEFAULT.detach())
    char_svg_var = gr.State(char_svg_DEFAULT)
    # # data for MDN
    mdn_words_var = gr.State(mdn_words_DEFAULT)
    all_word_mdn_Ws_var = gr.State([e.detach() for e in all_word_mdn_Ws_DEFAULT])
    all_word_mdn_Cs_var = gr.State([e.detach() for e in all_word_mdn_Cs_DEFAULT])
    clamp_mdn_var = gr.State(clamp_mdn_DEFAULT)
    scale_sd_var = gr.State(scale_sd_DEFAULT)
    mdn_svg_var = gr.State(mdn_svg_DEFAULT)

    with gr.Tabs():
        with gr.TabItem("Blend Writers"):
            target_word = gr.Textbox(label="Target Word", value=" ".join(writer_words_DEFAULT), max_lines=1)
            with gr.Row():
                left_ratio_options = ["Style " + str(id) for i, id in enumerate(writer_options) if i % 2 == 0]
                right_ratio_options = ["Style " + str(id) for i, id in enumerate(writer_options) if i % 2 == 1]
                with gr.Column():
                    writer1 = gr.Radio(left_ratio_options, value="Style 120", label="Style for first writer")
                with gr.Column():
                    writer2 = gr.Radio(right_ratio_options, value="Style 80", label="Style for second writer")
            with gr.Row():
                writer_slider = gr.Slider(0, 1, value=writer_weight_DEFAULT, label="Style 120 vs. Style 80")
            with gr.Row():
                writer_default_image = update_writer_slider(writer_weight_DEFAULT, writer_words_DEFAULT, writer_all_word_Ws_DEFAULT, writer_all_word_Cs_DEFAULT)
                writer_output = gr.HTML(writer_default_image[0]["value"])
            with gr.Row():
                writer_download_btn = gr.Button("Save to SVG file")
                writer_download_btn.style(full_width="true")
                writer_download = gr.File(interactive=False, show_label=False, visible=False)

            writer_slider.change(fn=update_writer_slider, 
                inputs=[writer_slider, writer_words_var, writer_all_word_Ws_var, writer_all_word_Cs_var], 
                outputs=[writer_output, writer_download, writer_download_btn, writer_weight_var, writer_svg_var], show_progress=False)
            target_word.submit(fn=update_writer_word, 
                inputs=[target_word, writer_mean_Ws_var, all_loaded_data_var, writer_weight_var], 
                outputs=[writer_words_var, writer_all_word_Ws_var, writer_all_word_Cs_var, writer_output, writer_download, writer_download_btn, writer_weight_var, writer_svg_var], show_progress=False)
            writer1.change(fn=update_chosen_writers, 
                inputs=[writer1, writer2, writer_weight_var, writer_words_var, all_loaded_data_var], 
                outputs=[writer_slider, chosen_writers_var, writer_mean_Ws_var, writer_words_var, writer_all_word_Ws_var, writer_all_word_Cs_var, writer_output, writer_download, writer_download_btn, writer_weight_var, writer_svg_var])
            writer2.change(fn=update_chosen_writers, 
                inputs=[writer1, writer2, writer_weight_var, writer_words_var, all_loaded_data_var], 
                outputs=[writer_slider, chosen_writers_var, writer_mean_Ws_var, writer_words_var, writer_all_word_Ws_var, writer_all_word_Cs_var, writer_output, writer_download, writer_download_btn, writer_weight_var, writer_svg_var])
            writer_download_btn.click(fn=update_writer_download, 
                inputs=[writer_svg_var], 
                outputs=[writer_download, writer_download_btn])

        with gr.TabItem("Blend Characters"):
            with gr.Row():
                with gr.Column():
                    char1 = gr.Dropdown(choices=avail_char_list, value=char_chosen_DEFAULT[0], label="Character 1")
                with gr.Column():
                    char2 = gr.Dropdown(choices=avail_char_list, value=char_chosen_DEFAULT[1], label="Character 2")
            with gr.Row():
                char_slider = gr.Slider(0, 1, value=char_weight_DEFAULT, label=f"'{char_chosen_DEFAULT[0]}' vs. '{char_chosen_DEFAULT[1]}'")
            with gr.Row():
                char_default_image = update_char_slider(char_weight_DEFAULT, char_Ws_DEFAULT, char_Cs_DEFAULT, char_chosen_DEFAULT)
                char_output = gr.HTML(char_default_image[0]["value"])
            with gr.Row():
                char_download_btn = gr.Button("Save to SVG file")
                char_download_btn.style(full_width="true")
                char_download = gr.File(interactive=False, show_label=False, visible=False)

            char_slider.change(fn=update_char_slider, 
                inputs=[char_slider, char_Ws_var, char_Cs_var, char_chosen_var], 
                outputs=[char_output, char_download, char_download_btn, char_weight_var, char_svg_var], show_progress=False)

            char1.change(fn=update_blend_chars, 
                inputs=[char1, char2, char_weight_var, char_Ws_var], 
                outputs=[char_slider, char_Cs_var, char_chosen_var, char_output, char_download, char_download_btn, char_weight_var, char_svg_var])
            char2.change(fn=update_blend_chars, 
                inputs=[char1, char2, char_weight_var, char_Ws_var], 
                outputs=[char_slider, char_Cs_var, char_chosen_var, char_output, char_download, char_download_btn, char_weight_var, char_svg_var])

            char_download_btn.click(fn=update_char_download, 
                inputs=[char_svg_var], 
                outputs=[char_download, char_download_btn], show_progress=True)

        with gr.TabItem("Add Randomness"):
            mdn_word = gr.Textbox(label="Target Word", value=" ".join(mdn_words_DEFAULT), max_lines=1)
            with gr.Row():
                with gr.Column():
                    max_rand = gr.Slider(0, 1, value=clamp_mdn_DEFAULT, label="Maximum Randomness")
                with gr.Column():
                    scale_rand = gr.Slider(0, 3, value=scale_sd_DEFAULT, label="Scale of Randomness")
            with gr.Row():
                mdn_sample_button = gr.Button(value="Resample")
            with gr.Row():
                default_im = sample_mdn(scale_sd_DEFAULT, clamp_mdn_DEFAULT, mdn_words_DEFAULT, all_word_mdn_Ws_DEFAULT, all_word_mdn_Cs_DEFAULT)
                mdn_output = gr.HTML(default_im[0]["value"])
            with gr.Row():
                randomness_download_btn = gr.Button("Save to SVG file")
                randomness_download = gr.File(interactive=False, show_label=False, visible=False)

            max_rand.change(fn=sample_mdn, 
                inputs=[scale_rand, max_rand, mdn_words_var, all_word_mdn_Ws_var, all_word_mdn_Cs_var], 
                outputs=[mdn_output, randomness_download, randomness_download_btn, clamp_mdn_var, scale_sd_var, mdn_svg_var], show_progress=False)
            scale_rand.change(fn=sample_mdn, 
                inputs=[scale_rand, max_rand, mdn_words_var, all_word_mdn_Ws_var, all_word_mdn_Cs_var], 
                outputs=[mdn_output, randomness_download, randomness_download_btn, clamp_mdn_var, scale_sd_var, mdn_svg_var], show_progress=False)
            mdn_sample_button.click(fn=sample_mdn, 
                inputs=[scale_rand, max_rand, mdn_words_var, all_word_mdn_Ws_var, all_word_mdn_Cs_var], 
                outputs=[mdn_output, randomness_download, randomness_download_btn, clamp_mdn_var, scale_sd_var, mdn_svg_var], show_progress=False)
                
            mdn_word.submit(fn=update_mdn_word, 
                inputs=[mdn_word, scale_sd_var, clamp_mdn_var], 
                outputs=[mdn_words_var, all_word_mdn_Ws_var, all_word_mdn_Cs_var, mdn_output, randomness_download, randomness_download_btn, clamp_mdn_var, scale_sd_var, mdn_svg_var], show_progress=False)

            randomness_download_btn.click(fn=update_mdn_download, 
                inputs=[mdn_svg_var],
                outputs=[randomness_download, randomness_download_btn])
            randomness_download_btn.style(full_width="true")
demo.launch()