"""Coordinate changes in state space models."""

import abc

try:
    # cached_property is only available in Python >=3.8
    from functools import cached_property
except ImportError:
    from cached_property import cached_property

import numpy as np
import scipy.special  # for vectorised factorial

from probnum import config, linops, randvars


def apply_precon(precon, rv):
    # public (because it is needed in some integrator implementations),
    # but not exposed to the 'randprocs' namespace
    # (i.e. not imported in any __init__.py).

    # There is no way of checking whether `rv` has its Cholesky factor computed already or not.
    # Therefore, since we need to update the Cholesky factor for square-root filtering,
    # we also update the Cholesky factor for non-square-root algorithms here,
    # which implies additional cost.
    # See Issues #319 and #329.
    # When they are resolved, this function here will hopefully be superfluous.

    new_mean = precon @ rv.mean
    new_cov_cholesky = precon @ rv.cov_cholesky  # precon is diagonal, so this is valid
    new_cov = new_cov_cholesky @ new_cov_cholesky.T

    return randvars.Normal(new_mean, new_cov, cov_cholesky=new_cov_cholesky)


class Preconditioner(abc.ABC):
    """Coordinate change transformations as preconditioners in state space models.

    For some models, this makes the filtering and smoothing steps more numerically
    stable.
    """

    @abc.abstractmethod
    def __call__(self, step) -> np.ndarray:
        # if more than step is needed, add them into the signature in the future
        raise NotImplementedError

    @cached_property
    def inverse(self) -> "Preconditioner":
        raise NotImplementedError


class NordsieckLikeCoordinates(Preconditioner):
    """Nordsieck-like coordinates.

    Similar to Nordsieck coordinates (which store the Taylor coefficients instead of the
    derivatives), but better for ODE filtering and smoothing. Used in integrator-transitions, e.g. in
    :class:`IntegratedWienerTransition`.
    """

    def __init__(self, powers, scales, dimension):
        # Clean way of assembling these coordinates cheaply,
        # because the powers and scales of the inverse
        # are better read off than inverted
        self.powers = powers
        self.scales = scales
        self.dimension = dimension

    @classmethod
    def from_order(cls, order, dimension):
        # used to conveniently initialise in the beginning
        powers = np.arange(order, -1, -1)
        scales = scipy.special.factorial(powers)
        return cls(
            powers=powers + 0.5,
            scales=scales,
            dimension=dimension,
        )

    def __call__(self, step):
        scaling_vector = np.abs(step) ** self.powers / self.scales
        if config.matrix_free:
            return linops.Kronecker(
                A=linops.Identity(self.dimension),
                B=linops.Scaling(factors=scaling_vector),
            )
        return np.kron(np.eye(self.dimension), np.diag(scaling_vector))

    @cached_property
    def inverse(self) -> "NordsieckLikeCoordinates":
        return NordsieckLikeCoordinates(
            powers=-self.powers,
            scales=1.0 / self.scales,
            dimension=self.dimension,
        )

[((33, 11, 33, 76), 'probnum.randvars.Normal', 'randvars.Normal', (), '', False, 'from probnum import config, linops, randvars\n'), ((72, 17, 72, 41), 'numpy.arange', 'np.arange', ({(72, 27, 72, 32): 'order', (72, 34, 72, 36): '-1', (72, 38, 72, 40): '-1'}, {}), '(order, -1, -1)', True, 'import numpy as np\n'), ((87, 23, 87, 45), 'numpy.eye', 'np.eye', ({(87, 30, 87, 44): 'self.dimension'}, {}), '(self.dimension)', True, 'import numpy as np\n'), ((87, 47, 87, 70), 'numpy.diag', 'np.diag', ({(87, 55, 87, 69): 'scaling_vector'}, {}), '(scaling_vector)', True, 'import numpy as np\n'), ((81, 25, 81, 37), 'numpy.abs', 'np.abs', ({(81, 32, 81, 36): 'step'}, {}), '(step)', True, 'import numpy as np\n'), ((84, 18, 84, 49), 'probnum.linops.Identity', 'linops.Identity', ({(84, 34, 84, 48): 'self.dimension'}, {}), '(self.dimension)', False, 'from probnum import config, linops, randvars\n'), ((85, 18, 85, 56), 'probnum.linops.Scaling', 'linops.Scaling', (), '', False, 'from probnum import config, linops, randvars\n')]