"""
Implement the ES-MDA algorithms.

@author: acollet
"""
from typing import Any, Callable, Dict, List, Optional, Sequence

import numpy as np
import numpy.typing as npt

from pyesmda.utils import (
    approximate_cov_mm,
    approximate_covariance_matrix_from_ensembles,
    check_nans_in_predictions,
    inflate_ensemble_around_its_mean,
)

# pylint: disable=C0103 # Does not conform to snake_case naming style


class ESMDA:
    r"""
    Ensemble Smoother with Multiple Data Assimilations.

    Implement the ES-MDA as proposed by  Emerick, A. A. and A. C. Reynolds
    :cite:p:`emerickEnsembleSmootherMultiple2013,
    emerickHistoryMatchingProductionSeismic2013`.

    Attributes
    ----------
    d_dim : int
        Number of observation values :math:`N_{obs}`, and consequently of
        predicted values.
    obs : npt.NDArray[np.float64]
        Obsevrations vector with dimensions (:math:`N_{obs}`).
    cov_obs: npt.NDArray[np.float64]
        Covariance matrix of observed data measurement errors with dimensions
        (:math:`N_{obs}`, :math:`N_{obs}`). Also denoted :math:`R`.
    d_obs_uc: npt.NDArray[np.float64]
        Vectors of pertubed observations with dimension
        (:math:`N_{e}`, :math:`N_{obs}`).
    d_pred: npt.NDArray[np.float64]
        Vectors of predicted values (one for each ensemble member)
        with dimensions (:math:`N_{e}`, :math:`N_{obs}`).
    d_history: List[npt.NDArray[np.float64]]
        List of vectors of predicted values obtained at each assimilation step.
    m_prior:
        Vectors of parameter values (one vector for each ensemble member) used in the
        last assimilation step. Dimensions are (:math:`N_{e}`, :math:`N_{m}`).
    m_bounds : npt.NDArray[np.float64]
        Lower and upper bounds for the :math:`N_{m}` parameter values.
        Expected dimensions are (:math:`N_{m}`, 2) with lower bounds on the first
        column and upper on the second one.
    m_history: List[npt.NDArray[np.float64]]
        List of successive `m_prior`.
    cov_md: npt.NDArray[np.float64]
        Cross-covariance matrix between the forecast state vector and predicted data.
        Dimensions are (:math:`N_{m}, N_{obs}`).
    cov_obsd: npt.NDArray[np.float64]
        Autocovariance matrix of predicted data.
        Dimensions are (:math:`N_{obs}, N_{obs}`).
    cov_mm: npt.NDArray[np.float64]
        Autocovariance matrix of estimated parameters.
        Dimensions are (:math:`N_{m}, N_{m}`).
    forward_model: callable
        Function calling the non-linear observation model (forward model)
        for all ensemble members and returning the predicted data for
        each ensemble member.
    forward_model_args: Tuple[Any]
        Additional args for the callable forward_model.
    forward_model_kwargs: Dict[str, Any]
        Additional kwargs for the callable forward_model.
    n_assimilations : int
        Number of data assimilations (:math:`N_{a}`).
    cov_obs_inflation_factors : List[float]
        List of multiplication factor used to inflate the covariance matrix of the
        measurement errors.
    cov_mm_inflation_factors: List[float]
        List of factors used to inflate the adjusted parameters covariance among
        iterations:
        Each realization of the ensemble at the end of each update step i,
        is linearly inflated around its mean.
        See :cite:p:`andersonExploringNeedLocalization2007`.
    dd_correlation_matrix : Optional[npt.NDArray[np.float64]]
        Correlation matrix based on spatial and temporal distances between
        observations and observations :math:`\rho_{DD}`. It is used to localize the
        autocovariance matrix of predicted data by applying an elementwise
        multiplication by this matrix.
        Expected dimensions are (:math:`N_{obs}`, :math:`N_{obs}`).
    md_correlation_matrix : Optional[npt.NDArray[np.float64]]
        Correlation matrix based on spatial and temporal distances between
        parameters and observations :math:`\rho_{MD}`. It is used to localize the
        cross-covariance matrix between the forecast state vector (parameters)
        and predicted data by applying an elementwise
        multiplication by this matrix.
        Expected dimensions are (:math:`N_{m}`, :math:`N_{obs}`).
    save_ensembles_history: bool
        Whether to save the history predictions and parameters over the assimilations.
    rng: np.random.Generator
        The random number generator used in the predictions perturbation step.
    is_forecast_for_last_assimilation: bool
        Whether to compute the predictions for the ensemble obtained at the
        last assimilation step.
    """
    # pylint: disable=R0902 # Too many instance attributes
    __slots__: List[str] = [
        "obs",
        "_cov_obs",
        "d_obs_uc",
        "d_pred",
        "d_history",
        "m_prior",
        "_m_bounds",
        "m_history",
        "cov_md",
        "cov_dd",
        "forward_model",
        "forward_model_args",
        "forward_model_kwargs",
        "_n_assimilations",
        "_assimilation_step",
        "_cov_obs_inflation_factors",
        "_cov_mm_inflation_factors",
        "_dd_correlation_matrix",
        "_md_correlation_matrix",
        "save_ensembles_history",
        "rng",
        "is_forecast_for_last_assimilation",
    ]

    def __init__(
        self,
        obs: npt.NDArray[np.float64],
        m_init: npt.NDArray[np.float64],
        cov_obs: npt.NDArray[np.float64],
        forward_model: Callable[..., npt.NDArray[np.float64]],
        forward_model_args: Sequence[Any] = (),
        forward_model_kwargs: Optional[Dict[str, Any]] = None,
        n_assimilations: int = 4,
        cov_obs_inflation_factors: Optional[Sequence[float]] = None,
        cov_mm_inflation_factors: Optional[Sequence[float]] = None,
        dd_correlation_matrix: Optional[npt.NDArray[np.float64]] = None,
        md_correlation_matrix: Optional[npt.NDArray[np.float64]] = None,
        m_bounds: Optional[npt.NDArray[np.float64]] = None,
        save_ensembles_history: bool = False,
        seed: Optional[int] = None,
        is_forecast_for_last_assimilation: bool = True,
    ) -> None:
        # pylint: disable=R0913 # Too many arguments
        # pylint: disable=R0914 # Too many local variables
        r"""Construct the instance.

        Parameters
        ----------
        obs : npt.NDArray[np.float64]
            Obsevrations vector with dimension :math:`N_{obs}`.
        m_init : npt.NDArray[np.float64]
            Initial ensemble of parameters vector with dimensions
            (:math:`N_{e}`, :math:`N_{m}`).
        cov_obs: npt.NDArray[np.float64]
            Covariance matrix of observed data measurement errors with dimensions
            (:math:`N_{obs}`, :math:`N_{obs}`). Also denoted :math:`R`.
        forward_model: callable
            Function calling the non-linear observation model (forward model)
            for all ensemble members and returning the predicted data for
            each ensemble member.
        forward_model_args: Optional[Tuple[Any]]
            Additional args for the callable forward_model. The default is None.
        forward_model_kwargs: Optional[Dict[str, Any]]
            Additional kwargs for the callable forward_model. The default is None.
        n_assimilations : int, optional
            Number of data assimilations (:math:`N_{a}`). The default is 4.
        cov_obs_inflation_factors : Optional[Sequence[float]]
            Multiplication factor used to inflate the covariance matrix of the
            measurement errors.
            Must match the number of data assimilations (:math:`N_{a}`).
            The default is None.
        cov_mm_inflation_factors: Optional[Sequence[float]]
            List of factors used to inflate the adjusted parameters covariance
            among iterations:
            Each realization of the ensemble at the end of each update step i,
            is linearly inflated around its mean.
            Must match the number of data assimilations (:math:`N_{a}`).
            See :cite:p:`andersonExploringNeedLocalization2007`.
            If None, the default is 1.0. at each iteration (no inflation).
            The default is None.
        dd_correlation_matrix : Optional[npt.NDArray[np.float64]]
            Correlation matrix based on spatial and temporal distances between
            observations and observations :math:`\rho_{DD}`. It is used to localize the
            autocovariance matrix of predicted data by applying an elementwise
            multiplication by this matrix.
            Expected dimensions are (:math:`N_{obs}`, :math:`N_{obs}`).
            The default is None.
        md_correlation_matrix : Optional[npt.NDArray[np.float64]]
            Correlation matrix based on spatial and temporal distances between
            parameters and observations :math:`\rho_{MD}`. It is used to localize the
            cross-covariance matrix between the forecast state vector (parameters)
            and predicted data by applying an elementwise
            multiplication by this matrix.
            Expected dimensions are (:math:`N_{m}`, :math:`N_{obs}`).
            The default is None.
        m_bounds : Optional[npt.NDArray[np.float64]], optional
            Lower and upper bounds for the :math:`N_{m}` parameter values.
            Expected dimensions are (:math:`N_{m}`, 2) with lower bounds on the first
            column and upper on the second one. The default is None.
        save_ensembles_history: bool, optional
            Whether to save the history predictions and parameters over
            the assimilations. The default is False.
        seed: Optional[int]
            Seed for the white noise generator used in the perturbation step.
            If None, the default :func:`numpy.random.default_rng()` is used.
            The default is None.
        is_forecast_for_last_assimilation: bool, optional
            Whether to compute the predictions for the ensemble obtained at the
            last assimilation step. The default is True.

        """
        self.obs: npt.NDArray[np.float64] = obs
        self.m_prior: npt.NDArray[np.float64] = m_init
        self.save_ensembles_history: bool = save_ensembles_history
        self.m_history: list[npt.NDArray[np.float64]] = []
        self.d_history: list[npt.NDArray[np.float64]] = []
        self.d_pred: npt.NDArray[np.float64] = np.zeros([self.n_ensemble, self.d_dim])
        self.cov_obs = cov_obs
        self.d_obs_uc: npt.NDArray[np.float64] = np.array([])
        self.cov_md: npt.NDArray[np.float64] = np.array([])
        self.cov_dd: npt.NDArray[np.float64] = np.array([])
        self.forward_model: Callable = forward_model
        self.forward_model_args: Sequence[Any] = forward_model_args
        if forward_model_kwargs is None:
            forward_model_kwargs = {}
        self.forward_model_kwargs: Dict[str, Any] = forward_model_kwargs
        self._set_n_assimilations(n_assimilations)
        self._assimilation_step: int = 0
        self.set_cov_obs_inflation_factors(cov_obs_inflation_factors)
        self.cov_mm_inflation_factors = cov_mm_inflation_factors
        self.dd_correlation_matrix = dd_correlation_matrix
        self.md_correlation_matrix = md_correlation_matrix
        self.m_bounds = m_bounds
        self.rng: np.random.Generator = np.random.default_rng(seed)
        self.is_forecast_for_last_assimilation = is_forecast_for_last_assimilation

    @property
    def n_assimilations(self) -> int:
        """Return the number of assimilations to perform. Read-only attribute."""
        return self._n_assimilations

    def _set_n_assimilations(self, n: int) -> None:
        """Set the number of assimilations to perform."""
        if not isinstance(n, int):
            raise TypeError("The number of assimilations must be a positive integer.")
        if n < 1:
            raise ValueError("The number of assimilations must be 1 or more.")
        self._n_assimilations = n

    @property
    def n_ensemble(self) -> int:
        """Return the number of ensemble members."""
        return self.m_prior.shape[0]

    @property
    def m_dim(self) -> int:
        """Return the length of the parameters vector."""
        return self.m_prior.shape[1]

    @property
    def d_dim(self) -> int:
        """Return the number of forecast data."""
        return len(self.obs)

    @property
    def cov_obs(self) -> npt.NDArray[np.float64]:
        """Get the observation errors covariance matrix."""
        return self._cov_obs

    @cov_obs.setter
    def cov_obs(self, s: npt.NDArray[np.float64]) -> None:
        """Set the observation errors covariance matrix."""
        if len(s.shape) != 2 or s.shape[0] != s.shape[1] or s.shape[0] != self.d_dim:
            raise ValueError(
                "cov_obs must be a 2D square matrix with "
                f"dimensions ({self.d_dim}, {self.d_dim})."
            )
        self._cov_obs: npt.NDArray[np.float64] = s

    @property
    def cov_mm(self) -> npt.NDArray[np.float64]:
        r"""
        Get the estimated parameters autocovariance matrix. It is a read-only attribute.

        The covariance matrice :math:`C^{l}_{MM}`
        is approximated from the ensemble in the standard way of EnKF
        :cite:p:`evensenDataAssimilationEnsemble2007,aanonsenEnsembleKalmanFilter2009`:

        .. math::
           C^{l}_{MM} = \frac{1}{N_{e} - 1} \sum_{j=1}^{N_{e}}\left(m^{l}_{j} -
           \overline{m^{l}}\right)\left(m^{l}_{j}
           - \overline{m^{l}} \right)^{T}

        with :math:`\overline{m^{l}}`, the parameters
        ensemble means, at iteration :math:`l`.
        """
        return approximate_cov_mm(self.m_prior)

    @property
    def m_bounds(self) -> npt.NDArray[np.float64]:
        """Get the parameter errors covariance matrix."""
        return self._m_bounds

    @m_bounds.setter
    def m_bounds(self, mb: Optional[npt.NDArray[np.float64]]) -> None:
        """Set the parameter errors covariance matrix."""
        if mb is None:
            # In that case, create an array of nan.
            self._m_bounds: npt.NDArray[np.float64] = np.empty(
                [self.m_dim, 2], dtype=np.float64
            )
            self._m_bounds[:] = np.nan
        elif mb.shape[0] != self.m_dim:
            raise ValueError(
                f"m_bounds is of shape {mb.shape} while it "
                f"should be of shape ({self.m_dim}, 2)"
            )
        else:
            self._m_bounds = mb

    @property
    def cov_obs_inflation_factors(self) -> List[float]:
        r"""
        Get the inlfation factors for the covariance matrix of the measurement errors.

        Single and multiple data assimilation are equivalent for the
        linear-Gaussian case as long as the factor :math:`\alpha_{l}` used to
        inflate the covariance matrix of the measurement errors satisfy the
        following condition:

        .. math::
            \sum_{l=1}^{N_{a}} \frac{1}{\alpha_{l}} = 1

        In practise, :math:`\alpha_{l} = N_{a}` is a good choice
        :cite:p:`emerickEnsembleSmootherMultiple2013`.
        """
        return self._cov_obs_inflation_factors

    def set_cov_obs_inflation_factors(self, a: Optional[Sequence[float]]) -> None:
        """Set the inflation factors the covariance matrix of the measurement errors."""
        if a is None:
            self._cov_obs_inflation_factors: List[float] = [
                1 / self.n_assimilations
            ] * self.n_assimilations
        elif len(a) != self.n_assimilations:
            raise ValueError(
                "The length of cov_obs_inflation_factors should match n_assimilations"
            )
        else:
            self._cov_obs_inflation_factors = list(a)

    @property
    def cov_mm_inflation_factors(self) -> List[float]:
        r"""
        Get the inlfation factors for the adjusted parameters covariance matrix.

        Covariance inflation is a method used to counteract the tendency of ensemble
        Kalman methods to underestimate the uncertainty because of either undersampling,
        inbreeding, or spurious correlations
        :cite:p:`todaroAdvancedTechniquesSolving2021`.
        The spread of the ensemble is artificially increased before the assimilation
        of the observations, according to scheme introduced by
        :cite:`andersonExploringNeedLocalization2007`:

        .. math::
            m^{l}_{j} \leftarrow r^{l}\left(m^{l}_{j} - \frac{1}{N_{e}}
            \sum_{j}^{N_{e}}m^{l}_{j}\right) + \frac{1}{N_{e}}\sum_{j}^{N_{e}}m^{l}_{j}

        where :math:`r` is the inflation factor for the assimilation step :math:`l`.
        """
        return list(self._cov_mm_inflation_factors)

    @cov_mm_inflation_factors.setter
    def cov_mm_inflation_factors(self, a: Optional[Sequence[float]]) -> None:
        """
        Set the inflation factors the adjusted parameters covariance matrix.

        If no values have been provided by the user, the default is 1.0. at
        each iteration (no inflation).
        """
        if a is None:
            self._cov_mm_inflation_factors: List[float] = [1.0] * self.n_assimilations
        elif len(a) != self.n_assimilations:
            raise ValueError(
                "The length of cov_mm_inflation_factors should match n_assimilations"
            )
        else:
            self._cov_mm_inflation_factors = list(a)

    @property
    def dd_correlation_matrix(self) -> Optional[npt.NDArray[np.float64]]:
        """Get the observations-observations localization matrix."""
        return self._dd_correlation_matrix

    @dd_correlation_matrix.setter
    def dd_correlation_matrix(self, s: Optional[npt.NDArray[np.float64]]) -> None:
        """Set the observations-observations localization matrix."""
        if s is None:
            self._dd_correlation_matrix = None
            return
        if len(s.shape) != 2 or s.shape[0] != s.shape[1] or s.shape[0] != self.d_dim:
            raise ValueError(
                "dd_correlation_matrix must be a 2D square matrix with "
                f"dimensions ({self.d_dim}, {self.d_dim})."
            )
        self._dd_correlation_matrix: Optional[npt.NDArray[np.float64]] = s

    @property
    def md_correlation_matrix(self) -> Optional[npt.NDArray[np.float64]]:
        """Get the parameters-observations localization matrix."""
        return self._md_correlation_matrix

    @md_correlation_matrix.setter
    def md_correlation_matrix(self, s: Optional[npt.NDArray[np.float64]]) -> None:
        """Set the parameters-observations localization matrix."""
        if s is None:
            self._md_correlation_matrix = None
            return
        if len(s.shape) != 2:
            raise ValueError(
                "md_correlation_matrix must be a 2D matrix with "
                f"dimensions ({self.m_dim}, {self.d_dim})."
            )
        if s.shape[0] != self.m_dim or s.shape[1] != self.d_dim:
            raise ValueError(
                "md_correlation_matrix must be a 2D matrix with "
                f"dimensions ({self.m_dim}, {self.d_dim})."
            )
        self._md_correlation_matrix: Optional[npt.NDArray[np.float64]] = s

    def solve(self) -> None:
        """Solve the optimization problem with ES-MDA algorithm."""
        if self.save_ensembles_history:
            self.m_history.append(self.m_prior)  # save m_init
        for self._assimilation_step in range(self.n_assimilations):
            print(f"Assimilation # {self._assimilation_step + 1}")
            self._forecast()
            self._pertrub(self.cov_obs_inflation_factors[self._assimilation_step])
            self._approximate_covariance_matrices()
            # Update the prior parameter for next iteration
            self.m_prior = self._apply_bounds(
                inflate_ensemble_around_its_mean(
                    self._analyse(
                        self.cov_obs_inflation_factors[self._assimilation_step]
                    ),
                    self.cov_mm_inflation_factors[self._assimilation_step],
                )
            )
            # Saving the parameters history
            if self.save_ensembles_history:
                self.m_history.append(self.m_prior)

        if self.is_forecast_for_last_assimilation:
            self._forecast()

    def _forecast(self) -> None:
        r"""
        Forecast step of ES-MDA.

        Run the forward model from time zero until the end of the historical
        period from time zero until the end of the historical period to
        compute the vector of predicted data

        .. math::
            d^{l}_{j}=g\left(m^{l}_{j}\right),\textrm{for }j=1,2,...,N_{e},

        where :math:`g(·)` denotes the nonlinear observation model, i.e.,
        :math:`d^{l}_{j}` is the :math:`N_{d}`-dimensional vector of predicted
        data obtained by running
        the forward model reservoir simulation with the model parameters given
        by the vector :math:`m^{l}_{j}` from time zero. Note that we use
        :math:`N_{d}` to denote the total number of measurements in the entire
        history.
        """
        self.d_pred = self.forward_model(
            self.m_prior, *self.forward_model_args, **self.forward_model_kwargs
        )
        if self.save_ensembles_history:
            self.d_history.append(self.d_pred)

        # Check if no nan values are found in the predictions.
        # If so, stop the assimilation
        check_nans_in_predictions(self.d_pred, self._assimilation_step)

    def _pertrub(self, inflation_factor: float) -> None:
        r"""
        Perturbation of the observation vector step of ES-MDA.

        Perturb the vector of observations

        .. math::
            d^{l}_{uc,j} = d_{obs} + \sqrt{\alpha_{l+1}}C_{D}^{1/2}Z_{d},
            \textrm{for } j=1,2,...,N_{e},

        where :math:`Z_{d} \sim \mathcal{N}(O, I_{N_{d}})`.

        Notes
        -----
        To get reproducible behavior, use a seed when creating the ESMDA instance.

        """
        self.d_obs_uc = np.zeros([self.n_ensemble, self.d_dim])
        for i in range(self.d_dim):
            self.d_obs_uc[:, i] = self.obs[i] + np.sqrt(
                inflation_factor
            ) * self.rng.normal(0, np.abs(self.cov_obs[i, i]), self.n_ensemble)

    def _approximate_covariance_matrices(self) -> None:
        r"""
        Approximate the covariance matrices.

        The covariance matrices :math:`C^{l}_{MD}` and :math:`C^{l}_{DD}`
        are approximated from the ensemble in the standard way of EnKF
        :cite:p:`evensenDataAssimilationEnsemble2007,aanonsenEnsembleKalmanFilter2009`:

        .. math::
           C^{l}_{MD} = \frac{1}{N_{e} - 1} \sum_{j=1}^{N_{e}}\left(m^{l}_{j} -
           \overline{m^{l}}\right)\left(d^{l}_{j}
           - \overline{d^{l}} \right)^{T}

        .. math::
           C^{l}_{DD} = \frac{1}{N_{e} - 1} \sum_{j=1}^{N_{e}}\left(d^{l}_{j}
           -\overline{d^{l}} \right)\left(d^{l}_{j}
           - \overline{d^{l}} \right)^{T}

        with :math:`\overline{m^{l}}` and :math:`\overline{d^{l}}`, the
        the ensemble means, at iteration :math:`l`, of parameters and predictions,
        respectively.

        If defined by the user, covariances localization is applied by element-wise
        multiplication (Schur product or Hadamard product) of the original
        covariance matrix and a distance dependent correlation function
        that smoothly reduces the correlations between points for increasing distances
        and cuts off long-range correlations above a specific distance:

        .. math::
           \tilde{C}^{l}_{MD} = \rho_{MD} \odot C^{l}_{MD}

        .. math::
           \tilde{C}^{l}_{DD} = \rho_{DD} \odot C^{l}_{DD}

        with :math:`\odot` the element wise multiplication.
        """
        self.cov_md = approximate_covariance_matrix_from_ensembles(
            self.m_prior, self.d_pred
        )
        self.cov_dd = approximate_covariance_matrix_from_ensembles(
            self.d_pred, self.d_pred
        )

        # Spatial and temporal localization: obs - obs
        if self.dd_correlation_matrix is not None:
            self.cov_dd *= self.dd_correlation_matrix
        # Spatial and temporal localization: parameters - obs
        if self.md_correlation_matrix is not None:
            self.cov_md *= self.md_correlation_matrix

    def _analyse(self, inflation_factor: float) -> npt.NDArray[np.float64]:
        r"""
        Analysis step of the ES-MDA.

        Update the vector of model parameters using

        .. math::
           m^{l+1}_{j} = m^{l}_{j} + C^{l}_{MD}\left(C^{l}_{DD}+\alpha_{l+1}
           C_{D}\right)^{-1} \left(d^{l}_{uc,j} - d^{l}_{j} \right),
           \textrm{for } j=1,2,...,N_{e}.

        Notes
        -----
        To avoid the inversion of :math:`\left(C^{l}_{DD}+\alpha_{l+1} C_{D}\right)`,
        the product :math:`\left(C^{l}_{DD}+\alpha_{l+1} C_{D}\right) ^{-1}
        \left(d^{l}_{uc,j} - d^{l}_{j} \right)`
        is solved linearly as :math:`A^{-1}b = x`
        which is equivalent to solve :math:`Ax = b`.

        """
        # predicted parameters
        m_pred: npt.NDArray[np.float64] = np.zeros([self.n_ensemble, self.m_dim])
        for j in range(self.n_ensemble):
            m_pred[j, :] = self.m_prior[j, :] + np.matmul(
                self.cov_md,
                np.linalg.solve(
                    self.cov_dd + inflation_factor * self.cov_obs,
                    self.d_obs_uc[j, :] - self.d_pred[j, :],
                ),
            )
        return m_pred

    def _apply_bounds(self, m_pred: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        """Apply bounds constraints to the adjusted parameters."""
        m_pred = np.where(
            m_pred < self.m_bounds[:, 0], self.m_bounds[:, 0], m_pred
        )  # lower bounds
        m_pred = np.where(
            m_pred > self.m_bounds[:, 1], self.m_bounds[:, 1], m_pred
        )  # upper bounds

        return m_pred
