from __future__ import annotations

import pathlib
import subprocess
import tempfile

# import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import numpy as np

from nmreval.data.points import Points
from nmreval.io.asciireader import AsciiReader
from nmreval.io.hdfreader import HdfReader
from nmreval.lib.logger import logger
from nmreval.utils.utils import get_temperature, roundrobin


class FCReader:
    def __init__(self, fname: (list | str | pathlib.Path)):
        if isinstance(fname, (str, pathlib.Path)):
            self.fnames = [fname]
        else:
            self.fnames = fname

        self.temperatures = []
        self.data = {}
        self.filenames = {}
        self.t_params = {}
        self.f_params = {}

    def __call__(self, fname: (list | str | pathlib.Path)):
        if isinstance(fname, (str, pathlib.Path)):
            self.fnames = [fname]
        else:
            self.fnames = fname

        self.temperatures = []
        self.data = {}
        self.filenames = {}
        self.t_params = {}
        self.f_params = {}

    def load_magnetization(self, region: tuple = None, overwrite: bool = True):
        for filename in self.fnames:
            filename = pathlib.Path(filename)
            found_temperature = get_temperature(filename.stem)
            if found_temperature == -1:
                found_temperature = filename.stem

            if filename.is_file():
                if (region is None) or (region == (None, None)):
                    _temp = self._read_from_hdf(filename)
                else:
                    _temp = self._read_signals(filename, region)

            elif filename.is_dir():
                _temp = self._read_from_dir(filename)

            else:
                raise TypeError(f'{filename} is of unknown type')

            if not _temp:
                raise OSError(-666, f'No magnetization found for {filename.name}.', filename.name)

            self.data[found_temperature] = _temp
            self.filenames[found_temperature] = filename

            fname_no_ext = filename.with_suffix('')
            data_path = fname_no_ext.joinpath('data')
            # data_path.mkdir(parents=True, exist_ok=True)

            for k, v in sorted(_temp.items()):
                save_name = data_path.joinpath(f'{filename.stem}_{k:011.2f}'.replace('.', 'p') + '.dat')
                if save_name.exists() and (not overwrite):
                    continue
                v.savetxt(save_name)

    def _read_from_hdf(self, filename: (str | pathlib.Path)) -> dict:
        _temp = {}
        reader = HdfReader(filename)

        mag_directory = None
        for m in ['Magnetization', 'mag']:
            if m in reader:
                mag_directory = m
                break

        if m is not None:
            for mag in reader.get_selected(mag_directory, dtype='points'):
                _temp[mag.value] = mag

        return _temp

    @staticmethod
    def _read_from_dir(filename) -> dict:
        fname_no_ext = filename.with_suffix('')
        data_path = fname_no_ext / 'data'
        _temp = {}

        for mag in data_path.glob('*.dat'):
            d = AsciiReader(mag).export()
            for v in d:
                _temp[v.value] = v
                break

        return _temp

    def _read_signals(self, filename, region: tuple = None) -> dict:
        reader = HdfReader(filename)

        # This is one set with attributes to find default start:stop values (with some
        p = reader.parameters('/ABS_ACC_FID')
        start = p.get('stainte') or p.get('start') or 0
        stop = p.get('stointe') or p.get('stop') or 30e-5

        if region == (None, None) or region is None:
            region = (start, stop)

        if region[0] is None:
            region = (start, region[1])
        if region[1] is None:
            region = (region[0], stop)

        _temp = {}
        bevo_groups = reader['data'].children
        no_bevo = False
        for grp in bevo_groups.values():
            for k in grp.data():
                try:
                    # look for parameter, should be missing in older scripts
                    bevo = k.parameter['bevo']
                except KeyError:
                    no_bevo = True
                    break

                if str(k.parent) != 'ABS_ACC_FID_sig (group)':
                    continue

                tevo = k.parameter['tevo']

                s = reader.make_signal(k, flag='fid', value='tevo', group='bevo')
                pts = s.points([region])

                if bevo not in _temp:
                    _temp[bevo] = []

                _temp[bevo].append((tevo, *[np.real(pp[1]) for pp in pts]))

            if no_bevo:
                break

        if no_bevo:
            # hopefully only for old scripts
            sig = reader.get_selected('/data/B=*/ABS_ACC_FID*', dtype='signal')
            _temp = {}
            for s in sig:
                pts = s.points([region])
                b = s.group
                if b not in _temp:
                    _temp[b] = []

                _temp[b].append([s.value, *[np.real(pp[1]) for pp in pts]])

        for b, m in sorted(_temp.items()):
            m = np.array(m)
            _temp[b] = Points(x=m[:, 0], y=m[:, 1], value=b, name=f'B={b}').sort()

        return _temp

    def fit(self, kww: bool = True, save_fits: bool = True, save_fig: bool = True):
        if kww:
            bounds = ([-np.inf, -np.inf, 0.0, 0.0], [np.inf, np.inf, np.inf, np.inf])
        else:
            bounds = ([-np.inf, -np.inf, 0.0, 0.99999], [np.inf, np.inf, np.inf, 1.0])

        for temperature, filename in self.filenames.items():
            fname_no_ext = filename.with_suffix('')

            if save_fits:
                fit_path = fname_no_ext.joinpath('fit')
                fit_path.mkdir(parents=True, exist_ok=True)

            if save_fig:
                data_path = fname_no_ext.joinpath('data')
                image_path = fname_no_ext.joinpath('png')
                image_path.mkdir(parents=True, exist_ok=True)

            header = 'm0\tt1\tbeta\toff\n'

            params = []
            errors = []
            freqs = []

            for k, v in sorted(self.data[temperature].items()):
                # fit
                p0 = [v.y[0], v.y[-1]-v.y[0], v.x[int(0.5*len(v.x) - 0.5)], 1]
                try:
                    p0, pcov = curve_fit(FCReader.kww, v.x, v.y, p0=p0, bounds=bounds, max_nfev=500*len(v))
                except Exception as e:
                    logger.error(f'Fit of field {k} failed with exception', exc_info=e)
                    continue

                freqs.append(k)

                perr = np.sqrt(np.diag(pcov))
                params.append(p0)
                errors.append(perr)

                if isinstance(temperature, float):
                    new_entry = list(roundrobin([temperature], p0, perr))

                    try:
                        self.f_params[k].append(new_entry)
                    except KeyError:
                        self.f_params[k] = [new_entry]

                if save_fits or save_fig:
                    xplot = np.geomspace(v.x[0], v.x[-1], num=10*len(v.x))
                    yplot = FCReader.kww(xplot, *p0)
                    save_name = f'{filename.stem}_{k:011.2f}'.replace('.', 'p') + '.dat'

                    if save_fits:
                        np.savetxt(fit_path.joinpath(save_name), np.c_[xplot, yplot],
                                   header=header+'\t'.join([f'{p}+/-{err}' for p, err in zip(p0, perr)]))

                    if save_fig:
                        img_file = image_path.joinpath(save_name).with_suffix(".png")

                        gnuplot_args = [
                            'set terminal png;',
                            f'set output "{img_file}";',
                            f'set title "f = {k:.4g} Hz\\n{p0[2]:.4g}(+/-{perr[2]:.4g}) beta: {p0[3]:.4g}(+/-{perr[3]:.4g})";',
                            'set xlabel "t / s";',
                            'set logscale x;',
                            'set format x "10^{{%L}}";',
                            'set ylabel "M";',
                            'set key off;',
                            f'plot "{data_path.joinpath(save_name)}" with points pointtype 5, "{fit_path.joinpath(save_name)}" with lines;',
                        ]

                        try:
                            proc = subprocess.Popen(
                                ['gnuplot', '-p'],
                                shell=True,
                                stdin=subprocess.PIPE,
                                encoding='utf8',
                            )
                            for args in gnuplot_args:
                                proc.stdin.write(args)
                            proc.stdin.write('quit\n')
                            proc.stdin.flush()
                        except Exception as e:
                            logger.error(f'saving image {save_name} failed', e)

            freqs = np.asanyarray(freqs)
            params = np.asanyarray(params)
            errors = np.asanyarray(errors)

            # das ist jetzt eher haesslich
            self.t_params[temperature] = np.c_[freqs,
                                               params[:, 0], errors[:, 0], params[:, 1], errors[:, 1],
                                               params[:, 2], errors[:, 2], params[:, 3], errors[:, 3]]

        for k, val in self.f_params.items():
            val = np.array(val)
            np.nan_to_num(val)
            self.f_params[k] = val

    def write_parameter(self, path, kind):
        path = pathlib.Path(path)
        path.mkdir(parents=True, exist_ok=True)

        if kind == 'temp':
            _params = self.t_params
            fmt = '3.2f'
        else:
            _params = self.f_params
            fmt = '011.2f'

        save_path = path.joinpath(kind)
        if not save_path.exists():
            save_path.mkdir(parents=True)

        for key, par in _params.items():
            try:
                np.savetxt(save_path.joinpath(f'fitparameter_{key:{fmt}}.dat'), par,
                           header=f'{key}\nM0\tM0_err\tOff\tOff_err\tT1\tT1_err\tbeta\tbeta_err')
            except ValueError:
                np.savetxt(save_path.joinpath(f'fitparameter_{key}.dat'), par,
                           header=f'{key}\nM0\tM0_err\tOff\tOff_err\tT1\tT1_err\tbeta\tbeta_err')

    @staticmethod
    def _write_parameter(key, parameter, path, kind, fmt):
        save_path = path.joinpath(kind)
        save_path.mkdir(parents=True, exist_ok=True)
        try:
            np.savetxt(save_path.joinpath(f'fitparameter_{key:{fmt}}.dat'), parameter,
                       header=f'{key}\nM0\tM0_err\tOff\tOff_err\tT1\tT1_err\tbeta\tbeta_err')
        except ValueError:
            np.savetxt(save_path.joinpath(f'fitparameter_{key}.dat'), parameter,
                       header=f'{key}\nM0\tM0_err\tOff\tOff_err\tT1\tT1_err\tbeta\tbeta_err')

    # @staticmethod
    # def _plot_parameter(key, param, fig_mag, fig_t1, fig_beta):
    #     ax_mag = fig_mag.get_axes()[0]
    #     ax_t1 = fig_t1.get_axes()[0]
    #     ax_beta = fig_beta.get_axes()[0]
    #
    #     pl, = ax_mag.plot(param[:, 0], param[:, 1], 'o', label=key)
    #     ax_mag.plot(param[:, 0], param[:, 3], 's', color=pl.get_color())
    #     ax_t1.plot(param[:, 0], param[:, 5], 'o', label=key)
    #     ax_beta.plot(param[:, 0], param[:, 7], 'o', label=key)

    @staticmethod
    def _save_parameter_plot(path, kind, fig_mag, fig_t1, fig_beta):
        ax_mag = fig_mag.get_axes()[0]
        ax_t1 = fig_t1.get_axes()[0]
        ax_beta = fig_beta.get_axes()[0]

        for a in [ax_mag, ax_t1, ax_beta]:
            if kind == 'freq':
                a.legend(loc='upper left', bbox_to_anchor=(1, 1), ncol=2)
                a.set_xlabel('T / K')
            else:
                a.set_xscale('log')
                a.legend(loc='upper left', bbox_to_anchor=(1, 1))
                a.set_xlabel('f / Hz')

        ax_t1.set_yscale('log')
        ax_t1.set_ylabel('T1 / s')
        ax_beta.set_ylabel('beta')
        ax_mag.set_ylabel('M0 (squares), Offset (circles)')

        fig_beta.savefig(path.joinpath(f'beta_{kind}.png'), bbox_inches="tight")
        fig_mag.savefig(path.joinpath(f'mag_{kind}.png'), bbox_inches="tight")
        fig_t1.savefig(path.joinpath(f't1_{kind}.png'), bbox_inches="tight")
        # plt.close(fig_mag)
        # plt.close(fig_beta)
        # plt.close(fig_t1)

    # def plot_parameter(self, path, kind):
    #     path = pathlib.Path(path)
    #     path.mkdir(parents=True, exist_ok=True)
    #
    #     fig_mag, ax_mag = plt.subplots()
    #     fig_t1, ax_t1 = plt.subplots()
    #     fig_beta, ax_beta = plt.subplots()
    #
    #     if kind == 'temp':
    #         _params = self.t_params
    #     else:
    #         _params = self.f_params
    #
    #     save_path = path.joinpath(kind)
    #     if not save_path.exists():
    #         save_path.mkdir(parents=True)
    #
    #     for key, par in _params.items():
    #         pl, = ax_mag.plot(par[:, 0], par[:, 1], 'o', label=key)
    #         ax_mag.plot(par[:, 0], par[:, 3], 's', color=pl.get_color())
    #         ax_t1.plot(par[:, 0], par[:, 5], 'o', label=key)
    #         ax_beta.plot(par[:, 0], par[:, 7], 'o', label=key)
    #
    #     for a in [ax_mag, ax_t1, ax_beta]:
    #         if kind == 'freq':
    #             a.legend(loc='upper left', bbox_to_anchor=(1, 1), ncol=2)
    #             a.set_xlabel('T / K')
    #         else:
    #             a.set_xscale('log')
    #             a.legend(loc='upper left', bbox_to_anchor=(1, 1))
    #             a.set_xlabel('f / Hz')
    #
    #     ax_t1.set_yscale('log')
    #     ax_t1.set_ylabel('T1 / s')
    #     ax_beta.set_ylabel('beta')
    #     ax_mag.set_ylabel('M0 (squares), Offset (circles)')
    #
    #     fig_beta.savefig(path.joinpath(f'beta_{kind}.png'), bbox_inches="tight")
    #     fig_mag.savefig(path.joinpath(f'mag_{kind}.png'), bbox_inches="tight")
    #     fig_t1.savefig(path.joinpath(f't1_{kind}.png'), bbox_inches="tight")
    #     plt.close(fig_mag)
    #     plt.close(fig_beta)
    #     plt.close(fig_t1)

    def get_parameter(self, parameter='all', kind='freq', path=None, write=True, plot=True):
        param_list = []
        if isinstance(parameter, str):
            parameter = [parameter]

        for p in parameter:
            plower = p.lower()
            if plower == 'all':
                param_list = [(0, 'M0'), (1, 'Off'), (2, 'T1'), (3, 'beta')]
                break

            for i, name in [(0, 'M0'), (1, 'Off'), (2, 'T1'), (3, 'beta')]:
                if plower == name.lower():
                    param_list.append((i, name))

        if write or plot:
            if path is None:
                raise ValueError('Please specify a path to write to.')
            else:
                path = pathlib.Path(path)
                path.mkdir(parents=True, exist_ok=True)

        if kind == 'temp':
            _params = self.t_params
            fmt = '3.2f'
        else:
            _params = self.f_params
            fmt = '011.2f'

        # if plot:
        #     fig_mag, ax_mag = plt.subplots()
        #     fig_t1, ax_t1 = plt.subplots()
        #     fig_beta, ax_beta = plt.subplots()

        ret_val = []
        for key, par in _params.items():
            try:
                value = float(key)
            except ValueError:
                value = None

            if write:
                self._write_parameter(key, par, path, kind, fmt)

            # if plot:
            #     self._plot_parameter(key, par, fig_mag, fig_t1, fig_beta)

            for i, name in param_list:
                ret_val.append(Points(x=par[:, 0], y=par[:, 2*i+1], y_err=par[:, 2*i+2],
                                      name=f'{key} ({name})', value=value))

        # if plot:
        #     self._save_parameter_plot(path, kind, fig_mag, fig_t1, fig_beta)

        return ret_val

    @staticmethod
    def kww(x, m0, off, t1, beta):
        return m0 * np.exp(-(x/t1)**beta) + off
