# Licensed under a 3-clause BSD style license - see LICENSE.rst

"""Spline models and fitters."""
# pylint: disable=line-too-long, too-many-lines, too-many-arguments, invalid-name

import warnings

import abc
import functools
import numpy as np

from astropy.utils.exceptions import (AstropyUserWarning,)
from astropy.utils import isiterable
from .core import (FittableModel, ModelDefinitionError)

from .parameters import Parameter


__all__ = ['Spline1D', 'SplineInterpolateFitter', 'SplineSmoothingFitter',
           'SplineExactKnotsFitter', 'SplineSplrepFitter']
__doctest_requires__ = {('Spline1D'): ['scipy']}


class _Spline(FittableModel):
    """Base class for spline models"""
    _knot_names = ()
    _coeff_names = ()

    optional_inputs = {}

    def __init__(self, knots=None, coeffs=None, degree=None, bounds=None,
                 n_models=None, model_set_axis=None, name=None, meta=None):

        super().__init__(
            n_models=n_models, model_set_axis=model_set_axis, name=name,
            meta=meta)

        self._user_knots = False
        self._init_tck(degree)

        # Hack to allow an optional model argument
        self._create_optional_inputs()

        if knots is not None:
            self._init_spline(knots, coeffs, bounds)
        elif coeffs is not None:
            raise ValueError("If one passes a coeffs vector one needs to also pass knots!")

    @property
    def param_names(self):
        """
        Coefficient names generated based on the spline's degree and
        number of knots.
        """

        return tuple(list(self._knot_names) + list(self._coeff_names))

    @staticmethod
    def _optional_arg(arg):
        return f'_{arg}'

    def _create_optional_inputs(self):
        for arg in self.optional_inputs:
            attribute = self._optional_arg(arg)
            if hasattr(self, attribute):
                raise ValueError(f'Optional argument {arg} already exists in this class!')
            else:
                setattr(self, attribute, None)

    def _intercept_optional_inputs(self, **kwargs):
        new_kwargs = kwargs
        for arg in self.optional_inputs:
            if (arg in kwargs):
                attribute = self._optional_arg(arg)
                if getattr(self, attribute) is None:
                    setattr(self, attribute, kwargs[arg])
                    del new_kwargs[arg]
                else:
                    raise RuntimeError(f'{arg} has already been set, something has gone wrong!')

        return new_kwargs

    def evaluate(self, *args, **kwargs):
        """ Extract the optional kwargs passed to call """

        optional_inputs = kwargs
        for arg in self.optional_inputs:
            attribute = self._optional_arg(arg)

            if arg in kwargs:
                # Options passed in
                optional_inputs[arg] = kwargs[arg]
            elif getattr(self, attribute) is not None:
                # No options passed in and Options set
                optional_inputs[arg] = getattr(self, attribute)
                setattr(self, attribute, None)
            else:
                # No options passed in and No options set
                optional_inputs[arg] = self.optional_inputs[arg]

        return optional_inputs

    def __call__(self, *args, **kwargs):
        """
        Make model callable to model evaluation
        """

        # Hack to allow an optional model argument
        kwargs = self._intercept_optional_inputs(**kwargs)

        return super().__call__(*args, **kwargs)

    def _create_parameter(self, name: str, index: int, attr: str, fixed=False):
        """
        Create a spline parameter linked to an attribute array.

        Parameters
        ----------
        name : str
            Name for the parameter
        index : int
            The index of the parameter in the array
        attr : str
            The name for the attribute array
        fixed : optional, bool
            If the parameter should be fixed or not
        """

        # Hack to allow parameters and attribute array to freely exchange values
        #   _getter forces reading value from attribute array
        #   _setter forces setting value to attribute array

        def _getter(value, model: "_Spline", index: int, attr: str):
            return getattr(model, attr)[index]

        def _setter(value, model: "_Spline", index: int, attr: str):
            getattr(model, attr)[index] = value
            return value
        getter = functools.partial(_getter, index=index, attr=attr)
        setter = functools.partial(_setter, index=index, attr=attr)

        default = getattr(self, attr)
        param = Parameter(name=name, default=default[index], fixed=fixed,
                          getter=getter, setter=setter)
        # setter/getter wrapper for parameters in this case require the
        # parameter to have a reference back to its parent model
        param.model = self
        param.value = default[index]

        # Add parameter to model
        self.__dict__[name] = param

    def _create_parameters(self, base_name: str, attr: str, fixed=False):
        """
        Create a spline parameters linked to an attribute array for all
        elements in that array

        Parameters
        ----------
        base_name : str
            Base name for the parameters
        attr : str
            The name for the attribute array
        fixed : optional, bool
            If the parameters should be fixed or not
        """
        names = []
        for index in range(len(getattr(self, attr))):
            name = f"{base_name}{index}"
            names.append(name)

            self._create_parameter(name, index, attr, fixed)

        return tuple(names)

    @abc.abstractmethod
    def _init_parameters(self):
        raise NotImplementedError("This needs to be implemented")

    @abc.abstractmethod
    def _init_data(self, knots, coeffs, bounds=None):
        raise NotImplementedError("This needs to be implemented")

    def _init_spline(self, knots, coeffs, bounds=None):
        self._init_data(knots, coeffs, bounds)
        self._init_parameters()

    def _init_tck(self, degree):
        self._c = None
        self._t = None
        self._degree = degree


class Spline1D(_Spline):
    """
    One dimensional Spline Model

    Parameters
    ----------
    knots :  optional
        Define the knots for the spline. Can be 1) the number of interior
        knots for the spline, 2) the array of all knots for the spline, or
        3) If both bounds are defined, the interior knots for the spline
    coeffs : optional
        The array of knot coefficients for the spline
    degree : optional
        The degree of the spline. It must be 1 <= degree <= 5, default is 3.
    bounds : optional
        The upper and lower bounds of the spline.

    Notes
    -----
    Much of the functionality of this model is provided by
    `scipy.interpolate.BSpline` which can be directly accessed via the
    bspline property.

    Fitting for this model is provided by wrappers for:
    `scipy.interpolate.UnivariateSpline`,
    `scipy.interpolate.InterpolatedUnivariateSpline`,
    and `scipy.interpolate.LSQUnivariateSpline`.

    If one fails to define any knots/coefficients, no parameters will
    be added to this model until a fitter is called. This is because
    some of the fitters for splines vary the number of parameters and so
    we cannot define the parameter set until after fitting in these cases.

    Since parameters are not necessarily known at model initialization,
    setting model parameters directly via the model interface has been
    disabled.

    Direct constructors are provided for this model which incorporate the
    fitting to data directly into model construction.

    Knot parameters are declared as "fixed" parameters by default to
    enable the use of other `astropy.modeling` fitters to be used to