#from matplotlib import pyplot as plt
from kivy.uix.popup import Popup
from clease.gui.fittingAlgorithmEditors import LassoEditor, L2Editor, BCSEditor
from clease.gui.fittingAlgorithmEditors import GAEditor, FitAlgEditor
from clease.gui.load_save_dialog import LoadDialog
from kivy.app import App
import json
from clease.gui.util import parse_max_cluster_dia
from clease.gui.constants import BACKGROUND_COLOR, FOREGROUND_TEXT_COLOR
from clease.gui.constants import ECI_GRAPH_COLORS
from threading import Thread
from kivy.uix.screenmanager import Screen
from kivy.utils import get_color_from_hex
from kivy_garden.graph import Graph, ScatterPlot, BarPlot, LinePlot
import numpy as np


class ECIOptimiser(object):
    fit_page = None
    evaluator = None

    def optimise(self):
        try:
            self.fit_page.eci = self.evaluator.get_eci_dict()

            e_ce = self.evaluator.cf_matrix.dot(self.evaluator.eci)

            self.fit_page.e_ce = e_ce
            self.fit_page.e_dft = self.evaluator.e_dft
            self.fit_page.e_pred_leave_out = self.evaluator.e_pred_loo
            cv = self.evaluator.get_cv_score()
            rmse = self.evaluator.rmse()*1000.0
            mae = self.evaluator.mae()*1000.0
            self.fit_page.set_cv(cv)
            self.fit_page.set_rmse(rmse)
            self.fit_page.set_mae(mae)
            App.get_running_app().root.ids.status.text = 'Idle'
            self.fit_page.fitting_in_progress = False
            self.fit_page.update_plots()
        except Exception as exc:
            App.get_running_app().root.ids.status.text = str(exc)
            self.fit_page.fitting_in_progress = False
            return


class GAClusterSelector(object):
    fit_page = None
    kwargs = None
    settings = None
    _pop_up = None

    def run(self):
        from clease import GAFit, LinearRegression, Evaluate

        try:
            ga = GAFit(self.settings, **self.kwargs)
            _ = ga.run()

            optimiser = ECIOptimiser()
            optimiser.fit_page = self.fit_page

            max_cluster_dia = self.kwargs['max_cluster_dia']
            max_cluster_size = self.kwargs['max_cluster_size']
            select_cond = self.kwargs['select_cond']
            min_weight = self.kwargs['min_weight']
            optimiser.evaluator = Evaluate(
                self.settings, max_cluster_dia=max_cluster_dia,
                max_cluster_size=max_cluster_size,
                select_cond=select_cond, min_weight=min_weight,
                fitting_scheme=LinearRegression())
            optimiser.optimise()
        except Exception as exc:
            App.get_running_app().root.ids.status.text.text = str(exc)


class FitPage(Screen):
    graphs_added = False
    fitting_params = {}
    eci = {}
    fit_result = {}
    _pop_up = None
    fitting_in_progress = False
    e_dft = None
    e_ce = None

    energy_plot = None
    eci_plots = []
    energy_graph = None
    eci_graph = None
    zero_line_energy = None

    def on_enter(self):
        if not self.graphs_added:
            self.energy_graph = Graph(
                xlabel='DFT energy (eV/atom)',
                ylabel="E_CE - E_DFT (meV/atom)",
                x_ticks_minor=0,
                x_ticks_major=10,
                y_ticks_major=10,
                y_grid_label=True,
                x_grid_label=True,
                padding=5,
                xlog=False,
                ylog=False,
                xmin=0.0,
                ymin=0.0,
                precision='%.2f')
            self.energy_plot = ScatterPlot(color=FOREGROUND_TEXT_COLOR,
                                           point_size=3)
            self.zero_line_energy = LinePlot(line_width=2,
                                             color=FOREGROUND_TEXT_COLOR)
            self.energy_graph.add_plot(self.energy_plot)
            self.energy_graph.add_plot(self.zero_line_energy)
            self.ids.energyPlot.add_widget(self.energy_graph)

            self.eci_graph = Graph(
                ylabel='ECI (eV/atom)',
                xlabel='Clusters',
                x_ticks_minor=0,
                x_ticks_major=10,
                y_ticks_major=10,
                y_grid_label=True,
                x_grid_label=True,
                padding=5,
                xlog=False,
                ylog=False,
                xmin=-1.0,
                ymin=0.0,
                precision='%.2f'
            )

            for i in range(8):
                color = ECI_GRAPH_COLORS[i % len(ECI_GRAPH_COLORS)]
                color = get_color_from_hex(color)
                plot = BarPlot(color=color, bar_width=4, bar_spacing=1)
                self.eci_plots.append(plot)
                self.eci_graph.add_plot(plot)

            self.ids.eciPlot.add_widget(self.eci_graph)
            self.graphs_added = True

    def dismiss_popup(self):

        if isinstance(self._pop_up.content, FitAlgEditor):
            self._pop_up.content.backup()
        self._pop_up.dismiss()
        self._pop_up = None

    def show_lasso_editor(self):
        content = LassoEditor(close=self.close_lasso_editor)
        self._pop_up = Popup(title="LASSO Editor", content=content,
                             pos_hint={'right': 0.95, 'top': 0.95},
                             size_hint=(0.9, 0.9))
        self._pop_up.open()

    def show_l2_editor(self):
        content = L2Editor(close=self.close_l2_editor)
        self._pop_up = Popup(title="L2 Editor", content=content,
                             pos_hint={'right': 0.95, 'top': 0.95},
                             size_hint=(0.9, 0.9))
        self._pop_up.open()

    def show_bcs_editor(self):
        content = BCSEditor(close=self.close_bcs_editor)

        self._pop_up = Popup(title="BCS Editor", content=content,
                             pos_hint={'right': 0.95, 'top': 0.95},
                             size_hint=(0.9, 0.9))
        self._pop_up.open()

    def show_ga_editor(self):
        content = GAEditor(close=self.close_ga_editor)

        self._pop_up = Popup(title="Genetic Algorithm Editor", content=content,
                             pos_hint={'right': 0.95, 'top': 0.95},
                             size_hint=(0.9, 0.9))
        self._pop_up.open()

    def close_l2_editor(self, alpha):
        self.fitting_params = {}
        self.fitting_params['algorithm'] = 'L2'
        self.fitting_params['alpha'] = float(alpha)
        self.dismiss_popup()

    def close_lasso_editor(self, alpha):
        self.fitting_params = {}
        self.fitting_params['algorithm'] = 'LASSO'
        self.fitting_params['alpha'] = float(alpha)
        self.dismiss_popup()

    def close_bcs_editor(self, shape_var, rate_var, shape_lamb, var_opt_start,
                         init_lamb, lamb_opt_start, max_iter, noise):
        self.fitting_params = {
            'algorithm': 'BCS',
            'shape_var': float(shape_var),
            'rate_var': float(rate_var),
            'shape_lamb': float(shape_lamb),
            'var_opt_start': float(var_opt_start),
            'init_lamb': float(init_lamb),
            'lamb_opt_start': float(lamb_opt_start),
            'max_iter': int(max_iter),
            'noise': float(noise)
        }
        self.dismiss_popup()

    def close_ga_editor(self, elitism, mut_prob, num_individuals, max_active,
                        cost_func, sparsity, sub_clust):
        self.fitting_params = {
            'algorithm': 'GA',
            'elitism': int(elitism),
            'mut_prob': float(mut_prob),
            'num_individuals': int(num_individuals),
            'max_active': int(max_active),
            'cost_func': cost_func,
            'sparsity': float(sparsity),
            'sub_clust': sub_clust == 'Yes'
        }
        self.dismiss_popup()

    def show_load_ECI_file_dialog(self):
        content = LoadDialog(load=self.load_eci_file,
                             cancel=self.dismiss_popup)

        self._pop_up = Popup(title="Load ECI file", content=content,
                             pos_hint={'right': 0.95, 'top': 0.95},
                             size_hint=(0.9, 0.9))
        self._pop_up.open()

    def load_eci_file(self, path, filename):
        if len(filename) == 0:
            fname = path
        else:
            fname = filename[0]

        self.ids.eciFileInput.text = fname
        self.dismiss_popup()

    def launch_fit_alg_editor(self):
        if self.ids.fitAlgSpinner.text == 'LASSO':
            self.show_lasso_editor()
        elif self.ids.fitAlgSpinner.text == 'L2':
            self.show_l2_editor()
        elif self.ids.fitAlgSpinner.text == 'BCS':
            self.show_bcs_editor()
        elif self.ids.fitAlgSpinner.text == 'Genetic Algorithm':
            self.show_ga_editor()

    def save_eci(self):
        fname = self.ids.eciFileInput.text

        if not fname.endswith('.json'):
            msg = 'ECI file should be a JSON file'
            App.get_running_app().root.ids.status.text = msg

        try:
            with open(fname, 'w') as out:
                json.dump(self.eci, out, separators=(',', ': '), indent=2)
            msg = 'ECIs saved to {}'.format(fname)
            App.get_running_app().root.ids.status.text = msg
        except Exception as exc:
            App.get_running_app().root.ids.status.text = str(exc)

    def fit_eci(self):
        if self.fitting_in_progress:
            # We already optimising ECIs
            return

        from clease import Evaluate
        settings = App.get_running_app().root.settings

        if settings is None:
            msg = 'Settings not set. Call Update settings first.'
            App.get_running_app().root.ids.status.text = msg
            return

        if self.ids.maxClusterDiaCut.text == '':
            max_cluster_dia_cut = None
        else:
            try:
                max_cluster_dia_cut = \
                    parse_max_cluster_dia(self.ids.maxClusterDiaCut.text)
            except Exception as exc:
                App.get_running_app().root.ids.status.text = str(exc)
                return

        if self.ids.maxClusterSizeCut.text == '':
            max_cluster_size_cut = None
        else:
            max_cluster_size_cut = int(self.ids.maxClusterSizeCut.text)

        k_fold = self.ids.kFoldInput.text
        if k_fold == '':
            msg = 'K-fold has to be given'
            App.get_running_app().root.ids.status.text = msg
            return
        k_fold = int(k_fold)

        num_rep = self.ids.numRepititionsInput.text
        if num_rep == '':
            num_rep = 1
        else:
            num_rep = int(num_rep)

        scheme = self.fitting_params.get('algorithm', None)
        if scheme is None:
            msg = 'Open the fitting scheme editor prior to the fit'
            App.get_running_app().root.ids.status.text = msg
            return

        scheme = scheme.lower()
        alpha = 0.0
        if scheme in ['lasso', 'l2']:
            alpha = self.fitting_params['alpha']
        elif scheme == 'bcs':
            from clease import BayesianCompressiveSensing
            scheme = BayesianCompressiveSensing(
                shape_var=self.fitting_params['shape_var'],
                rate_var=self.fitting_params['rate_var'],
                shape_lamb=self.fitting_params['shape_lamb'],
                variance_opt_start=self.fitting_params['var_opt_start'],
                lamb_opt_start=self.fitting_params['lamb_opt_start'],
                maxiter=self.fitting_params['max_iter'],
                noise=self.fitting_params['noise'],
                init_lamb=self.fitting_params['init_lamb']
            )
        elif scheme == 'ga':
            ga_runner = GAClusterSelector()
            ga_runner.fit_page = self
            ga_runner.kwargs = {
                'max_cluster_size': max_cluster_size_cut,
                'max_cluster_dia': max_cluster_dia_cut,
                'select_cond': None,
                'min_weight': 1.0,
                'mutation_prob': self.fitting_params['mut_prob'],
                'elitism': self.fitting_params['elitism'],
                'num_individuals': self.fitting_params['num_individuals'],
                'max_num_in_init_pool': self.fitting_params['max_active'],
                'cost_func': self.fitting_params['cost_func'].lower(),
                'sparsity_slope': self.fitting_params['sparsity'],
                'include_subclusters': self.fitting_params['sub_clust']
            }
            ga_runner.settings = settings

            # GA behaves a bit different from the other schems
            # therefore we have a separate runner and return
            # after the runner is finished...
            msg = 'Selecting clusters with GA..'
            App.get_running_app().root.ids.status.text = msg
            Thread(target=ga_runner.run).start()
            return

        if k_fold == -1:
            scoring_scheme = 'loocv'
        else:
            scoring_scheme = 'k-fold'

        try:
            evaluator = Evaluate(
                settings, fitting_scheme=scheme, alpha=alpha,
                max_cluster_size=max_cluster_size_cut,
                max_cluster_dia=max_cluster_dia_cut, nsplits=k_fold,
                num_repetitions=num_rep, scoring_scheme=scoring_scheme)

            App.get_running_app().root.ids.status.text = 'Optimizing ECIs...'
            eci_optimiser = ECIOptimiser()
            eci_optimiser.fit_page = self
            eci_optimiser.evaluator = evaluator
            self.fitting_in_progress = True

            Thread(target=eci_optimiser.optimise).start()
        except Exception as exc:
            App.get_running_app().root.ids.status.text = str(exc)

    def set_cv(self, cv):
        self.ids.cvLabel.text = 'CV: {:.3f} meV/atom'.format(cv)

    def set_rmse(self, rmse):
        self.ids.rmseLabel.text = 'RMSE: {:.3f} meV/atom'.format(rmse)

    def set_mae(self, mae):
        self.ids.maeLabel.text = 'MAE: {:.3f} meV/atom'.format(mae)

    def _eci_has_been_fitted(self):
        return self.e_dft is not None and self.e_ce is not None

    def update_energy_plot(self, e_dft, e_ce):
        if not self._eci_has_been_fitted():
            msg = 'ECIs has not been fitted yet'
            App.get_running_app().root.ids.status.text = msg
            return

        xmin = np.min(e_dft)
        xmax = np.max(e_dft)
        ymin = np.min(e_ce - e_dft)*1000.0
        ymax = np.max(e_ce - e_dft)*1000.0

        x_rng = xmax - xmin
        xmin -= 0.05*x_rng
        xmax += 0.05*x_rng

        y_rng = ymax - ymin
        ymin -= 0.05*y_rng
        ymax += 0.05*y_rng

        self.energy_graph.xmin = float(xmin)
        self.energy_graph.xmax = float(xmax)
        self.energy_graph.ymin = float(ymin)
        self.energy_graph.ymax = float(ymax)
        self.energy_graph.x_ticks_major = float(xmax - xmin)/10.0
        self.energy_graph.y_ticks_major = float(ymax - ymin)/10.0
        self.energy_graph.y_grid_label = True
        self.energy_graph.x_grid_label = True
        self.energy_plot.points = [(y, (x - y)*1000.0)
                                   for x, y in zip(e_ce, e_dft)]
        self.zero_line_energy.points = [(xmin, 0.0), (xmax, 0.0)]

    def update_eci_plot(self, eci):
        eci_by_size = {}

        for k, v in eci.items():
            size = int((k[1]))
            if size < 2:
                continue
            if size not in eci_by_size.keys():
                eci_by_size[size] = []
            eci_by_size[size].append(v)

        sorted_keys = sorted(list(eci_by_size.keys()))
        prev = 0

        xmax = len(list(eci.keys())) + len(sorted_keys)
        ymin = min([v for _, v in eci.items()])
        ymax = max([v for _, v in eci.items()])
        for i, k in enumerate(sorted_keys):
            values = eci_by_size[k]
            indx = range(prev, prev + len(values))
            self.eci_plots[i].points = list(zip(indx, values))
            prev = indx[-1]+2

        y_rng = ymax - ymin
        ymin -= 0.05*y_rng
        ymax += 0.05*y_rng

        self.eci_graph.xmax = int(xmax)
        self.eci_graph.ymin = float(ymin)
        self.eci_graph.ymax = float(ymax)
        self.eci_graph.x_ticks_major = float(xmax)/10.0
        self.eci_graph.y_ticks_major = float(ymax - ymin)/10.0

    def update_plots(self):
        if self.fitting_in_progress:
            msg = 'Performing ECI opmisation in the background. '
            msg += 'Wait until the process is finished.'
            App.get_running_app().root.ids.status.text = msg

        self.update_energy_plot(self.e_dft, self.e_ce)
        self.update_eci_plot(self.eci)

    def to_dict(self):
        return {
            'eci_file': self.ids.eciFileInput.text,
            'db_select_cond': self.ids.dbSelectCondInput.text,
            'max_cluster_size': self.ids.maxClusterSizeCut.text,
            'max_cluster_dia': self.ids.maxClusterDiaCut.text,
            'k_fold': self.ids.kFoldInput.text,
            'num_repetitions': self.ids.numRepititionsInput.text,
            'fit_alg': self.ids.fitAlgSpinner.text
        }

    def from_dict(self, data):
        self.ids.eciFileInput.text = data.get('eci_file', '')
        self.ids.dbSelectCondInput.text = data.get('db_select_cond', '')
        self.ids.maxClusterSizeCut.text = data.get('max_cluster_size', '')
        self.ids.maxClusterDiaCut.text = data.get('max_cluster_dia', '')
        self.ids.kFoldInput.text = data.get('k_fold', '10')
        self.ids.numRepititionsInput.text = data.get('num_repetitions', '1')
        self.ids.fitAlgSpinner.text = data.get('fit_alg', 'LASSO')
