import warnings
from itertools import product

import numpy as np
import scipy.linalg as la
from scipy import optimize
import scipy.odr as odr

from . import EPS
from .data import Data
from .model import Model
from .parameter import Parameters
from .result import FitResultCreator

__all__ = ['FitRoutine', 'FitAbortException']


class FitAbortException(Exception):
    pass


class FitRoutine(object):
    def __init__(self, mode='lsq'):
        self._fitmethod = mode
        self.data = []
        self.fit_model = None
        self._no_own_model = []
        self.parameter = Parameters()
        self.result = []
        self.linked = []
        self._abort = False

    def add_data(self, x, y=None, we=None, idx=None):
        if isinstance(x, Data):
            d = x

        else:
            _x = np.asarray(x)
            if _x.ndim == 2 and _x.shape[1] == 2:
                d = Data(*x, we=we)

            elif y is None:
                raise ValueError('First argument must be of type Data, 2d-array, '
                                 'or second argument must be given')

            else:
                d = Data(x, y, we=we, idx=idx)

        if idx is not None:
            d.idx = idx
        d.minimizer = self
        self.data.append(d)
        self.result.append(None)

        return d

    def remove_data(self, data):
        try:
            idx = self.data.index(data)
            _ = self.data.pop(idx)
            self.result.pop(idx)

        except ValueError:
            raise IndexError(f'Data {data} not found')

    def set_model(self, func, *args, idx=None, **kwargs):
        if isinstance(func, Model):
            model = func
        else:
            model = Model(func, *args, **kwargs)

        if idx is not None:
            for i, data in enumerate(self.data):
                if data.idx == idx:
                    data.set_model(model)
        else:
            self.fit_model = model

        return self.fit_model

    def set_link_parameter(self, parameter: tuple, replacement: tuple):
        if isinstance(replacement[0], Model):
            if replacement[1] not in replacement[0].global_parameter:
                raise KeyError(f'Parameter at pos {replacement[1]} of '
                               f'model {str(replacement[0])} is not global')

        if isinstance(parameter[0], Model):
            warnings.warn(f'Replaced parameter at pos {parameter[1]} in {str(parameter[0])} '
                          f'becomes global with linkage.')

        self.linked.append((*parameter, *replacement))

    def prepare_links(self):
        self._no_own_model = []
        self.parameter = Parameters()
        _found_models = {}
        linked_sender = {}

        for v in self.data:
            linked_sender[v] = set()
            self.parameter.update(v.parameter.copy())

            # set temporaray model
            if v.model is None:
                v.model = self.fit_model
                self._no_own_model.append(v)

            # register model
            if v.model not in _found_models:
                _found_models[v.model] = []
                m_param = v.model.parameter.copy()
                self.parameter.update(m_param)

            _found_models[v.model].append(v)

            if v.model not in linked_sender:
                linked_sender[v.model] = set()

        linked_parameter = {}
        for par, par_parm, repl, repl_par in self.linked:
            if isinstance(par, Data):
                if isinstance(repl, Data):
                    linked_parameter[par.para_keys[par_parm]] = repl.para_keys[repl_par]
                else:
                    linked_parameter[par.para_keys[par_parm]] = repl.global_parameter[repl_par]

            else:
                if isinstance(repl, Data):
                    par.global_parameter[par_parm] = repl.para_keys[repl_par]
                else:
                    par.global_parameter[par_parm] = repl.global_parameter[repl_par]

            linked_sender[repl].add(par)
            linked_sender[par].add(repl)

        for mm, m_data in _found_models.items():
            if mm.global_parameter:
                for dd in m_data:
                    linked_sender[mm].add(dd)
                    linked_sender[dd].add(mm)

        coupled_data = []
        visited_data = []
        for s in linked_sender.keys():
            if s in visited_data:
                continue
            sub_graph = []
            self.find_paths(s, linked_sender, sub_graph, visited_data)
            if sub_graph:
                coupled_data.append(sub_graph)

        return coupled_data, linked_parameter

    def find_paths(self, start, graph, coupled_nodes=None, visited_nodes=None):
        visited_nodes.append(start)
        if isinstance(start, Data):
            coupled_nodes.append(start)

        for neighbor in graph[start]:
            if neighbor in visited_nodes:
                continue

            self.find_paths(neighbor, graph, coupled_nodes, visited_nodes)

    def abort(self):
        print('ABORT ???')
        self._abort = True

    def run(self, mode='lsq'):
        self._abort = False
        self.parameter = Parameters()

        fit_groups, linked_parameter = self.prepare_links()

        for data_groups in fit_groups:
            if len(data_groups) == 1 and not self.linked:
                data = data_groups[0]
                # get variable parameter for fitter
                p0_k, lb_k, ub_k, var_pars_k = self._prep_data(data)

                if mode == 'lsq':
                    self._least_squares_single(data, p0_k, lb_k, ub_k, var_pars_k)

                elif mode == 'nm':
                    self._nm_single(data, p0_k, lb_k, ub_k, var_pars_k)

                elif mode == 'odr':
                    # ODR takes no bounds
                    self._odr_single(data, p0_k, var_pars_k)

            else:
                data_pars, p0, lb, ub, var_pars = self._prep_global(data_groups, linked_parameter)

                if mode == 'lsq':
                    self._least_squares_global(data_groups, p0, lb, ub, var_pars, data_pars)

                elif mode == 'nm':
                    self._nm_global(data_groups, p0, lb, ub, var_pars, data_pars)

                elif mode == 'odr':
                    self._odr_global(data_groups, p0, var_pars, data_pars)

        self.unprep_run()

        return self.result

    def _prep_data(self, data):
        if data.get_model() is None:
            data._model = self.fit_model
            self._no_own_model.append(data)

        return self._prep_parameter(data.parameter)

    @staticmethod
    def _prep_parameter(parameter):
        vals = []
        var_pars = []
        for p_k, v_k in parameter.items():
            if v_k.var:
                vals.append([v_k.scaled_value, v_k.lb / v_k.scale, v_k.ub / v_k.scale])
                var_pars.append(p_k)

        pp, lb, ub = zip(*vals)

        return pp, lb, ub, var_pars

    def _prep_global(self, data_group, linked):
        p0 = []
        lb = []
        ub = []
        var = []
        data_pars = []

        # loopyloop over data that belong to one fit (linked or global)
        for data in data_group:
            actual_pars = []
            for i, (p_k, v_k) in enumerate(data.parameter.items()):
                p_k_used = p_k
                v_k_used = v_k

                # is parameter replaced by global parameter?
                if i in data.model.global_parameter:
                    p_k_used = data.model.global_parameter[i]
                    v_k_used = self.parameter[p_k_used]

                # links trump global parameter
                if p_k_used in linked:
                    p_k_used = linked[p_k_used]
                    v_k_used = self.parameter[p_k_used]

                actual_pars.append(p_k_used)
                # parameter is variable and was not found before as shared parameter
                if v_k_used.var and p_k_used not in var:
                    p0.append(v_k_used.scaled_value)
                    lb.append(v_k_used.lb / v_k_used.scale)
                    ub.append(v_k_used.ub / v_k_used.scale)
                    var.append(p_k_used)

            data_pars.append(actual_pars)

        return data_pars, p0, lb, ub, var

    def unprep_run(self):
        for d in self._no_own_model:
            d._model = None

        self._no_own_model = []

    # COST FUNCTIONS: f(x) - y (least_square, minimize), and f(x) (ODR)
    def __cost_scipy(self, p, data, varpars, used_pars):
        for keys, values in zip(varpars, p):
            self.parameter[keys].scaled_value = values

        actual_parameters = [self.parameter[keys].value for keys in used_pars]
        return data.cost(actual_parameters)

    def __cost_odr(self, p, data, varpars, used_pars):
        for keys, values in zip(varpars, p):
            self.parameter[keys].scaled_value = values

        actual_parameters = [self.parameter[keys].value for keys in used_pars]

        return data.func(actual_parameters, data.x)

    def __cost_scipy_glob(self, p, data, varpars, used_pars):
        # replace values
        for keys, values in zip(varpars, p):
            self.parameter[keys].scaled_value = values

        r = []
        # unpack parameter and calculate y values and concatenate all
        for values, p_idx in zip(data, used_pars):
            actual_parameters = [self.parameter[keys].value for keys in p_idx]
            r = np.r_[r, values.cost(actual_parameters)]

        return r

    def __cost_odr_glob(self, p, data, varpars, used_pars):
        # replace values
        for keys, values in zip(varpars, p):
            self.parameter[keys].scaled_value = values

        r = []
        # unpack parameter and calculate y values and concatenate all
        for values, p_idx in zip(data, used_pars):
            actual_parameters = [self.parameter[keys].value for keys in p_idx]
            r = np.r_[r, values.func(actual_parameters, values.x)]

        return r

    def _least_squares_single(self, data, p0, lb, ub, var):
        def cost(p):
            if self._abort:
                raise FitAbortException(f'Fit aborted by user')

            return self.__cost_scipy(p, data, var, data.para_keys)

        with np.errstate(all='ignore'):
            res = optimize.least_squares(cost, p0, bounds=(lb, ub), max_nfev=1000 * len(p0))

        err, corr, partial_corr = self._calc_error(res.jac, np.sum(res.fun**2), *res.jac.shape)
        self.make_results(data, res.x, var, data.para_keys, res.jac.shape,
                          err=err, corr=corr, partial_corr=partial_corr)

    def _least_squares_global(self, data, p0, lb, ub, var, data_pars):
        def cost(p):
            if self._abort:
                raise FitAbortException(f'Fit aborted by user')
            return self.__cost_scipy_glob(p, data, var, data_pars)

        with np.errstate(all='ignore'):
            res = optimize.least_squares(cost, p0, bounds=(lb, ub), max_nfev=1000 * len(p0))

        err, corr, partial_corr = self._calc_error(res.jac, np.sum(res.fun**2), *res.jac.shape)
        for v, var_pars_k in zip(data, data_pars):
            self.make_results(v, res.x, var, var_pars_k, res.jac.shape,
                              err=err, corr=corr, partial_corr=partial_corr)

    def _nm_single(self, data, p0, lb, ub, var):
        def cost(p):
            if self._abort:
                raise FitAbortException(f'Fit aborted by user')
            return (self.__cost_scipy(p, data, var, data.para_keys)**2).sum()

        with np.errstate(all='ignore'):
            res = optimize.minimize(cost, p0, bounds=[(b1, b2) for (b1, b2) in zip(lb, ub)],
                                    method='Nelder-Mead', options={'maxiter': 1000 * len(p0)})

        self.make_results(data, res.x, var, data.para_keys, (len(data), len(p0)))

    def _nm_global(self, data, p0, lb, ub, var, data_pars):
        def cost(p):
            if self._abort:
                raise FitAbortException(f'Fit aborted by user')
            return (self.__cost_scipy_glob(p, data, var, data_pars)**2).sum()

        with np.errstate(all='ignore'):
            res = optimize.minimize(cost, p0, bounds=[(b1, b2) for (b1, b2) in zip(lb, ub)],
                                    method='Nelder-Mead', options={'maxiter': 1000 * len(p0)})

        for v, var_pars_k in zip(data, data_pars):
            self.make_results(v, res.x, var, var_pars_k, (sum(len(d) for d in data), len(p0)))

    def _odr_single(self, data, p0, var_pars):
        odr_data = odr.Data(data.x, data.y)

        def func(p, _):
            if self._abort:
                raise FitAbortException(f'Fit aborted by user')
            return self.__cost_odr(p, data, var_pars, data.para_keys)

        odr_model = odr.Model(func)

        o = odr.ODR(odr_data, odr_model, beta0=p0)
        res = o.run()

        corr = res.cov_beta / (res.sd_beta[:, None] * res.sd_beta[None, :]) * res.res_var
        try:
            corr_inv = np.linalg.inv(corr)
            corr_inv_diag = np.diag(np.sqrt(1 / np.diag(corr_inv)))
            partial_corr = -1. * np.dot(np.dot(corr_inv_diag, corr_inv), corr_inv_diag)  # Partial correlation matrix
            partial_corr[np.diag_indices_from(partial_corr)] = 1.
        except np.linalg.LinAlgError:
            partial_corr = corr

        self.make_results(data, res.beta, var_pars, data.para_keys, (len(data), len(p0)),
                          err=res.sd_beta, corr=corr, partial_corr=partial_corr)

    def _odr_global(self, data, p0, var, data_pars):
        def func(p, _):
            if self._abort:
                raise FitAbortException(f'Fit aborted by user')
            return self.__cost_odr_glob(p, data, var, data_pars)

        x = []
        y = []
        for d in data:
            x = np.r_[x, d.x]
            y = np.r_[y, d.y]

        odr_data = odr.Data(x, y)
        odr_model = odr.Model(func)

        o = odr.ODR(odr_data, odr_model, beta0=p0, ifixb=var)
        res = o.run()

        corr = res.cov_beta / (res.sd_beta[:, None] * res.sd_beta[None, :]) * res.res_var
        try:
            corr_inv = np.linalg.inv(corr)
            corr_inv_diag = np.diag(np.sqrt(1 / np.diag(corr_inv)))
            partial_corr = -1. * np.dot(np.dot(corr_inv_diag, corr_inv), corr_inv_diag)  # Partial correlation matrix
            partial_corr[np.diag_indices_from(partial_corr)] = 1.
        except np.linalg.LinAlgError:
            partial_corr = corr

        for v, var_pars_k in zip(data, data_pars):
            self.make_results(v, res.beta, var, var_pars_k, (sum(len(d) for d in data), len(p0)),
                              err=res.sd_beta, corr=corr, partial_corr=partial_corr)

    def make_results(self, data, p, var_pars, used_pars, shape,
                     err=None, corr=None, partial_corr=None):

        if err is None:
            err = [0] * len(p)

        # update parameter values
        for keys, p_value, err_value in zip(var_pars, p, err):
            self.parameter[keys].scaled_value = p_value
            self.parameter[keys].scaled_error = err_value

        combinations = list(product(var_pars, var_pars))
        actual_parameters = []
        corr_idx = []

        for i, p_i in enumerate(used_pars):
            actual_parameters.append(self.parameter[p_i])
            for j, p_j in enumerate(used_pars):
                try:
                    # find the position of the parameter combinations
                    corr_idx.append(combinations.index((p_i, p_j)))
                except ValueError:
                    pass

        # reshape the correlation matrices
        if corr is None:
            actual_corr = None
            actual_pcorr = None
        else:
            indexes = np.unravel_index(corr_idx, corr.shape)
            dim_one = int(np.sqrt(len(corr_idx)))
            actual_corr = corr[indexes].reshape((dim_one, -1))
            actual_pcorr = partial_corr[indexes].reshape((dim_one, -1))

        idx = self.data.index(data)
        model = data.get_model()

        self.result[idx] = FitResultCreator.make_with_model(model, data.x, data.y,
                                                            actual_parameters, data.fun_kwargs, data.idx,
                                                            *shape, corr=actual_corr, pcorr=actual_pcorr)

        return self.result

    @staticmethod
    def _calc_error(jac, chi, nobs, nvars):
        # copy of scipy.curve_fit to calculate covariance
        # noinspection PyTupleAssignmentBalance
        _, s, vt = la.svd(jac, full_matrices=False)
        threshold = EPS * max(jac.shape) * s[0]
        s = s[s > threshold]
        vt = vt[:s.size]
        pcov = np.dot(vt.T / s**2, vt) * chi / (nobs - nvars)

        if pcov is None:
            _err = np.zeros(nvars)
            corr = np.zeros((nvars, nvars))
        else:
            _err = np.sqrt(np.diag(pcov))
            corr = pcov / (_err[:, None] * _err[None, :])

        corr = corr.astype(np.float64)
        try:
            corr_inv = np.linalg.inv(corr)
            corr_inv_diag = np.diag(np.sqrt(1 / np.diag(corr_inv)))
            partial_corr = -1. * np.dot(np.dot(corr_inv_diag, corr_inv), corr_inv_diag)  # Partial correlation matrix
            partial_corr[np.diag_indices_from(partial_corr)] = 1.
        except np.linalg.LinAlgError:
            partial_corr = corr

        return _err, corr, partial_corr
