# Copyright 2022-2024 Jean-Baptiste Delisle
# Licensed under the EUPL-1.2 or later

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.linalg import block_diag, cholesky, solve, solve_triangular

__all__ = ['CBV', 'Bspline']


@dataclass
class LinearMarginalizer:
  r"""
  Linear model for which we want to marginalize the likelihood over its parameters.

  Parameters
  ----------
  size : float
    Size (n) of the dataset.
  nparam : int
    Number of linear parameters.
  indices: (c, 2) ndarray
    Start/end indices of the parameters corresponding to each component of the model.
  icovprior : (nparam, nparam) ndarray
    Inverse of the covariance matrix of the prior on linear parameters.
  logsqdetcovprior: float
    log of the square-root of the determinant of the covariance matrix.
  """

  size: int
  nparam: int
  indices: np.ndarray
  At: np.ndarray
  icovprior: np.ndarray
  logsqdetcovprior: float

  def __add__(self, b: LinearMarginalizer) -> LinearMarginalizer:
    assert self.size == b.size
    return LinearMarginalizer(
      self.size,
      self.nparam + b.nparam,
      np.concatenate((self.indices, b.indices + self.nparam)),
      np.concatenate((self.At, b.At)),
      block_diag(self.icovprior, b.icovprior),
      self.logsqdetcovprior + b.logsqdetcovprior,
    )

  def loglike(self, y: np.ndarray, sig: np.ndarray, m: np.ndarray | float = 1) -> float:
    r"""
    Marginalized log-likelihood over all linear parameters.

    Parameters
    ----------
    y : (n,) ndarray
      Timeseries to be detrended.
    sig : (n,) ndarray
      Measurements errors.
    m : (n,) ndarray or float
      Multiplicative model.

    Returns
    -------
    loglike : float
      The marginalized log-likelihood.
    """
    Gt = self.At * m / sig
    u = y / sig
    icov = Gt @ Gt.T + self.icovprior
    x = Gt @ u
    L = cholesky(icov, lower=True)
    z = solve_triangular(L, x, lower=True)
    return (
      -np.sum(np.log(np.diag(L)))
      - self.logsqdetcovprior
      - np.sum(np.log(sig))
      - y.size / 2 * np.log(2 * np.pi)
      + np.sum(z * z) / 2
      - np.sum(u * u) / 2
    )

  def coefficients(
    self,
    y: np.ndarray,
    sig: np.ndarray,
    m: np.ndarray | float = 1,
    split: bool = False,
    cov: bool = False,
  ) -> np.ndarray | tuple:
    r"""
    Conditional expectation of the model coefficients given y and m.

    Parameters
    ----------
    y : (n,) ndarray
      Timeseries to be detrended.
    sig : (n,) ndarray
      Measurements errors.
    m : (n,) ndarray or float
      Multiplicative model.
    split: bool
      Wheter to split the contributions of each component of the model or to merge them.
    cov: bool
      Whether to compute the conditional covariance in addition to the mean.

    Returns
    -------
    mu_eta : (m,) ndarray or tuple of (m,) ndarrays.
      Conditional expectation of the model coefficients,
      in a single array or split by component (if split is True).
    cov_eta : optional (m, m) ndarray or tuple of (m, m) ndarrays.
      Conditional covariance of the model coefficients (only provided if cov is True).
      If both cov and split are True, the output is the tuple of
      couples (mu_eta, cov_eta) for each component.
    """
    Gt = self.At * m / sig
    u = y / sig
    icov = Gt @ Gt.T + self.icovprior
    x = Gt @ u
    if cov:
      cov_eta = np.linalg.inv(icov)
      mu_eta = cov_eta @ x
      if split:
        return tuple((mu_eta[sk:ek], cov_eta[sk:ek, sk:ek]) for sk, ek in self.indices)
      else:
        return (mu_eta, cov_eta)
    else:
      mu_eta = solve(icov, x)
      if split:
        return tuple(mu_eta[sk:ek] for sk, ek in self.indices)
      else:
        return mu_eta

  def mean(
    self,
    y: np.ndarray,
    sig: np.ndarray,
    m: np.ndarray | float = 1,
    split: bool = False,
    cov: bool = False,
  ) -> np.ndarray | tuple:
    r"""
    Conditional expectation of the model given y and m.

    Parameters
    ----------
    y : (n,) ndarray
      Timeseries to be detrended.
    sig : (n,) ndarray
      Measurements errors.
    m : (n,) ndarray or float
      Multiplicative model.
    split: bool
      Wheter to split the contributions of each component of the model or to merge them.
    cov: bool
      Whether to compute the conditional covariance in addition to the mean.

    Returns
    -------
    mu : (n,) ndarray or tuple of (n,) ndarrays
      Conditional expectation of the model, or of each component (if split is True).
    cov : optional (n, n) ndarray or tuple of (n, n) ndarrays
      Conditional covariance of the model (only provided if cov is True).
      If both cov and split are True, the output is the tuple of
      couples (mu, cov) for each component.
    """
    mu_eta = self.coefficients(y, sig, m, False, cov)
    if cov:
      mu_eta, cov_eta = mu_eta
      if split:
        return tuple(
          (
            mu_eta[sk:ek] @ self.At[sk:ek],
            self.At[sk:ek].T @ cov_eta[sk:ek, sk:ek] @ self.At[sk:ek],
          )
          for sk, ek in self.indices
        )
      else:
        return (mu_eta @ self.At, self.At.T @ cov_eta @ self.At)
    else:
      if split:
        return tuple(mu_eta[sk:ek] @ self.At[sk:ek] for sk, ek in self.indices)
      else:
        return mu_eta @ self.At


class CBV(LinearMarginalizer):
  r"""
  Cotrending basis vectors.

  Parameters
  ----------
  cbv: (m, n) ndarray
    Matrix of m cotrending basis vectors for a time series of size n.
  prior : (m, m) ndarray or (m,) ndarray or float or None
    Scale of the Gaussian prior on the CBV parameters.
    If a matrix is provided, this should be the covariance matrix.
    If a vector (or single value) is provided,
    this should be the scales for each (all) coefficient.
    If None, we assume unbounded uniform priors and ignore the normalization factor.
  """

  def __init__(self, cbv: np.ndarray, prior: np.ndarray | float | None = None):
    nparam = cbv.shape[0]
    if isinstance(prior, np.ndarray):
      if prior.ndim == 2:
        self._add(nparam, cbv, np.linalg.inv(prior), np.linalg.slogdet(prior)[1] / 2)
      else:
        self._add(nparam, cbv, np.diag(1 / prior**2), np.sum(np.log(prior)))
    else:
      icovprior = np.zeros((nparam, nparam))
      logsqdetcovprior = 0
      if prior is not None:
        icovprior[np.diag_indices(nparam)] = 1 / prior**2
        logsqdetcovprior = nparam * np.log(prior)
    super().__init__(
      size=cbv.shape[1],
      nparam=nparam,
      indices=np.array([[0, nparam]]),
      At=cbv,
      icovprior=icovprior,
      logsqdetcovprior=logsqdetcovprior,
    )


class Bspline(LinearMarginalizer):
  r"""
  Uniform cubic B-spline.

  Parameters
  ----------
  t : (n,) ndarray
    Times of measurements.
  tau : float or None
    Timescale of the B-spline (lag between two knots).
    This parameter is ignored if P (and nknot) is provided.
  P : float or None
    Period of the B-spline.
    If provided, the parameter tau is ignored but nknot should be provided instead.
    If None, the spline is not periodic.
  nknot : int or None:
    Number of knots over a period for a periodic B-spline.
  prior : float or None
    Scale of the Gaussian prior on the B-spline parameters.
    If None, we assume unbounded uniform priors and ignore the normalization factor.
  """

  def __init__(
    self,
    t: np.ndarray,
    tau: float | None = None,
    P: float | None = None,
    nknot: int | None = None,
    prior: float | None = None,
  ):
    if P is None:
      t0 = tau * np.floor(t.min() / tau)
      s = (t - t0) / tau
    else:
      tau = P / nknot
      s = (t % P) / tau

    k = np.floor(s).astype(int)
    ds = s - k
    ds2 = ds * ds
    ds3 = ds2 * ds

    uk = np.unique(k)
    if P is None:
      ukdef = np.unique(np.concatenate((uk, uk + 1, uk + 2, uk + 3)))
    else:
      ukdef = np.unique(
        np.concatenate((uk, (uk + 1) % nknot, (uk + 2) % nknot, (uk + 3) % nknot))
      )
    nparam = ukdef.size
    iuk = np.empty(ukdef[-1] + 1, dtype=int)
    iuk[ukdef] = np.arange(nparam)
    ik = iuk[k]
    inds = np.arange(t.size)
    At = np.zeros((nparam, t.size))
    At[ik, inds] = 1 - ds3 + 3 * (ds2 - ds)
    At[(ik + 1) % nparam, inds] = 3 * ds3 - 6 * ds2 + 4
    At[(ik + 2) % nparam, inds] = 3 * (ds + ds2 - ds3) + 1
    At[(ik + 3) % nparam, inds] = ds3

    icovprior = np.zeros((nparam, nparam))
    logsqdetcovprior = 0
    if prior is not None:
      icovprior[np.diag_indices(nparam)] = 1 / prior**2
      logsqdetcovprior = nparam * np.log(prior)

    super().__init__(
      size=t.size,
      nparam=nparam,
      indices=np.array([[0, nparam]]),
      At=At,
      icovprior=icovprior,
      logsqdetcovprior=logsqdetcovprior,
    )
