# lsqfitgp/_linalg/_decomp.py
#
# Copyright (c) 2023, Giacomo Petrillo
#
# This file is part of lsqfitgp.
#
# lsqfitgp is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# lsqfitgp is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with lsqfitgp.  If not, see <http://www.gnu.org/licenses/>.

"""

Copy-pasted from the notes:

2023-02-14
==========

My current decomposition system is a mess. I can't take reverse gradients.
I can't straightforwardly implement optimized algorithms that compute together
likelihood, gradient, and fisher. Jax patterns break down unpredictably. I
have to redesign it from scratch.

Guidelines and requirements:

  - Sufficient modularity to implement composite decompositions (Woodbury,
    Block)

  - Does not mess up jax in any way

  - Caches decompositions

  - Favors optimizing together the likelihood and its derivatives

Operations (in the following I indicate with lowercase inputs which are
typically vectors or tall matrices, and uppercase inputs which are typically
large matrices, since optimization requires taking it into account):

    pinv_bilinear(A, r) => A'K⁺r (for the posterior mean)
    pinv_bilinear_robj(A, r) same but r can be gvars
    ginv_quad(A) => A'K⁻A (for the posterior covariance)
        I want the pseudoinverse for the mean because the data may not be
        in the span and I want to project it orthogonally, while for the
        covariance I expect A and K to come from a pos def matrix so they are
        coherent
    ginv_diagquad(A) => diag(A'K⁻A) (for the posterior variance)
    minus_log_normal_density(
        r: 1d array,      # the residuals (data - prior mean)
        dr_vjp: callable, # x -> x_i ∂r_i/∂p_j,   gradrev and fishvec
        dK_vjp: callable, # x -> x_ij ∂K_ij/∂p_k, gradrev and fishvec
        dr_jvp: callable, # x -> ∂r_i/∂r_j x_j,  fishvec
        dK_jvp: callable, # x -> ∂K_ij/∂p_k x_k, fishvec
        dr: 2d array,     # ∂r_i/∂p_j,  gradfwd and fisher
        dK: 3d array,     # ∂K_ij/∂p_k, gradfwd and fisher
        vec: 1d array,    # input vector of fishvec, same size as params
        value: bool,
        gradrev: bool,
        gradfwd: bool,
        fisher: bool,
        fishvec: bool,
    )
        This computes on request
            value: 1/2 tr(KK⁺) log 2π
                 + 1/2 tr(I-KK⁺) log 2π
                 + 1/2 log pdet K
                 + 1/2 tr(I-KK⁺) log ε
                 + 1/2 r'(K⁺+(I-KK⁺)/ε)r
            gradrev,
            gradfwd: 1/2 tr(K⁺dK)
                    + r'(K⁺+(I-KK⁺)/ε) dr
                    - 1/2 r'(K⁺+2(I-KK⁺)/ε)dKK⁺r
            fisher: 1/2 tr(K⁺dK(K⁺+2(I-KK⁺)/ε)d'K)
                  - 2 tr(K⁺dK(I-KK⁺)d'KK⁺)
                  + dr'(K⁺+(I-KK⁺)/ε)d'r
            fishvec: fisher matrix times vec
        There should be options for omitting the pieces with ε. I also need a
        way to make densities with different values of ε comparable with each
        other (may not be possible, if it is, it probably requires a history of
        ranks and ε). gradfwd/rev form K⁺ explicitly to compute tr(K⁺dK) for
        efficiency.
    correlate(x)
        Zx where K = ZZ'.
    back_correlate(X):
        Z'X, this is used by Sandwich and Woodbury.

Since I also want to compute the Student density, I could split
minus_log_normal_density's return value into logdet and quad. The gradient
splits nicely between the two terms, but I have to redo the calculation of the
Fisher matrix for the Student distribution. Alternatively, I could use the the
Normal Fisher. => See Lange et al. (1989, app. B). => I think I can split the
gradient and Fisher matrix too.

2023-03-07
==========

To compute a Fisher-vector product when there are many parameters, do

    tr(K+ dK K+ dK) v =
    = K_vjp(K+ K_jvp(v) K+)

"""

# TODO to automatize this further, I could take in a function that generates K
# (or its pieces) and the arguments to the function. But how would this play
# together with passing decomposition objects as pieces?

# TODO split this file by class

# TODO Consider using lineax for implementing non-materialized decomps

import abc
import functools

import numpy
import jax
from jax import numpy as jnp
from jax.scipy import linalg as jlinalg
from jax import lax

from .. import _jaxext
from . import _pytree

class Decomposition(_pytree.AutoPyTree, abc.ABC):
    """
    Abstract base class for decompositions of positive semidefinite matrices.
    """

    @abc.abstractmethod
    def __init__(self, *args, **kw):
        """ Decompose the input matrix """
        pass

    @abc.abstractmethod
    def matrix(self):
        """ The input matrix """
        pass

    @abc.abstractmethod
    def ginv_linear(self, X):
        """ Compute K⁻X """
        pass

    @abc.abstractmethod
    def pinv_bilinear(self, A, r):
        """Compute A'K⁺r."""
        pass

    @abc.abstractmethod
    def pinv_bilinear_robj(self, A, r):
        """Compute A'K⁺r, where r can be an array of objects."""
        pass

    @abc.abstractmethod
    def ginv_quad(self, A):
        """Compute A'K⁻A."""
        pass

    @abc.abstractmethod
    def ginv_diagquad(self, A):
        """Compute diag(A'K⁻A)."""
        pass

    @abc.abstractmethod
    def correlate(self, x):
        """ Compute Zx where K = ZZ' """
        pass

    @abc.abstractmethod
    def back_correlate(self, X):
        """ Compute Z'X """
        pass

    @abc.abstractmethod
    def pinv_correlate(self, x):
        """ Compute Z⁺x """
        pass

    @abc.abstractmethod
    def minus_log_normal_density(self,
        r,
        *,
        dr_vjp=None,
        dK_vjp=None,
        dr_jvp=None,
        dK_jvp=None,
        dr=None,
        dK=None,
        value=False,
        gradrev=False,
        gradfwd=False,
        fisher=False,
        fishvec=False,
        ):
        """
        Compute minus log a Normal density and its derivatives, with covariance
        matrix K.

        If an input derivative is not specified, it is assumed to be zero.

        Parameters
        ----------
        r: 1d array
            The residuals (value - mean)
        dr_vjp: callable
            x -> x_i ∂r_i/∂p_j, for gradrev and fishvec
        dK_vjp: callable
            x -> x_ij ∂K_ij/∂p_k, for gradrev and fishvec
        dr_jvp_vec: 1d array
            ∂r_i/∂r_j vec_j, for fishvec
        dK_jvp_vec: 2d array
            ∂K_ij/∂p_k vec_k, for fishvec
        dr: 2d array
            ∂r_i/∂p_j for gradfwd and fisher
        dK: 3d array
            ∂K_ij/∂p_k, for gradfwd and fisher
        value: bool
        gradrev: bool
        gradfwd: bool
        fisher: bool
        fishvec: bool
            These parameters indicate which of the return values to compute.
            Default all False.

        Returns
        -------
        value: 1/2 tr(KK⁺) log 2π
             + 1/2 tr(I-KK⁺) log 2π
             + 1/2 log pdet K
             + 1/2 tr(I-KK⁺) log ε
             + 1/2 r'(K⁺+(I-KK⁺)/ε)r
        gradrev,
        gradfwd: 1/2 tr(K⁺dK)
                + r'(K⁺+(I-KK⁺)/ε) dr
                - 1/2 r'(K⁺+2(I-KK⁺)/ε)dKK⁺r
        fisher: 1/2 tr(K⁺dK(K⁺+2(I-KK⁺)/ε)d'K)
              - 2 tr(K⁺dK(I-KK⁺)d'KK⁺)
              + dr'(K⁺+(I-KK⁺)/ε)d'r
        fishvec: fisher matrix @ vec
        """
        pass

    def _parseeps(self, K, epsrel, epsabs, maxeigv=None):
        """ Determine eps from input arguments """
        machine_eps = jnp.finfo(_jaxext.float_type(K)).eps
        if epsrel == 'auto':
            epsrel = len(K) * machine_eps
        if epsabs == 'auto':
            epsabs = machine_eps
        if maxeigv is None:
            maxeigv = eigval_bound(K)
        self._eps = epsrel * maxeigv + epsabs
        return self._eps

    @property
    def eps(self):
        """
        The threshold below which eigenvalues are too small to be determined.
        """
        return self._eps

    @property
    @abc.abstractmethod
    def n(self):
        """ Number of rows/columns of the matrix """
        pass

    @property
    @abc.abstractmethod
    def m(self):
        """ Number of columns of Z """
        pass

    def ginv(self):
        """ Compute K⁻ """
        return self.ginv_quad(jnp.eye(self.n))

def solve_triangular_python(a, b, *, lower=False):
    """
    Pure python implementation of scipy.linalg.solve_triangular for when
    a or b are object arrays.
    """
    # TODO maybe commit this to gvar.linalg
    a = numpy.asarray(a)
    x = numpy.copy(b)

    vec = x.ndim < 2
    if vec:
        x = x[:, None]

    n = a.shape[-1]
    assert x.shape[-2] == n

    if not lower:
        a = a[..., ::-1, ::-1]
        x = x[..., ::-1, :]

    x[..., 0, :] /= a[..., 0, 0, None]
    for i in range(1, n):
        x[..., i:, :] -= x[..., None, i - 1, :] * a[..., i:, i - 1, None]
        x[..., i, :] /= a[..., i, i, None]

    if not lower:
        x = x[..., ::-1, :]

    if vec:
        x = numpy.squeeze(x, -1)
    return x

def solve_triangular_batched(a, b, *, lower=False):
    """ Version of jax.scipy.linalg.solve_triangular that batches matmul-like """
    a = jnp.asarray(a)
    b = jnp.asarray(b)
    vec = b.ndim < 2
    if vec:
        b = b[:, None]

    batch_shape = jnp.broadcast_shapes(a.shape[:-2], b.shape[:-2])
    a_shape = batch_shape + a.shape[-2:]
    b_shape = batch_shape + b.shape[-2:]
    result = lax.linalg.triangular_solve(
        jnp.broadcast_to(a, a_shape), jnp.broadcast_to(b, b_shape),
        left_side=True, lower=lower,
    )
    assert result.shape == b_shape

    if vec:
        result = result.squeeze(-1)
    return result

def solve_batched(a, b, **kw):
    """ Version of jax.scipy.linalg.solve that batches matmul-like """
    a = jnp.asarray(a)
    b = jnp.asarray(b)
    vec = b.ndim < 2
    if vec:
        b = b[:, None]

    @functools.partial(jnp.vectorize, signature='(i,j),(j,k)->(i,k)')
    def solve_batched(a, b):
        return jlinalg.solve(a, b, **kw)
    result = solve_batched(a, b)

    if vec:
        result = result.squeeze(-1)
    return result

def eigval_bound(K):
    """
    Upper bound on the largest magnitude eigenvalue of the matrix, from
    Gershgorin's theorem.
    """
    return jnp.max(jnp.sum(jnp.abs(K), axis=1))

def diag_scale_pow2(K):
    """
    Compute a vector s of powers of 2 such that diag(K / outer(s, s)) ~ 1.
    """
    d = jnp.diag(K)
    return jnp.where(d, jnp.exp2(jnp.rint(0.5 * jnp.log2(d))), 1)

def transpose(x):
    """ swap the last two axes of array x, corresponds to matrix tranposition
    with the broadcasting convention of matmul """
    if x.ndim < 2:
        return x
    elif isinstance(x, jnp.ndarray):
        return jnp.swapaxes(x, -2, -1)
    else:
        # need to support numpy because this function is used with gvars
        return numpy.swapaxes(x, -2, -1)

class Chol(Decomposition):
    """Cholesky decomposition. The matrix is regularized adding a small multiple
    of the identity."""

    def __init__(self, K, *, epsrel='auto', epsabs=0):
        # K <- K + Iε
        # K = LL'
        self._K = K
        s = diag_scale_pow2(K)
        K = K / s / s[:, None]
        eps = self._parseeps(K, epsrel, epsabs)
        K = K.at[jnp.diag_indices_from(K)].add(eps)
        L = jlinalg.cholesky(K, lower=True)
        with _jaxext.skipifabstract():
            if not jnp.all(jnp.isfinite(L)):
                # TODO check that jax fills with nan after failed row, detect
                # and report minor index like scipy
                raise numpy.linalg.LinAlgError('cholesky decomposition not finite, probably matrix not pos def numerically')
        self._L = L * s[:, None]
        self._eps = eps * jnp.min(s * s)

    def matrix(self):
        return self._K

    def ginv_linear(self, X):
        # = K⁻¹X
        # K⁻¹ = L'⁻¹L⁻¹
        # K⁻¹X = L'⁻¹(L⁻¹X)
        invLX = jlinalg.solve_triangular(self._L, X, lower=True)
        return jlinalg.solve_triangular(self._L.T, invLX, lower=False)

    def pinv_bilinear(self, A, r):
        # = A'K⁻¹r = A'L'⁻¹L⁻¹r = (L⁻¹A)'(L⁻¹r)
        invLr = jlinalg.solve_triangular(self._L, r, lower=True)
        invLA = jlinalg.solve_triangular(self._L, A, lower=True)
        return invLA.T @ invLr

    def pinv_bilinear_robj(self, A, r):
        # = A'K⁻¹r
        invLr = solve_triangular_python(self._L, r, lower=True)
        invLA = jlinalg.solve_triangular(self._L, A, lower=True)
        return numpy.asarray(invLA).T @ invLr

    def ginv_quad(self, A):
        # = A'K⁻¹A = A'K⁻¹A = A'L'⁻¹L⁻¹A = (L⁻¹A)'(L⁻¹A)
        invLA = jlinalg.solve_triangular(self._L, A, lower=True)
        return invLA.T @ invLA

    def ginv_diagquad(self, A):
        # = diag(A'K⁻¹A)
        # X = L⁻¹A
        # diag(A'K⁻¹A)_i = diag(X'X)_i = ∑_j X'_ij X_ji = ∑_j X_ji X_ji
        invLA = jlinalg.solve_triangular(self._L, A, lower=True)
        return jnp.einsum('ji,ji->i', invLA, invLA)

    def correlate(self, x):
        # = Lx
        return self._L @ x

    def back_correlate(self, X):
        # = L'X
        return self._L.T @ X

    def pinv_correlate(self, x):
        # = L⁻¹x
        return jlinalg.solve_triangular(self._L, x, lower=True)

    def minus_log_normal_density(self,
        r,               # 1d array, the residuals (data - prior mean)
        *,
        dr_vjp=None,     # callable, x -> x_i ∂r_i/∂p_j,   gradrev and fishvec
        dK_vjp=None,     # callable, x -> x_ij ∂K_ij/∂p_k, gradrev and fishvec
        dr_jvp_vec=None, # 1d array, ∂r_i/∂r_j v_j,  fishvec
        dK_jvp_vec=None, # 2d array, ∂K_ij/∂p_k v_k, fishvec
        dr=None,         # 2d array, ∂r_i/∂p_j,  gradfwd and fisher
        dK=None,         # 3d array, ∂K_ij/∂p_k, gradfwd and fisher
        value=False,
        gradrev=False,
        gradfwd=False,
        fisher=False,
        fishvec=False,
    ):

        L = self._L

        out = {}

        # compute shared factors
        grad = (
            (gradrev and (dK_vjp is not None or dr_vjp is not None))
            or (gradfwd and (dK is not None or dr is not None))
        )
        if value or grad:
            invLr = jlinalg.solve_triangular(L, r, lower=True)
        if grad:
            invKr = jlinalg.solve_triangular(L.T, invLr, lower=False)
        if (gradrev and dK_vjp is not None) or (gradfwd and dK is not None):
            invL = jlinalg.solve_triangular(L, jnp.eye(len(L)), lower=True)
            invK = invL.T @ invL

        if value:
            # = 1/2 n log 2π
            #   + 1/2 log det K
            #   + 1/2 r'K⁻¹r
            # K = LL'
            # K⁻¹ = L'⁻¹L⁻¹
            # det K = (det L)² =
            #       = (∏_i L_ii)²
            # r'K⁻¹r = r'L'⁻¹L⁻¹r =
            #        = (L⁻¹r)'(L⁻¹r)
            out['value'] = 1/2 * (
                len(L) * jnp.log(2 * jnp.pi) +
                2 * jnp.sum(jnp.log(jnp.diag(L))) +
                invLr @ invLr
            )
        else:
            out['value'] = None

        if gradrev:
            # = 1/2 tr(K⁻¹dK)
            #   + r'K⁻¹dr
            #   - 1/2 r'K⁻¹dKK⁻¹r
            # tr(K⁻¹dK) = K⁻¹_ij dK_ji =
            #           = K⁻¹_ij dK_ij =
            #           = dK_vjp(K⁻¹)
            # r'K⁻¹dr = r_i K⁻¹_ij dr_j =
            #         = (K⁻¹r)_j dr_j =
            #         = dr_vjp(K⁻¹r)
            # r'K⁻¹dKK⁻¹r = r_i K⁻¹_ij dK_jl K⁻¹_lm r_m =
            #             = (K⁻¹r)_j dK_jl (K⁻¹r)_l =
            #             = dK_vjp((K⁻¹r) ⊗ (K⁻¹r))
            out['gradrev'] = 0
            if dK_vjp is not None:
                tr_invK_dK = dK_vjp(invK)
                r_invK_dK_invK_r = dK_vjp(jnp.outer(invKr, invKr))
                out['gradrev'] += 1/2 * (tr_invK_dK - r_invK_dK_invK_r)
            if dr_vjp is not None:
                r_invK_dr = dr_vjp(invKr)
                out['gradrev'] += r_invK_dr
        else:
            out['gradrev'] = None

        if gradfwd:
            # = 1/2 tr(K⁻¹dK)
            #   + r'K⁻¹dr
            #   - 1/2 r'K⁻¹dKK⁻¹r
            # tr(K⁻¹dK)_k = K⁻¹_ij dK_ijk
            # r'K⁻¹dr = (K⁻¹r)'dr
            # (r'K⁻¹dKK⁻¹r)_k = r_i K⁻¹_ij dK_jlk K⁻¹_lm r_m =
            #                 = (K⁻¹r)_j dK_jlk (K⁻¹r)_l
            out['gradfwd'] = 0
            if dK is not None:
                tr_invK_dK = jnp.einsum('ij,ijk->k', invK, dK)
                r_invK_dK_invK_r = jnp.einsum('i,ijk,j->k', invKr, dK, invKr)
                out['gradfwd'] += 1/2 * (tr_invK_dK - r_invK_dK_invK_r)
            if dr is not None:
                r_invK_dr = invKr @ dr
                out['gradfwd'] += r_invK_dr
        else:
            out['gradfwd'] = None

        if fisher:
            # = 1/2 tr(K⁻¹dKK⁻¹d'K)
            #   + dr'K⁻¹d'r
            # tr(K⁻¹dKK⁻¹d'K)_ij = tr(L'⁻¹L⁻¹dKL'⁻¹L⁻¹d'K)_ij =
            #                    = tr(L⁻¹dKL'⁻¹L⁻¹d'KL'⁻¹)_ij =
            #                    = (L⁻¹dKL'⁻¹)_kli (L⁻¹dKL'⁻¹)_klj
            # (L⁻¹dKL'⁻¹)_ijk = L⁻¹_il dK_lmk L'⁻¹_mj =
            #                 = L⁻¹_il L⁻¹_jm dK_lmk
            # (dr'K⁻¹d'r)_kq = dr'_k L'⁻¹L⁻¹dr_q =
            #                = (L⁻¹dr_k)_i (L⁻¹dr_q)_i
            out['fisher'] = 0
            if dK is not None:
                invL_dK = solve_triangular_batched(L,
                    jnp.moveaxis(dK, 2, 0),
                    lower=True) # kim: L⁻¹_il dK_lmk
                invL_dK_invL = solve_triangular_batched(L,
                    jnp.swapaxes(invL_dK, 1, 2),
                    lower=True) # kji: L⁻¹_jm (L⁻¹_il dK_lmk)
                tr_invK_dK_invK_dK = jnp.einsum('kij,qij->kq', invL_dK_invL, invL_dK_invL)
                out['fisher'] += 1/2 * tr_invK_dK_invK_dK
            if dr is not None:
                invLdr = jlinalg.solve_triangular(L, dr, lower=True)
                dr_invK_dr = invLdr.T @ invLdr
                out['fisher'] += dr_invK_dr
        else:
            out['fisher'] = None

        if fishvec:
            # = 1/2 tr(K⁻¹dKK⁻¹d'K) v
            #   + dr'K⁻¹d'r v
            # tr(K⁻¹dKK⁻¹d'K) v = K_vjp(K⁻¹K_jvp(v)K⁻¹) =
            #                   = K_vjp(L'⁻¹L⁻¹ K_jvp(v) L'⁻¹L⁻¹)
            # dr'K⁻¹d'r v = dr'K⁻¹dr_jvp(v) =
            #             = dr_vjp(K⁻¹dr_jvp(v)) =
            #             = dr_vjp(L'⁻¹L⁻¹ dr_jvp(v))
            out['fishvec'] = 0
            if not (dK_jvp_vec is None and dK_vjp is None):
                invL_dKv = jlinalg.solve_triangular(L, dK_jvp_vec, lower=True)
                invK_dKv = jlinalg.solve_triangular(L.T, invL_dKv, lower=False)
                invL_dKv_invK = jlinalg.solve_triangular(L, invK_dKv.T, lower=True)
                invK_dKv_invK = jlinalg.solve_triangular(L.T, invL_dKv_invK, lower=False)
                tr_invK_dK_invK_dK_v = dK_vjp(invK_dKv_invK)
                out['fishvec'] += 1/2 * tr_invK_dK_invK_dK_v
            if not (dr_jvp_vec is None and dr_vjp is None):
                invL_drv = jlinalg.solve_triangular(L, dr_jvp_vec, lower=True)
                invK_drv = jlinalg.solve_triangular(L.T, invL_drv, lower=False)
                dr_invK_drv_v = dr_vjp(invK_drv)
                out['fishvec'] += dr_invK_drv_v
        else:
            out['fishvec'] = None

        return tuple(out.values())

    @classmethod
    def make_derivs(cls,
        K_fun, r_fun, primal,
        *,
        args=(),
        kw={},
        vec=None,
        value=False,
        gradrev=False,
        gradfwd=False,
        fisher=False,
        fishvec=False,
    ):
        """
        Prepares arguments for `minus_log_normal_density`.

        Parameters
        ----------
        K_fun, r_fun : callable
            Functions with signature ``f(primal, *args, **kw)`` that produce the
            `K` init argument and the `r` `minus_log_normal_density` argument.
        primal : 1d array
            The first argument to `K_fun` and `r_fun`.
        args : tuple
            Additional positional arguments to `K_fun` and `r_fun`.
        kw : dict
            Keyword arguments to `K_fun` and `r_fun`.
        vec : 1d array
            A tangent vector to compute the jacobian-vector products.
        value, gradrev, gradfwd, fisher, fishvec : bool
            Arguments to `minus_log_normal_density`, used to determine which
            derivatives are needed.

        Returns
        -------
        K : 2d array
            Output of `K_fun`.
        r : 1d array
            Output of `r_fun`.
        out : dict
            Dictionary with derivative arguments to `minus_log_normal_density`.
        """

        partial = lambda f: lambda x: f(x, *args, **kw)
        K_fun = partial(K_fun)
        r_fun = partial(r_fun)

        out = {}

        if gradrev or fishvec:
            K, dK_vjp = jax.vjp(K_fun, primal)
            r, dr_vjp = jax.vjp(r_fun, primal)
            out['dK_vjp'] = lambda x: dK_vjp(x)[0]
            out['dr_vjp'] = lambda x: dr_vjp(x)[0]
        else:
            K = K_fun(primal)
            r = r_fun(primal)
        if fishvec:
            _, out['dK_jvp_vec'] = jax.jvp(K_fun, (primal,), (vec,))
            _, out['dr_jvp_vec'] = jax.jvp(r_fun, (primal,), (vec,))
        if gradfwd or fisher:
            out['dK'] = jax.jacfwd(K_fun)(primal)
            out['dr'] = jax.jacfwd(r_fun)(primal)

        return K, r, out

    @property
    def n(self):
        return len(self._L)

    m = n
