from __future__ import annotations

import re
from collections import OrderedDict
from io import StringIO
from pathlib import Path
from typing import Any

import numpy as np
from scipy.stats import f as fdist
from scipy.interpolate import interp1d

from ._meta import MultiModel
from .model import Model
from .parameter import Parameter
from ..data.points import Points
from ..data.signals import Signal
from ..utils.text import convert


class FitResultCreator:
    @staticmethod
    def make_from_session(x_orig: np.ndarray, y_orig: np.ndarray, idx: int, kwargs: dict[Any]) -> FitResult:
        params = OrderedDict()

        for key, pbest, err in zip(kwargs['pnames'], kwargs['parameter'], kwargs['error']):
            params[key] = Parameter(pbest)
            params[key].error = err

        _x = kwargs['x']
        _y = kwargs['y']

        if len(_x) != len(x_orig):
            f = interp1d(_x, _y)
            rng = (x_orig >= np.min(_x)) * (x_orig <= np.max(_x))
            resid = f(x_orig[rng]) - y_orig[rng]
        else:
            resid = kwargs['y'] - y_orig

        stats = calc_statistics(resid, _y)

        return FitResult(
            x=kwargs['x'],
            y=kwargs['y'],
            x_data=x_orig,
            y_data=y_orig,
            params=params,
            fun_kwargs=dict(kwargs['choice']),
            resid=resid,
            nobs=0,
            nvar=0,
            we='',
            name=kwargs['name'],
            stats=stats,
            idx=idx,
        )

    @staticmethod
    def make_with_model(
            model: 'Model',
            x_orig: np.ndarray,
            y_orig: np.ndarray,
            p: list,
            fun_kwargs: dict,
            we: str,
            idx: str | None,
            nobs: int,
            nvar: int,
            corr:  np.ndarray,
            pcorr: np.ndarray,
            data_mode: int,
    ) -> FitResult:
        if np.all(x_orig > 0) and (np.max(x_orig) > 100 * np.min(x_orig)):
            islog = True
        else:
            islog = False

        if len(x_orig) < 51:
            if islog:
                _x = np.geomspace(np.min(x_orig), np.max(x_orig), num=10*x_orig.size-9)
            else:
                _x = np.linspace(np.min(x_orig), np.max(x_orig), num=10*x_orig.size-9)
        else:
            _x = x_orig

        try:
            pnames = model.pnames
        except AttributeError:
            pnames = model.params

        parameters = OrderedDict([(k, v) for k, v in zip(pnames, p)])
        p_final = [p.value for p in parameters.values()]

        resid = model.func(p_final, x_orig, **fun_kwargs) - y_orig

        actual_mode = -1
        if 'complex_mode' in fun_kwargs:
            actual_mode = fun_kwargs['complex_mode']
            fun_kwargs['complex_mode'] = 0

        _y = check_complex(model.func(p_final, _x, **fun_kwargs), actual_mode, data_mode)

        fun_kwargs['complex_mode'] = actual_mode

        stats = calc_statistics(_y, resid, nobs, nvar)
        varied = [p.var for p in parameters.values()]

        if corr is None:
            correlation = np.eye(len(p_final), len(p_final))
            partial_correlation = np.eye(len(p_final), len(p_final))

        else:
            if len(corr) < len(p_final):
                correlation = np.eye(len(p_final), len(p_final))
                partial_correlation = np.eye(len(p_final), len(p_final))
                # probably some fixed variables
                j = 0
                for i in range(len(corr)):
                    while not varied[j]:
                        j += 1
                    correlation[j, varied] = corr[i]
                    partial_correlation[j, varied] = pcorr[i]
                    j += 1
            else:
                correlation = corr
                partial_correlation = pcorr

        return FitResult(
            x=_x,
            y=_y,
            x_data=x_orig,
            y_data=y_orig,
            params=parameters,
            fun_kwargs=fun_kwargs,
            resid=resid,
            nobs=nobs,
            nvar=nvar,
            we=we,
            name=model.name,
            stats=stats,
            idx=idx,
            corr=correlation,
            pcorr=partial_correlation,
            islog=islog,
            func=model,
            data_complex=data_mode,
        )


class FitResult(Points):

    def __init__(
            self: FitResult,
            x: np.ndarray,
            y: np.ndarray,
            x_data: np.ndarray,
            y_data: np.ndarray,
            params: dict,
            fun_kwargs: dict,
            resid: np.ndarray,
            nobs: int,
            nvar: int,
            we: str,
            name: str,
            stats: dict,
            idx: int = None,
            corr: np.ndarray = None,
            pcorr: np.ndarray = None,
            islog: bool = False,
            func=None,
            data_complex: int = 1,
            **kwargs,
    ):

        self.parameter, name = self._prepare_names(params, name)
        label = kwargs.get('label', name)
        super().__init__(x=x, y=y, name=label, **kwargs)

        self.residual = resid
        self.idx = idx
        self.statistics = stats
        self.nobs = nobs
        self.nvar = nvar
        self.we = we
        self.fun_kwargs = fun_kwargs
        self.correlation = corr
        self.partial_correlation = pcorr
        self.islog = islog
        self.iscomplex = np.iscomplexobj(self.y)
        self.x_data = x_data
        self.y_data = y_data
        self._model_name = name
        self._func = func
        self._data_complex = data_complex

    @staticmethod
    def _prepare_names(parameter: dict, modelname: str):
        pattern = re.compile(r'(\\?\w[\\\w .-]*(?:_{[\w\\ .-]})?)(\(\d+\))')

        split_funcs = {g2: g1 for (g1, g2) in pattern.findall(modelname)}

        parameter_dic = {}
        for pname, pvalue in parameter.items():
            nice_name = pname
            nice_func = ''
            m = pattern.match(pname)
            if m:
                func_number = m.group(2)
                nice_name = m.group(1)
                if func_number in split_funcs:
                    nice_func = split_funcs[func_number]
            pvalue.fix()
            pvalue.name = nice_name
            pvalue.function = nice_func
            parameter_dic[pname] = pvalue

        modelname = re.sub(r'\(\d+\)', '', modelname)
        if modelname[0] == '(' and modelname[-1] == ')':
            modelname = modelname[1:-1]

        return parameter_dic, modelname

    @property
    def model_name(self):
        return self._model_name

    def __repr__(self):
        try:
            return 'Fit: ' + self.name
        except AttributeError:
            return 'FitObject'

    @property
    def func(self):
        if isinstance(self._func, MultiModel):
            return self._func.func
        else:
            return self._func

    @property
    def p_final(self):
        return [pp.value for pp in self.parameter.values()]

    @property
    def dof(self):
        return self.nobs-self.nvar

    def pprint(self, statistics=True, correlations=True):
        s = StringIO()
        s.write('Fit result:\n')
        s.write(f'  model : {self.name}\n')
        s.write(f'  #data : {self.nobs}\n')
        s.write(f'  #var  : {self.nvar}\n')
        s.write(f'  #weight: {self.we}\n')
        s.write('\nParameter\n')
        s.write(self.parameter_string())

        if statistics:
            s.write('\nStatistics\n')
            for k, v in self.statistics.items():
                s.write(f'  {k} : {v:.4f}\n')

        if correlations and self.correlation is not None:
            s.write('\nCorrelation (partial corr.)\n')
            s.write(self._correlation_string())
            s.write('\n')

        print(s.getvalue())

    def parameter_string(self):
        ret_val = ''

        for pkey, pval in self.parameter.items():
            ret_val += convert(str(pval), old='tex', new='str') + '\n'

        if self.fun_kwargs:
            for k, v in self.fun_kwargs.items():
                ret_val += f'  {k}: {v}\n'

        return ret_val

    def _correlation_string(self):
        ret_val = ''
        for p_i, p_j, corr_ij, pcorr_ij in self.correlation_list():
            ret_val += f"  {convert(p_i, old='tex', new='str')} / {convert(p_j, old='tex', new='str')} : {corr_ij:.4f} ({pcorr_ij:.4f})\n"
        return ret_val

    def correlation_list(self, limit=0.1):
        correlations = []

        if self.correlation is not None:
            pnames = list(self.parameter.keys())

            corr = np.triu(self.correlation, k=1)

            for i, j in zip(*np.unravel_index(np.argsort(np.abs(corr), axis=None)[::-1],
                                              self.correlation.shape)):
                if i == j:
                    continue

                if abs(corr[i, j]) < limit:
                    break

                correlations.append((pnames[i], pnames[j], self.correlation[i, j], self.partial_correlation[i, j]))

        return correlations

    def savetxt(self, fname: str | Path, err: bool = True):
        header = self.name + '\n'
        for pval in self.parameter.values():
            header += '{}\t{}\t{}\n'.format(convert(pval.name, old='tex', new='str'), pval.value, pval.error)
        if self.fun_kwargs:
            for k, v in self.fun_kwargs.items():
                header += '{}\t{}\n'.format(convert(k, old='tex', new='str'),
                                            convert(str(v), old='tex', new='str'))

        if self.iscomplex == 'complex':
            np.savetxt(fname, np.c_[self.x, self.y.real, self.y.imag], header=header)
        else:
            np.savetxt(fname, np.c_[self.x, self.y], header=header)

    def save_parameter(self, fname: str | Path, label: str = None, overwrite: bool = False):
        path = Path(fname)
        if not path.is_file():
            overwrite = True

        writemode = 'w' if overwrite else 'a'

        if label is None:
            label = self.value

        with path.open(writemode) as f:
            if overwrite or not path.exists():
                f.write('# xvalue(1)\t')
                for i, pname in enumerate(self.parameter.keys()):
                    raw_name = convert(pname, old='tex', new='str')
                    f.write(f'{raw_name}({2*i+2})\t{raw_name}_err({2*i+3})\t')
                f.write('\n')

            f.write(str(label).replace(' ', '') + '\t')
            for p in self.parameter.values():
                if p.error is not None:
                    err = p.error
                else:
                    err = 0.
                f.write(f'{p.value:.8e}\t{err:.8e}\t')

            f.write('# ')
            if self.fun_kwargs:
                for k, v in self.fun_kwargs.items():
                    f.write(f"{convert(k, old='tex', new='str')}: {convert(str(v), old='tex', new='str')}\t")

            f.write(f'weight: {self.we}\t')

            f.write('\n')
            f.write(f'# line above from: {self.name} (model: {self.model_name})')
            f.write('\n')

    def f_test(self, chi2: float, dof: float):
        if 'red. chi^2' not in self.statistics or dof == self.dof:
            f_value = 1e318
        else:
            f_value = (chi2-self.statistics['chi^2']) / (dof-self.dof) / self.statistics['red. chi^2']

        try:
            prob_f = 1-fdist.cdf(f_value, dof-self.dof, self.dof)
        except:
            prob_f = 0

        return f_value, prob_f

    def get_state(self):
        state = super().get_state()

        for attr in ['idx', 'fun_kwargs', 'nobs', 'nvar', 'we', 'islog', 'iscomplex', 'x_data', 'y_data']:
            state[attr] = getattr(self, attr)

        state['name'] = self._model_name
        state['label'] = self.name
        state['corr'] = self.correlation
        state['pcorr'] = self.partial_correlation
        state['stats'] = self.statistics
        state['resid'] = self.residual
        state['params'] = {k: v.get_state() for k, v in self.parameter.items()}

        state['mode'] = 'fit'

        return state

    @staticmethod
    def set_state(state, **kwargs):
        state['params'] = {k: Parameter.set_state(v) for k, v in state.pop('params').items()}
        if 'we' not in state:
            state['we'] = '?'
        data = FitResult(**state)

        return data

    def with_new_x(self, x_values):
        if self.func is None:
            raise ValueError('no fit function available to calculate new y values')

        new_fit = self.copy()
        fun_kwargs = {k: v for k, v in self.fun_kwargs.items()}
        if self.fun_kwargs.get('complex_mode', -1) == -1:
            fun_kwargs.pop('complex_mode', None)
        y_values = self.func.func(self.p_final, x_values, **fun_kwargs)
        y_values = check_complex(y_values, self.fun_kwargs.get('complex_mode', -1), self._data_complex)

        new_fit.set_data(x_values, y_values, y_err=0.0)

        return new_fit

    def sub(self, x_values):
        if self.func is None:
            raise ValueError('no fit function available to calculate new y values')

        part_functions = []
        actual_mode = self.fun_kwargs.get('complex_mode', -1)
        fun_kwargs = {k: v for k, v in self.fun_kwargs.items()}
        if self.fun_kwargs.get('complex_mode', -1) == -1:
            fun_kwargs.pop('complex_mode', None)

        for sub_name, sub_y in zip(self.func.sub_name(), self.func.sub(self.p_final, x_values, **self.fun_kwargs)):
            sub_y = check_complex(sub_y, actual_mode, self._data_complex)

            if np.iscomplexobj(sub_y):
                part_functions.append(Signal(x_values, sub_y, name=sub_name))
            else:
                part_functions.append(Points(x_values, sub_y, name=sub_name))

        if actual_mode > 0:
            self.fun_kwargs['complex_mode'] = actual_mode

        return part_functions


def check_complex(y, model_complex, data_complex):
    if not np.iscomplexobj(y):
        return y

    if model_complex == 1:
        y.imag = 0
        if data_complex == 1:
            y = y.real
    elif model_complex == 2:
        y.real = 0
        if data_complex == 1:
            y = y.imag

    return y


def calc_statistics(y, residual, nobs=None, nvar=None):
    chi = (residual**2).sum()
    try:
        r = 1 - chi/((y-np.mean(y))**2).sum()
    except RuntimeWarning:
        r = -9999

    if nobs is None:
        nobs = 1

    if nvar is None:
        nvar = 0

    dof = nobs - nvar
    loglikehood = nobs * np.log(chi / nobs)

    stats = {
        'chi^2': chi,
        'R^2': r,
        'AIC': loglikehood + 2 * nvar,
        'BIC': loglikehood + np.log(nobs) * nvar,
        'adj. R^2': 1 - (nobs-1) / (dof+1e-13) * (1-r),
        'red. chi^2': chi / (dof + 1e-13),
    }

    stats['AICc'] = stats['AIC'] + 2*(nvar+1)*nvar / (dof - 1 + 1e-13)

    return stats
