import os
import numpy as np
import matplotlib.pyplot as plt
from PyQt5 import QtWidgets

import ana_cont.continuation as cont
from gui.gui_backend import RealFrequencyGrid, InputData, TextInputData, OutputData
from gui.maxent_ui import Ui_MainWindow


class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
    """The Main Window of the graphical user interface.

    The class MainWindow inherits from Ui_MainWindow, which is
    defined in maxent_ui.py. The latter file is autogenerated
    by pyuic from maxent_ui.ui [`pyuic5 maxent_ui.ui -o maxent_ui.py`]
    The ui file can be edited by the QtDesigner.
    """

    def __init__(self, *args, obj=None, **kwargs):
        """Connect the widgets, instantiate the main classes."""
        super(MainWindow, self).__init__(*args, **kwargs)
        self.setupUi(self)

        # real-frequency grid
        self.realgrid = RealFrequencyGrid(wmax=float(self.max_real_freq.text()),
                                          nw=int(self.num_real_freq.text()),
                                          type=str(self.grid_type_combo.currentText()))
        self.connect_realgrid_button()
        self.connect_wmax()
        self.connect_nw()
        self.connect_grid_type()

        # input data
        self.input_data = InputData(fname=str(self.inp_file_name.text()),
                                    iter_type=str(self.iteration_type_combo.currentText()),
                                    iter_num=str(self.iteration_number.text()),
                                    data_type=str(self.inp_data_type.currentText()),
                                    atom=str(self.atom_number.text()),
                                    orbital=str(self.orbital_number.text()),
                                    spin=str(self.spin_type_combo.currentText()),
                                    num_mats=str(self.num_mats_freq.text()))
        self.connect_select_button()
        self.connect_load_button()
        self.connect_show_button()
        self.connect_load_button_text()
        self.connect_show_button_2()
        self.connect_select_button_2()

        # text display field and "Do it" button
        self.text_output.setReadOnly(True)
        self.connect_doit_button()

        # output data
        self.output_data = OutputData()
        self.connect_select_output_button()
        self.connect_save_button()


    def connect_realgrid_button(self):
        self.gen_real_grid_button.clicked.connect(lambda: self.realgrid.create_grid())

    def connect_wmax(self):
        self.max_real_freq.returnPressed.connect(
            lambda: self.realgrid.update_wmax(float(self.max_real_freq.text())))
        self.max_real_freq.editingFinished.connect(
            lambda: self.realgrid.update_wmax(float(self.max_real_freq.text())))

    def connect_nw(self):
        self.num_real_freq.returnPressed.connect(
            lambda: self.realgrid.update_nw(int(self.num_real_freq.text())))
        self.num_real_freq.editingFinished.connect(
            lambda: self.realgrid.update_nw(int(self.num_real_freq.text())))

    def connect_grid_type(self):
        self.grid_type_combo.activated.connect(
            lambda: self.realgrid.update_type(str(self.grid_type_combo.currentText()))
        )

    def preset_fnames(self, fname):
        self.inp_file_name.setText(fname)
        self.inp_file_name_2.setText(fname)

    def connect_fname_input(self):
        self.inp_file_name.editingFinished.connect(
            lambda: self.input_data.update_fname(str(self.inp_file_name.text())))
        self.inp_file_name.textChanged.connect(
            lambda: self.input_data.update_fname(str(self.inp_file_name.text())))

    def get_fname(self):
        self.inp_file_name.setText(
            QtWidgets.QFileDialog.getOpenFileName(self,
                    'Open file', os.getcwd(), "HDF5 files (*.hdf5)")[0])

    def connect_select_button(self):
        self.select_file_button.clicked.connect(self.get_fname)

    def get_fname_text(self):
        self.inp_file_name_2.setText(
            QtWidgets.QFileDialog.getOpenFileName(self,
                    'Open file', os.getcwd(), "text files (*.dat *.txt)")[0])

    def connect_select_button_2(self):
        self.select_file_button_2.clicked.connect(self.get_fname_text)

    def connect_data_type(self):
        self.inp_data_type.activated.connect(
            lambda: self.input_data.update_data_type(str(self.inp_data_type.currentText())))

    def connect_iteration_type(self):
        self.iteration_type_combo.activated.connect(
            lambda: self.input_data.update_iter_type(str(self.iteration_type_combo.currentText())))

    def connect_iteration_number(self):
        self.iteration_number.editingFinished.connect(
            lambda: self.input_data.update_iter_num(str(self.iteration_number.text())))

    def connect_atom(self):
        self.atom_number.editingFinished.connect(
            lambda: self.input_data.update_atom(int(self.atom_number.text())))

    def connect_orbital(self):
        self.orbital_number.editingFinished.connect(
            lambda: self.input_data.update_orbital(int(self.orbital_number.text())))

    def connect_spin(self):
        self.spin_type_combo.activated.connect(
            lambda: self.input_data.update_spin(str(self.spin_type_combo.currentText())))

    def connect_num_mats(self):
        self.num_mats_freq.editingFinished.connect(
            lambda: self.input_data.update_num_mats(int(self.num_mats_freq.text()))
        )

    def connect_show_button(self):
        self.show_data_button.clicked.connect(
            lambda: self.input_data.plot()
        )

    def connect_show_button_2(self):
        self.show_data_button_2.clicked.connect(
            lambda: self.input_data.plot()
        )

    def load_w2dynamics_data(self):
        self.input_data = InputData(fname=str(self.inp_file_name.text()),
                                    iter_type=str(self.iteration_type_combo.currentText()),
                                    iter_num=str(self.iteration_number.text()),
                                    data_type=str(self.inp_data_type.currentText()),
                                    atom=str(self.atom_number.text()),
                                    orbital=str(self.orbital_number.text()),
                                    spin=str(self.spin_type_combo.currentText()),
                                    num_mats=str(self.num_mats_freq.text()),
                                    ignore_real_part=self.ignore_checkbox.isChecked())
        self.input_data.load_data()

    def connect_load_button(self):
        self.load_data_button.clicked.connect(self.load_w2dynamics_data)

    def load_text_data(self):
        self.input_data = TextInputData(fname=str(self.inp_file_name_2.text()),
                                        data_type=str(self.inp_data_type_text.currentText()),
                                        n_skip=str(self.n_skip.text()),
                                        num_mats=str(self.num_mats_freq_text.text()))
        self.input_data.read_data()

    def connect_load_button_text(self):
        self.load_data_button_2.clicked.connect(self.load_text_data)

    def get_preblur(self):
        preblur_checked = self.preblur_checkbox.isChecked()
        try:
            bw = float(self.blur_width.text()) if preblur_checked else 0.
        except ValueError:
            print('Invalid input for blur width, setting to 0.')
            bw = 0.
        preblur = preblur_checked and bw > 0.
        return preblur, bw

    def main_function(self):
        """Main function for the analytic continuation procedure.

        This function is called when the "Do it" button is clicked.
        It performs an analytical continuation for the present settings
        and shows a plot.
        """
        self.ana_cont_probl = cont.AnalyticContinuationProblem(im_axis=self.input_data.mats,
                                                               im_data=self.input_data.value,
                                                               re_axis=self.realgrid.grid,
                                                               kernel_mode='freq_fermionic')
        model = np.ones_like(self.realgrid.grid)
        model /= np.trapz(model, self.realgrid.grid)

        preblur, bw = self.get_preblur()

        sol = self.ana_cont_probl.solve(method='maxent_svd',
                                        optimizer='newton',
                                        alpha_determination='chi2kink',
                                        model=model,
                                        stdev=self.input_data.error,
                                        interactive=False, alpha_start=1e14, alpha_end=1e-3,
                                        preblur=preblur, blur_width=bw)

        inp_str = 'atom {}, orb {}, spin {}, blur {}: '.format(self.input_data.atom,
                                                               self.input_data.orbital,
                                                               self.input_data.spin,
                                                               bw)
        all_chis = np.isfinite(np.array([s.chi2 for s in sol[1]]))
        res_str = 'alpha_opt={:3.2f}, chi2(alpha_opt)={:3.2f}, min(chi2)={:3.2f}'.format(
            sol[0].alpha, sol[0].chi2, np.amin(all_chis)
        )
        self.text_output.append(inp_str + res_str)
        alphas = [s.alpha for s in sol[1]]
        chis = [s.chi2 for s in sol[1]]

        self.output_data.update(self.realgrid.grid, sol[0].A_opt, self.input_data)

        fig, ax = plt.subplots(ncols=2, nrows=2, figsize=(11.75, 8.25))  # A4 paper size
        ax[0, 0].loglog(alphas, chis, marker='s', color='black')
        ax[0, 0].loglog(sol[0].alpha, sol[0].chi2, marker='*', color='red', markersize=15)
        ax[0, 0].set_xlabel(r'$\alpha$')
        ax[0, 0].set_ylabel(r'$\chi^2(\alpha)$')

        ax[1, 0].plot(self.realgrid.grid, sol[0].A_opt)
        ax[1, 0].set_xlabel(r'$\omega$')
        ax[1, 0].set_ylabel('spectrum')

        ax[0, 1].plot(self.input_data.mats, self.input_data.value.real,
                      color='blue', ls=':', marker='x', markersize=5,
                      label='Re[data]')
        ax[0, 1].plot(self.input_data.mats, self.input_data.value.imag,
                      color='green', ls=':', marker='+', markersize=5,
                      label='Im[data]')
        ax[0, 1].plot(self.input_data.mats, sol[0].backtransform.real,
                      ls='--', color='gray', label='Re[fit]')
        ax[0, 1].plot(self.input_data.mats, sol[0].backtransform.imag,
                      color='gray', label='Im[fit]')
        ax[0, 1].set_xlabel(r'$\nu_n$')
        ax[0, 1].set_ylabel(self.input_data.data_type)
        ax[0, 1].legend()

        ax[1, 1].plot(self.input_data.mats, (self.input_data.value - sol[0].backtransform).real,
                      ls='--', label='real part')
        ax[1, 1].plot(self.input_data.mats, (self.input_data.value - sol[0].backtransform).imag,
                      label='imaginary part')
        ax[1, 1].set_xlabel(r'$\nu_n$')
        ax[1, 1].set_ylabel('data $-$ fit')
        ax[1, 1].legend()
        plt.tight_layout()
        plt.show()

    def connect_doit_button(self):
        self.doit_button.clicked.connect(lambda: self.main_function())

    def connect_fname_output(self):
        self.out_file_name.editingFinished.connect(
            lambda: self.output_data.update_fname(str(self.out_file_name.text())))
        self.inp_file_name.textChanged.connect(
            lambda: self.output_data.update_fname(str(self.out_file_name.text())))

    def get_fname_output(self):
        fname_out = QtWidgets.QFileDialog.getSaveFileName(self,
                    'Save as', '/'.join(self.input_data.fname.split('/')[:-1]), "DAT files (*.dat)")[0]
        self.out_file_name.setText(fname_out)
        self.output_data.update_fname(fname_out)

    def connect_select_output_button(self):
        self.output_directory_button.clicked.connect(self.get_fname_output)

    def save_output(self):
        fname_out = str(self.out_file_name.text())
        if fname_out == '':
            print('Error in saving: First you have to specify the output file name.')
            return 1

        self.output_data.update_fname(fname_out)
        try:
            self.output_data.save(interpolate=self.interpolate_checkbox.isChecked(),
                                 n_reg=int(self.n_interpolation.text()))
        except AttributeError:
            print('Error in saving: First you have to specify the output file name.')

    def connect_save_button(self):
        self.save_button.clicked.connect(lambda: self.save_output())

