# Source code for qspec.models._base

    # -*- coding: utf-8 -*-
    """
    qspec.models._base

    Created on 14.03.2022

    @author: Patrick Mueller

    Base classes for lineshape models.
    """

    import numpy as np

    from qspec.tools import merge_intervals

    __all__ = ['Model', 'Empty', 'NPeak', 'Offset', 'Amplifier', 'Custom', 'YPars', 'Listed', 'Summed', 'Linked']

    np_version = np.version.version.split('.')

    def _is_unc(fix):
        if isinstance(fix, float):
            return True
        if not isinstance(fix, str):
            return False

        j, k = fix.find('('), fix.find(')')
        try:
            if j == -1 or k == -1:
                raise ValueError('Not a value with uncertainty.')
            _ = '{}({})'.format(float(fix[:j]), float(fix[j+1:k]))
        except ValueError:
            return False
        return True

    def _val_fix_to_val(val, fix):
        if isinstance(fix, float):
            return val
        j = fix.find('(')
        return float(fix[:j])

    def _fix_to_unc(fix):
        if isinstance(fix, float):
            return fix
        j, k = fix.find('('), fix.find(')')
        return float(fix[j+1:k])

    def _args_ordered(args, order):
        return [args[i] for i in order]

    def _poly(x, *args):
        return np.sum([args[n] * x ** n for n in range(len(args))], axis=0)

    class Model:
        """
        Base class for all models.
        """

        def __init__(self, model=None):
            self.model = model
            self.type = 'Model'

        def __call__(self, x, *args, **kwargs):
            return self.evaluate(x, *self.update_args(args), **kwargs)

        def evaluate(self, x, *args, **kwargs):  # Reimplement this function in subclasses.
            pass

        def _add_arg(self, name, val, fix, link):
            if name in self.names:
                raise ValueError('Parameter {} already exists.'.format(name))
            self.names.append(name)
            self.vals.append(val)
            self.fixes.append(fix)
            self.links.append(link)
            self.expressions.append('args[{}]'.format(self._index))

            self.p[name] = self._index

            self._index += 1
            self._size += 1

        @property
        def description(self):
            label = ''
            super_model = self
            while super_model is not None:
                if isinstance(super_model, Listed):
                    label += super_model.type + '[0].'
                    super_model = super_model.models[0]
                else:
                    label += super_model.type + '.'
                    super_model = super_model.model
            return label[:-1]

        @property
        def size(self):
            """
            :returns: The number of parameters required by the model.
            """
            return self._size

        @property
        def dx(self):
            return 0.1 if self.model is None else self.model.dx

        @property
        def error(self):
            return self._error

        @error.setter
        def error(self, value):
            self._error = value

        @property
        def model(self):
            return self._model

        @model.setter
        def model(self, value):
            self._model = value
            if self._model is None:
                self.names, self.vals, self.fixes, self.links = [], [], [], []
                self.expressions = []
                self.p = {}
                self._index = 0
                self._size = 0
                self._error = ''
            else:
                self.names, self.vals, self.fixes, self.links = \
                    self._model.names, self._model.vals, self._model.fixes, self._model.links
                self.expressions = self._model.expressions
                self.p = self._model.p
                self._index = len(self._model.names)
                self._size = len(self._model.names)
                self._error = self._model.error

        def get_pars(self):
            return zip(self.names, self.vals, self.fixes, self.links)

        def set_pars(self, pars, force=False):
            for i, p in enumerate(pars):
                self.set_val(i, p[0], force=force)
                self.set_fix(i, p[1], force=force)
                self.set_link(i, p[2], force=force)

        def set_vals(self, vals, force=False):
            for i, val in enumerate(vals):
                self.set_val(i, val, force=force)

        def set_fixes(self, fixes, force=False):
            for i, fix in enumerate(fixes):
                self.set_fix(i, fix, force=force)

        def set_links(self, links, force=False):
            for i, link in enumerate(links):
                self.set_link(i, link, force=force)

        def set_val(self, i, val, force=False):
            if force or isinstance(val, int) or isinstance(val, float):
                if self.model is None:
                    self.vals[i] = val
                else:
                    self.model.set_val(i, val, force=True)  # Set val for all sub-models
                    # to ensure set_val is called for a ListedModel if one is part of the sub-models.
                    # This is only needed for the vals since these may be predicted by the specific model.
                    # Everything else is handled by the top-model.

        def set_fix(self, i, fix, force=False):
            if force:
                self.fixes[i] = fix
                return
            if isinstance(fix, int) or isinstance(fix, float):
                if isinstance(fix, bool):
                    fix = bool(fix)
                elif fix <= 0 or np.isinf(fix):
                    fix = fix <= 0
                else:
                    fix = float(fix)
                expr = 'args[{}]'.format(i)
            elif isinstance(fix, str):
                j, k = fix.find('('), fix.find(')')
                try:
                    if j == -1 or k == -1:
                        raise ValueError('Not a value with uncertainty.')
                    fix = '{}({})'.format(float(fix[:j]), float(fix[j+1:k]))
                    expr = 'args[{}]'.format(i)
                except ValueError:
                    _fix = fix
                    for j, name in enumerate(self.names):
                        _fix = _fix.replace(name, 'eval(self.expressions[{}])'.format(j))
                    expr = _fix
                    try:
                        self._eval_zero_division(self.vals, expr)
                    except (ValueError, TypeError, SyntaxError, NameError) as e:
                        print('Invalid expression for parameter \'{}\': {}. Got a {}.'.format(self.names[i], fix, repr(e)))
                        return
            elif isinstance(fix, list):
                if len(fix) == 0:
                    fix = [0, 1]
                elif len(fix) == 1:
                    fix = [0, fix[0]]
                else:
                    fix = fix[:2]
                expr = 'args[{}]'.format(i)
            else:
                return
            temp_expr = self.expressions[i]
            temp_fix = self.fixes[i]
            self.expressions[i] = compile(expr, '<string>', 'eval', optimize=2)  # Compile beforehand to save time.
            self.fixes[i] = fix
            try:
                self.update_args(self.vals)
            except RecursionError as e:
                print('Expression for {} with fix {} form a loop. Got a {}.'.format(self.names[i], fix, repr(e)))
                self.expressions[i] = temp_expr
                self.fixes[i] = temp_fix

        def set_link(self, i, link, force=False):
            if force:
                self.links[i] = link
                return
            if isinstance(link, int) or isinstance(link, float):
                self.links[i] = bool(link)

        def _eval_zero_division(self, args, expr):
            try:
                return eval(expr, {}, {'self': self, 'args': args})
            except ZeroDivisionError:
                return 0.

        def update_args(self, args):
            return tuple(self._eval_zero_division(args, expr) for expr in self.expressions)

        def update(self):
            self.set_vals(self.update_args(self.vals), force=False)

        def min(self):
            return -1. if self.model is None else self.model.min()

        def max(self):
            return 1. if self.model is None else self.model.max()

        def intervals(self):
            return [[self.min(), self.max()]] if self.model is None else self.model.intervals()

        def x(self):
            return np.concatenate([np.arange(i[0], i[1], self.dx, dtype=float) for i in self.intervals()], axis=0)

        def fit_prepare(self):
            bounds = (-np.inf, np.inf)
            fixed = [fix for fix in self.fixes]
            b_lower = []
            b_upper = []
            _bounds = False
            for i, fix in enumerate(self.fixes):
                if isinstance(fix, bool):
                    b_lower.append(-np.inf)
                    b_upper.append(np.inf)
                elif _is_unc(fix):
                    b_lower.append(-np.inf)
                    b_upper.append(np.inf)
                    fixed[i] = False
                elif isinstance(fix, list):
                    _bounds = True
                    b_lower.append(fix[0])
                    b_upper.append(fix[1])
                    fixed[i] = False
                elif isinstance(fix, str):
                    b_lower.append(-np.inf)
                    b_upper.append(np.inf)
                    fixed[i] = True
                    _fix = fix
                else:
                    raise TypeError('The type {} of the element {} in self.fixes is not supported.'
                                    .format(type(fix), fix))
            if _bounds:
                bounds = (b_lower, b_upper)
            return fixed, bounds

    class Empty(Model):
        """
        An empty model, returning zeros with the same shape as x.
        """
        def __init__(self):
            super().__init__(model=None)
            self.type = 'Empty'

        def evaluate(self, x, *args, **kwargs):
            return np.zeros_like(x)

    class NPeak(Model):
        """
        Evaluates the given 'model' at the positions x\ :sub:`i` with scalings p\ :sub:`i` where i < 'n_peaks'.
        """
        def __init__(self, model, n_peaks=1):
            super().__init__(model=model)
            self.type = 'NPeak'
            self.n_peaks = int(n_peaks)
            for n in range(self.n_peaks):
                self._add_arg('x{}'.format(n), 0., n == 0, False)
                self._add_arg('p{}'.format(n), 1., n == 0, False)

        def evaluate(self, x, *args, **kwargs):
            return np.sum([args[self.model.size + 2 * n + 1]
                           * self.model.evaluate(x - args[self.model.size + 2 * n], *args[:self.model.size])
                           for n in range(self.n_peaks)], axis=0)

        def min(self):
            min_center = min(self.vals[self.p['x{}'.format(n)]] for n in range(self.n_peaks))
            return min_center + self.model.min()

        def max(self):
            max_center = max(self.vals[self.p['x{}'.format(n)]] for n in range(self.n_peaks))
            return max_center + self.model.max()

        def intervals(self):
            return merge_intervals([[i[0] + self.vals[self.model.size + 2 * n],
                                     i[1] + self.vals[self.model.size + 2 * n]]
                                    for i in self.model.intervals() for n in range(self.n_peaks)])

    class Offset(Model):
        """
        Cuts the x-axis and adds y-axis offsets to every segment.
        """
        def __init__(self, model=None, x_cuts=None, offsets=None):
            """
            :param model: The model the offset will be added to. If None, the offset will be added to zero.
            :param x_cuts: x values where to cut the x-axis.
            :param offsets: A list of maximally considered polynomial orders for each slice.
             The list must have length len(x_cuts) + 1.
            """
            super().__init__(model=model)
            self.type = 'Offset'
            if x_cuts is None:
                x_cuts = []
            self.x_cuts = sorted(list(x_cuts))
            self.offsets = offsets
            if self.offsets is None:
                self.offsets = [0]
            if len(self.offsets) != len(self.x_cuts) + 1:
                raise ValueError('The parameter offset must be a list of size \'len(x_cuts) + 1\''
                                 ' and contain the maximally considered polynomial order for each slice.')

            self.offset_map = []
            self.offset_masks = []
            self.update_on_call = True

            self.gen_offset_map()

        def evaluate(self, x, *args, **kwargs):
            if self._model is None:
                return self._offset(x, *args)
            return self.model.evaluate(x, *args[:self.model.size]) + self._offset(x, *args)

        def set_x_cuts(self, x_cuts):
            x_cuts = list(x_cuts)
            if len(x_cuts) != len(self.x_cuts):
                raise ValueError('\'x_cuts\' must not change its size.')
            self.x_cuts = sorted(list(x_cuts))

        def _offset(self, x, *args):
            if self.update_on_call:
                self.gen_offset_masks(x)
            ret = np.zeros_like(x)
            for i, mask in enumerate(self.offset_masks):
                ret[mask] = _poly(x[mask], *_args_ordered(args, self.offset_map[i]))
            return ret

        def gen_offset_map(self):
            self.offset_map = []
            for i, n in enumerate(self.offsets):
                self.offset_map.append([])
                for k in range(n + 1):
                    self.offset_map[-1].append(self._index)
                    self._add_arg('off{}e{}'.format(i, k), 0., False, False)

        """ Preprocessing """

        def gen_offset_masks(self, x):
            self.offset_masks = []
            for x0, x1 in zip([np.min(x) - 1., ] + self.x_cuts, self.x_cuts + [np.max(x) + 1., ]):
                x_mean = 0.5 * (x0 + x1)
                self.offset_masks.append(np.abs(x - x_mean) < x1 - x_mean)

        def guess_offset(self, x, y):
            for i, mask in enumerate(self.offset_masks):
                self.vals[self.p['off{}e0'.format(i)]] = 0.5 * (y[mask][0] + y[mask][-1])
                try:
                    self.vals[self.p['off{}e1'.format(i)]] = (y[mask][-1] - y[mask][0]) / (x[mask][-1] - x[mask][0])
                except KeyError:
                    return

    class Amplifier(Model):
        """
        A polynomial of order 'order'.
        """
        def __init__(self, order=None):
            super().__init__(model=None)
            self.type = 'Amplifier'
            if order is None:
                order = 1
            self.order = order
            for n in range(order + 1):
                self._add_arg('a{}'.format(n), 1. if n == 1 else 0., False, False)
            self._min = -10
            self._max = 10

        def evaluate(self, x, *args, **kwargs):
            self._min = np.min(x)
            self._max = np.max(x)
            return _poly(x, *args)

        @property
        def dx(self):
            return 1e-2

        def min(self):
            return self._min

        def max(self):
            return self._max

    class Custom(Model):
        def __init__(self, model=None, parameters=None):
            super().__init__(model=model)
            self.type = 'Custom'
            if parameters is None:
                parameters = []
            self.parameters = parameters

            for p in self.parameters:
                self._add_arg(p, 0., False, False)

        def evaluate(self, x, *args, **kwargs):
            if self.model is None:
                return np.array(args, dtype=float)
            return self.model.evaluate(x, *args[:self.model.size], **kwargs)

    class YPars(Model):
        def __init__(self, model):
            super().__init__(model=model)
            self.type = 'YPars'

            self.p_y = [i for i, fix in enumerate(self.model.fixes) if _is_unc(fix)]

        def evaluate(self, x, *args, **kwargs):
            return np.concatenate([self.model.evaluate(x, *args, **kwargs),
                                   np.array([args[p_y] for p_y in self.p_y], dtype=float)], axis=0)

    class Listed(Model):
        def __init__(self, models, labels=None):
            super().__init__(model=None)
            self.type = 'Listed'

            self.models = models  # models is just a reference to a list defined somewhere else.
            self.labels = labels
            if self.labels is None:
                if len(self.models) == 1:
                    self.labels = ['', ]
                else:
                    self.labels = ['__{}'.format(i) for i in range(len(self.models))]

            self.slices = []
            self.model_map = []
            self.index_map = []
            for i, (model, label) in enumerate(zip(self.models, self.labels)):
                self.slices.append(slice(self._index, self._index + model.size, 1))
                for j, (name, val, fix, link) in enumerate(model.get_pars()):
                    self.model_map.append(i)
                    self.index_map.append(j)
                    if isinstance(fix, str) and not _is_unc(fix):
                        for _name in model.names:
                            fix = fix.replace(_name, '{}{}'.format(_name, label))
                    self._add_arg('{}{}'.format(name, label), val, fix, link)
            self.set_fixes(list(self.fixes))

        def set_val(self, i, val, force=False):
            super().set_val(i, val, force=force)
            if i < len(self.model_map):
                self.models[self.model_map[i]].set_val(self.index_map[i], self.vals[i], force=True)

        def set_fix(self, i, fix, force=False):
            super().set_fix(i, fix, force=force)
            if i < len(self.model_map):
                self.models[self.model_map[i]].set_fix(self.index_map[i], self.fixes[i], force=True)

        def set_link(self, i, link, force=False):
            super().set_link(i, link, force=force)
            if i < len(self.model_map):
                self.models[self.model_map[i]].set_link(self.index_map[i], self.links[i], force=True)

        def inherit_vals(self, force=False):
            self.set_vals([val for model in self.models for val in model.vals], force=force)

        def inherit_fixes(self, force=False):
            self.set_fixes([fix for model in self.models for fix in model.fixes], force=force)

        def inherit_links(self, force=False):
            self.set_links([link for model in self.models for link in model.links], force=force)

    class Summed(Listed):
        def __init__(self, models, labels=None):
            super().__init__(models, labels=labels)
            self.type = 'Summed'

            self.indices_add = []
            for n, (model, label) in enumerate(zip(self.models, self.labels)):
                self.indices_add.append([self._index, self._index + 1])
                self._add_arg('center{}'.format(label), 0., False, False)
                self._add_arg('int{}'.format(label), 1., False, False)

        def evaluate(self, x, *args, **kwargs):
            return np.sum([args[i[1]] * model.evaluate(x - args[i[0]], *args[_slice], **kwargs)
                           for model, _slice, i in zip(self.models, self.slices, self.indices_add)], axis=0)

        @property
        def dx(self):
            return min(model.dx for model in self.models)

        def min(self):
            return min(model.min() for model in self.models)

        def max(self):
            return max(model.max() for model in self.models)

        def intervals(self):
            return merge_intervals([[i[0] + self.vals[j[0]], i[1] + self.vals[j[0]]]
                                    for model, j in zip(self.models, self.indices_add) for i in model.intervals()])

    class Linked(Listed):
        def __init__(self, models):
            super().__init__(models, labels=None)
            self.type = 'Linked'

            for i, (name, val, fix, link) in enumerate(self.get_pars()):
                if link and (not fix or isinstance(fix, list) or _is_unc(fix)):
                    _name = name[:name.rfind('__')]
                    for j, model in enumerate(self.models):
                        if j < self.model_map[i] and _name in model.names:
                            self.set_fix(i, '{}__{}'.format(_name, j))
                            break

        def evaluate(self, x, *args, **kwargs):
            return np.concatenate(tuple(model.evaluate(_x, *args[_slice], **kwargs)
                                        for model, _slice, _x in zip(self.models, self.slices, x)), axis=0)

        def set_fix(self, i, fix, force=False):
            super(Listed, self).set_fix(i, fix, force=force)

        def set_link(self, i, link, force=False):
            super(Listed, self).set_link(i, link, force=force)