# Bokeh imports
from bokeh.layouts import row, layout
from bokeh.models.widgets import TextInput,  MultiSelect, Button, CheckboxGroup, Slider, PreText, Div
from bokeh.models import HoverTool
from bokeh.plotting import figure

import bokeh.palettes as palettes

# Other imports
from os import listdir
from os.path import isfile
import re
import numpy as np


class FilterSelector:
    def __init__(self, files):
        self.text_input = TextInput(value="", title="Regex filenmae filter:")
        self.text_input.on_change("value", self.ti_handler)
        self.multiselector = MultiSelect(title="Select files to plot:", options=[(file, file) for file in files])
        self.renderer = row(self.text_input, self.multiselector)
        self.files = files

    def ti_handler(self, attr, old, new):
        pattern = new
        filtered_files = [file for file in self.files if re.search(pattern, file)]
        self.multiselector.options = [(file, file) for file in filtered_files]

    def value(self):
        return self.multiselector.value


def results_widget(data_path):
    def update_plot():
        file_names = selector.value()
        # file_names = ["exps_m_thetao_accuracies.txt", 'exps_m_heatc_accuracies.txt']
        colors = palettes.cividis(len(file_names))
        for index, filename in enumerate(file_names):
            x, y = load_file_data(filename)
            color = colors[index]

            # Computing sizes
            diff = np.array(y[:-1]) - np.array(y[1:])
            positive_mask = diff > 0
            diff[diff == 0] = 1
            exponents = np.log10(np.abs(diff))
            exponents = (exponents + np.abs(min(exponents)) + 1) * 5
            exponents[np.logical_not(positive_mask)] = 1
            size = list(exponents) + [1]

            line = f.line(x, y, name=filename, legend=filename, color=color, alpha=0.3, hover_line_alpha=1,
                          line_width=2, hover_line_color=color)
            circle = f.circle(x, y, name=filename, legend=filename, color=color, size=size, alpha=0.3, hover_fill_alpha=1,
                            hover_fill_color=color, selection_color="firebrick",)

            line.hover_glyph.line_width = 4
        f.legend.location = "top_left"
        f.legend.click_policy = "hide"
        hover = HoverTool()
        hover.tooltips = [
            ("name", "$name"),
            ("x", "$x"),
            ("y", "$y")]
        f.add_tools(hover)

    def open_file(filename):
        filepath = "%s/%s" % (data_dir, filename)
        with open(filepath) as f:
            filedata = f.read()
        return filedata

    def filter_data(file_data):
        data = [l for l in file_data.split("\n")]
        x = [from_file_path_to_index(l.split()[0]) for l in data]
        y = [float(l.split()[1]) for l in data]
        return x, y

    def load_file_data(file_name):
        filedata = open_file(file_name)
        x, y = filter_data(filedata)
        return x, y

    def from_file_path_to_index(file_path):
        m = re.search("[a-z]([0-9]+)", file_path)
        return int(m.group(1))

    def threshold_handler(attr, old, new):
        for index, var_name in enumerate(var_names):
            if index in new:
                threshold_lines[var_name].visible = True
            else:
                threshold_lines[var_name].visible = False

    def diff_handler(attr, old , new):
        namelist_path = "/home/Earth/otintopr/RPE-Canal/AutoRPE/a1an_Analysis/IndividualVariables/"

        pattern = "([a-z])([0-9]+)"
        m = re.findall(pattern, new)
        if m:
            if len(m) == 1:
                exp, exp_number = m[0]
                exp_number = int(exp_number)
                try:
                    diffs = namelistDiff(namelist_path, exp, exp_number)
                except AssertionError:
                    diff_text.text = "Files not found!:\n%s" % new
                    return
                name = "%s%i vs %s%i" % (exp, exp_number, exp, exp_number+1)
            elif len(m) == 2:
                exp1, exp_number1 = m[0]
                exp_number1 = int(exp_number1)
                exp2, exp_number2 = m[1]
                exp_number2 = int(exp_number2)
                try:
                    diffs = namelistDiff(namelist_path, exp1, exp_number1, exp2, exp_number2)
                except AssertionError:
                    diff_text.text = "Files not found!:\n%s" % new
                    return
                name = "%s%i vs %s%i" % (exp1, exp_number1, exp2, exp_number2)
            else:
                diff_text.text = "Case not found!:\n%s" % new
                return

            colors = [precisions for ids, precisions in diffs.items()]
            colors = ["red" if a < b else "green" for a, b in colors]

            ids = [_id for _id, precision in diffs.items()]
            ids_colors = zip(ids, colors)

            text_elements = ['<font color="%s">%s</font>' % (color, _id) for _id , color in ids_colors]

            formated_variables = " ".join(text_elements)

            to_print = "%s<br>%s" % (name, formated_variables)
            diff_text.text = to_print
        else:
            diff_text.text = "Case not found!:\n%s" % new

    data_dir = data_path
    files = listdir(data_dir)
    selector = FilterSelector(files)

    button = Button(label="Plot Files!")
    button.on_click(update_plot)

    f = figure(y_axis_type="log", tools="tap, box_zoom, reset, save, pan")

    var_names = ["thetao", "so", "heatc", "saltc"]

    thresholds = {"thetao": 1.49*10**-6, "so": 6.41*10**-6, "heatc": 2.74*10**-6, "saltc": 5.47*10**-9}
    threshold_selector = CheckboxGroup(labels=var_names, active=[])
    threshold_selector.on_change("active", threshold_handler)

    threshold_lines = {}
    for var_name in var_names:
        threshold_lines[var_name] = f.line([1, 51], [thresholds[var_name], thresholds[var_name]], name=var_name+" threshold")
        threshold_lines[var_name].visible = False


    diff_input = TextInput(title="Check namelist differences (1 or 2 experiment ids)")
    diff_input.on_change("value", diff_handler)

    diff_text = Div()
    diff_handler(None, None, "n30 n31")

    return layout([selector.renderer, button], [f, threshold_selector], [diff_input], [diff_text], toolbar_location="below")


def load_namelist_precisions(namelist_path):
    pattern = "emulator_variable_precisions\(([0-9]+)\) = ([0-9]+)"
    namelist = {}
    with open(namelist_path) as f:
        data = [l for l in f]
        for l in data:
            m = re.search(pattern, l)
            if m:
                var_id = int(m.group(1))
                precision = int(m.group(2))
                namelist[var_id] = precision
    return namelist


def namelistDiff(data_path, exp, case, exp2=None, case2=None):

    namelist1 = "namelist_precisions_%s%i" % (exp, case)
    namelist1_path = "%s/%s" % (data_path, namelist1)
    assert isfile(namelist1_path)

    if exp2 is not None and case2 is not None:
        namelist2 = "namelist_precisions_%s%i" % (exp2, case2)
        namelist2_path = "%s/%s" % (data_path, namelist2)
        assert isfile(namelist2_path)
    else:
        case = case + 1
        namelist2 = "namelist_precisions_%s%i" % (exp, case)
        namelist2_path = "%s/%s" % (data_path, namelist2)

    assert isfile(namelist2_path)

    n1 = load_namelist_precisions(namelist1_path)
    n2 = load_namelist_precisions(namelist2_path)

    differences = {}
    for key in n1.keys():
        if n1[key] != n2[key]:
            differences[key] = (n1[key], n2[key])
    return differences
