"""
Contains (publicly available, but not officially supported) experimental
features of iterative_ensemble_smoother
"""

import logging
import numbers
import warnings
from typing import Callable, List, Optional, Sequence, Tuple, TypeVar, Union

import numpy as np
import numpy.typing as npt
import scipy as sp  # type: ignore
from joblib import Parallel, delayed

from iterative_ensemble_smoother import ESMDA
from iterative_ensemble_smoother.esmda import BaseESMDA
from iterative_ensemble_smoother.esmda_inversion import (
    empirical_cross_covariance,
    singular_values_to_keep,
)
from iterative_ensemble_smoother.utils import (
    calc_max_number_of_layers_per_batch_for_distance_localization,
    memory_usage_decorator,
)

logger = logging.getLogger(__name__)

T = TypeVar("T")

ENABLE_MEMORY_LOGGING_FOR_DISTANCE_BASED_LOCALIZATION = True


class AdaptiveESMDA(BaseESMDA):
    @staticmethod
    def correlation_threshold(ensemble_size: int) -> float:
        """Return a number that determines whether a correlation is significant.

        Default threshold taken from luo2022,
        Continuous Hyper-parameter OPtimization (CHOP) in an ensemble Kalman filter
        Section 2.3 - Localization in the CHOP problem

        Examples
        --------
        >>> AdaptiveESMDA.correlation_threshold(0)
        1.0
        >>> AdaptiveESMDA.correlation_threshold(9)
        1.0
        >>> AdaptiveESMDA.correlation_threshold(16)
        0.75
        >>> AdaptiveESMDA.correlation_threshold(36)
        0.5
        """
        return float(min(1, max(0, 3 / np.sqrt(ensemble_size))))

    @staticmethod
    def compute_cross_covariance_multiplier(
        *,
        alpha: float,
        C_D: npt.NDArray[np.double],
        D: npt.NDArray[np.double],
        Y: npt.NDArray[np.double],
        cov_YY: Optional[npt.NDArray[np.double]] = None,
    ) -> npt.NDArray[np.double]:
        """Compute transition matrix T such that:

            X + cov_XY @ transition_matrix

        In the notation of Emerick et al, recall that the update equation is:

            X_posterior = X_prior + cov_XY @ inv(C_DD + alpha * C_D) @ (D - Y)

        This function computes a transition matrix T, defined as:

            T := inv(C_DD + alpha * C_D) @ (D - Y)

        Note that this might not be the most efficient way to compute the update,
        since cov_XY := center(X) @ center(Y) / (N - 1), and we can avoid creating
        this (huge) covariance matrix by performing multiplications right-to-left:

            cov_XY @ inv(C_DD + alpha * C_D) @ (D - Y)
            center(X) @ center(Y) / (N - 1) @ inv(C_DD + alpha * C_D) @ (D - Y)
            center(X) @ [center(Y) / (N - 1) @ inv(C_DD + alpha * C_D) @ (D - Y)]

        The advantage of forming T instead is that if we already have computed
        cov_XY, then we can apply the same T to a reduced number of rows (parameters)
        in cov_XY.
        """
        # Compute cov(Y, Y) if it was not passed to the function.
        # Pre-computation might be faster, since covariance is commutative with
        # respect to indexing, ie, cov(Y[mask, :], YY[mask, :]) = cov(Y, Y)[mask, mask]
        C_DD = empirical_cross_covariance(Y, Y) if cov_YY is None else cov_YY

        assert C_DD.shape[0] == C_DD.shape[1]
        assert C_DD.shape[0] == Y.shape[0]

        # Arguments for sp.linalg.solve
        solver_kwargs = {
            "overwrite_a": True,
            "overwrite_b": True,
            "assume_a": "pos",  # Assume positive definite matrix (use cholesky)
            "lower": False,  # Only use the upper part while solving
        }

        # Compute T := sp.linalg.inv(C_DD + alpha * C_D) @ (D - Y)
        # by solving the system (C_DD + alpha * C_D) @ T = (D - Y)
        if C_D.ndim == 2:
            # C_D is a covariance matrix
            C_DD += alpha * C_D  # Save memory by mutating
        elif C_D.ndim == 1:
            # C_D is an array, so add it to the diagonal without forming diag(C_D)
            np.fill_diagonal(C_DD, C_DD.diagonal() + alpha * C_D)

        return sp.linalg.solve(C_DD, D - Y, **solver_kwargs)  # type: ignore

    def _cov_to_corr_inplace(
        self,
        cov_XY: npt.NDArray[np.double],
        stds_X: npt.NDArray[np.double],
        stds_Y: npt.NDArray[np.double],
    ) -> None:
        """Convert a covariance matrix to a correlation matrix in-place."""

        # Divide each element of cov_XY by the corresponding standard deviations
        cov_XY /= stds_X[:, np.newaxis]
        cov_XY /= stds_Y[np.newaxis, :]

        # Perform checks and clip values to [-1, 1]
        eps = 1e-8
        if not ((cov_XY.max() <= 1 + eps) and (cov_XY.min() >= -1 - eps)):
            msg = "Cross-correlation matrix has entries not in [-1, 1]."
            msg += f"The min and max values are: {cov_XY.min()} and {cov_XY.max()}"
            warnings.warn(msg)

        return np.clip(cov_XY, a_min=-1, a_max=1, out=cov_XY)

    def _corr_to_cov_inplace(
        self,
        corr_XY: npt.NDArray[np.double],
        stds_X: npt.NDArray[np.double],
        stds_Y: npt.NDArray[np.double],
    ) -> None:
        """Convert a correlation matrix to a covariance matrix in-place."""
        # Multiply each element of corr_XY by the corresponding standard deviations
        corr_XY *= stds_X[:, np.newaxis]
        corr_XY *= stds_Y[np.newaxis, :]
        return corr_XY

    @staticmethod
    def _update_single_parameter(
        param_num: int,
        correlated_responses_mask: npt.NDArray[np.bool_],
        cov_XY_row: npt.NDArray[np.double],
        Y: npt.NDArray[np.double],
        D: npt.NDArray[np.double],
        C_D: npt.NDArray[np.double],
        cov_YY: npt.NDArray[np.double],
        alpha: float,
        compute_cross_covariance_multiplier: Callable,
    ) -> Tuple[int, npt.NDArray[np.double]]:
        """
        Worker function to compute the update for a single parameter.
        """

        cov_XY_subset = cov_XY_row[correlated_responses_mask].reshape(1, -1)

        Y_subset = Y[correlated_responses_mask, :]
        D_subset = D[correlated_responses_mask, :]

        cov_YY_mask = np.ix_(correlated_responses_mask, correlated_responses_mask)
        cov_YY_subset = cov_YY[cov_YY_mask]

        C_D_subset = (
            C_D[correlated_responses_mask] if C_D.ndim == 1 else C_D[cov_YY_mask]
        )

        T = compute_cross_covariance_multiplier(
            alpha=alpha,
            C_D=C_D_subset,
            D=D_subset,
            Y=Y_subset,
            cov_YY=cov_YY_subset,
        )

        update_vector = cov_XY_subset @ T

        return param_num, update_vector

    def assimilate(
        self,
        *,
        X: npt.NDArray[np.double],
        Y: npt.NDArray[np.double],
        D: npt.NDArray[np.double],
        overwrite: bool = False,
        alpha: float,
        correlation_threshold: Union[Callable[[int], float], float, None] = None,
        cov_YY: Optional[npt.NDArray[np.double]] = None,
        progress_callback: Optional[Callable[[Sequence[T]], Sequence[T]]] = None,
        correlation_callback: Optional[Callable[[npt.NDArray[np.double]], None]] = None,
        n_jobs: int = 1,
    ) -> npt.NDArray[np.double]:
        """Assimilate data and return an updated ensemble X_posterior.

            X_posterior = smoother.assimilate(X, Y, D, alpha)

        This method first computes the cross-covariance and cross-correlation
        matrices between X and Y. Then it sets correlations that are below the
        threshold to zero. It then loops over parameters in X that are deemed
        significant with respect to Y (based on the threshold), and updates
        these groups together.

        Parameters
        ----------
        X : np.ndarray
            A 2D array of shape (num_parameters, ensemble_size). Each row
            corresponds to a parameter in the model, and each column corresponds
            to an ensemble member (realization).
        Y : np.ndarray
            2D array of shape (num_observations, ensemble_size), containing
            responses when evaluating the model at X. In other words, Y = g(X),
            where g is the forward model.
        D : np.ndarray
            2D array of shape (num_observations, ensemble_size), containing
            perturbed observations. D = observations + mv_normal(0, covariance),
            and D may be computed by the method `perturb_observations`.
        alpha : float
            The covariance inflation factor. The sequence of alphas should
            obey the equation sum_i (1/alpha_i) = 1. However, this is NOT
            enforced in this method. The user/caller is responsible for this.
        overwrite: bool
            If True, X will be overwritten and mutated.
            If False, the method will not mutate inputs in any way.
            Setting this to True saves memory.
        correlation_threshold : callable or float or None
            Either a callable with signature f(ensemble_size) -> float, or a
            float in the range [0, 1]. Entries in the covariance matrix that
            are lower than the correlation threshold will be set to zero.
            If None, the default 3/sqrt(ensemble_size) is used.
        cov_YY : np.ndarray or None
            A 2D array of shape (num_observations, num_observations) with the
            empirical covariance of Y. If passed, this is not computed in the
            method call, potentially saving time and computation.
        progress_callback : Callable[[Sequence[T]], Sequence[T]] or None
            A callback function that can be used to wrap the iteration over
            parameters for progress reporting.
            It should accept an iterable as input and return an iterable.
            This allows for integration with progress reporting tools like tqdm,
            which can provide visual feedback on the progress of the
            assimilation process.
            If None, no progress reporting is performed.
        correlation_callback : Optional[Callable]
            A callback function that is called with the correlation matrix (2D array)
            as its argument after the correlation matrix computation is complete.
            The callback should handle or process the correlation matrix, such as
            saving or logging it. The callback should not return any value.
        n_jobs : int
            The number of parallel jobs to run. If 1, no parallel processing
            is used. If -1, all available CPU cores are used. Default is 1.

        Returns
        -------
        X_posterior : np.ndarray
            2D array of shape (num_parameters, ensemble_size).
        """
        assert X.shape[1] == Y.shape[1]

        # Check the correlation threshold
        is_callable = callable(correlation_threshold)
        is_float = (
            isinstance(correlation_threshold, numbers.Real)
            and correlation_threshold >= 0
            and correlation_threshold <= 1
        )
        is_None = correlation_threshold is None
        if not (is_callable or is_float or is_None):
            raise TypeError(
                "`correlation_threshold` must be a callable or a float in [0, 1]"
            )

        # Do not overwrite input arguments
        if not overwrite:
            X = np.copy(X)

        # Create `correlation_threshold` if the argument is a float
        if is_float:
            corr_threshold: float = correlation_threshold

            def correlation_threshold(ensemble_size: int) -> float:
                return corr_threshold

        # Default correlation threshold function
        if correlation_threshold is None:
            correlation_threshold = self.correlation_threshold
        assert callable(correlation_threshold), (
            "`correlation_threshold` should be callable"
        )

        if progress_callback is None:

            def progress_callback(x):
                return x  # A simple pass-through function

        # Step 1: # Compute cross-correlation between parameters X and responses Y
        # Note: let the number of parameters be n and the number of responses be m.
        # This step requires both O(mn) computation and O(mn) storage, which is
        # larger than the O(n + m^2) computation used in ESMDA, which never explicitly
        # forms the cross-covariance matrix between X and Y. However, if we batch
        # X by parameter group (rows), the storage requirement decreases
        cov_XY = empirical_cross_covariance(X, Y)
        assert cov_XY.shape == (X.shape[0], Y.shape[0])

        stds_X = np.std(X, axis=1, ddof=1)
        stds_Y = np.std(Y, axis=1, ddof=1)

        corr_XY = self._cov_to_corr_inplace(cov_XY, stds_X, stds_Y)

        # Determine which elements in the cross covariance matrix that will
        # be set to zero
        threshold = correlation_threshold(X.shape[1])
        significant_corr_XY = np.abs(corr_XY) > threshold

        cov_XY = self._corr_to_cov_inplace(corr_XY, stds_X, stds_Y)

        # Pre-compute the covariance cov(Y, Y) here, and index on it later
        if cov_YY is None:
            cov_YY = empirical_cross_covariance(Y, Y)
        else:
            assert cov_YY.ndim == 2, "'cov_YY' must be a 2D array"
            assert cov_YY.shape == (Y.shape[0], Y.shape[0])

        # Identify rows with at least one significant correlation.
        significant_rows = np.any(significant_corr_XY, axis=1)

        # Loop only over rows with significant correlations
        params_to_update = np.where(significant_rows)[0]

        results = Parallel(n_jobs=n_jobs)(
            delayed(self._update_single_parameter)(
                param_num,
                significant_corr_XY[param_num],
                cov_XY[param_num],
                # These are sent in full, which is efficient as they are the same
                # for all workers. Joblib handles this well.
                Y,
                D,
                self.C_D,
                cov_YY,
                alpha,
                self.compute_cross_covariance_multiplier,
            )
            for param_num in progress_callback(params_to_update)
        )

        for param_num, update_vector in results:
            X[[param_num], :] += update_vector

        if correlation_callback is not None:
            corr_XY = self._cov_to_corr_inplace(cov_XY, stds_X, stds_Y)
            correlation_callback(corr_XY)

        return X


class RowScaling:
    # Illustration of how row scaling works, `multiply` is the important part
    # For the actual implementation, which is more involved, see:
    # https://github.com/equinor/ert/blob/84aad3c56e0e52be27b9966322d93dbb85024c1c/src/clib/lib/enkf/row_scaling.cpp
    def __init__(self, alpha=1.0):
        """Alpha is the strength of the update."""
        assert 0 <= alpha <= 1.0
        self.alpha = alpha

    def multiply(self, X, K):
        """Takes a matrix X and a matrix K and performs alpha * X @ K."""
        # This implementation merely mimics how RowScaling::multiply behaves
        # in the C++ code. It mutates the input argument X instead of returning.
        X[:, :] = X @ (K * self.alpha)


def ensemble_smoother_update_step_row_scaling(
    *,
    covariance: npt.NDArray[np.double],
    observations: npt.NDArray[np.double],
    X_with_row_scaling: List[Tuple[npt.NDArray[np.double], RowScaling]],
    Y: npt.NDArray[np.double],
    seed: Union[np.random._generator.Generator, int, None] = None,
    inversion: str = "exact",
    truncation: float = 1.0,
):
    """Perform a single ESMDA update (ES) with row scaling.
    The matrices in X_with_row_scaling WILL BE MUTATED.
    See the ESMDA class for information about input arguments.


    Explanation of row scaling
    --------------------------

    The ESMDA update can be written as:

        X_post = X_prior + X_prior @ K

    where K is a transition matrix. The core of the row scaling approach is that
    for each row i in the matrix X, we apply an update with strength alpha:

        X_post = X_prior + alpha * X_prior @ K
        X_post = X_prior @ (I + alpha * K)

    Clearly 0 <= alpha <= 1 denotes the 'strength' of the update; alpha == 1
    corresponds to a normal smoother update and alpha == 0 corresponds to no
    update. With the per row transformation of X the operation is no longer matrix
    multiplication but the pseudo code looks like:

        for i in rows:
            X_i_post = X_i_prior @ (I + alpha_i * K)

    See also original code:
        https://github.com/equinor/ert/blob/84aad3c56e0e52be27b9966322d93dbb85024c1c/src/clib/lib/enkf/row_scaling.cpp#L51

    """

    # Create ESMDA instance and set alpha=1 => run single assimilation (ES)
    smoother = ESMDA(
        covariance=covariance,
        observations=observations,
        seed=seed,
        inversion=inversion,
        alpha=1,
    )

    # Create transition matrix - common to all parameters in X
    transition_matrix = smoother.compute_transition_matrix(
        Y=Y, alpha=1, truncation=truncation
    )

    # The transition matrix K is a matrix such that
    #     X_posterior = X_prior + X_prior @ K
    # but the C++ code in ERT requires a transition matrix F that obeys
    #     X_posterior = X_prior @ F
    # To accomplish this, we add the identity to the transition matrix in place
    np.fill_diagonal(transition_matrix, transition_matrix.diagonal() + 1)

    # Loop over groups of rows (parameters)
    for X, row_scale in X_with_row_scaling:
        # In the C++ code, multiply() will transform the transition matrix F as
        #    F_new = F * alpha + I * (1 - alpha)
        # but the transition matrix F that we pass below is F := K + I, so:
        #    F_new = (K + I) * alpha + I * (1 - alpha)
        #    F_new = K * alpha + I * alpha + I - I * alpha
        #    F_new = K * alpha + I
        # And the update becomes : X_posterior = X_prior @ F_new
        # The update in the C++ code is equivalent to
        #    X_posterior = X_prior + alpha * X_prior @ K
        # if we had used the original transition matrix K that is returned from
        # ESMDA.compute_transition_matrix
        row_scale.multiply(X, transition_matrix)

    return X_with_row_scaling


class DistanceESMDA(ESMDA):
    def __init__(
        self,
        covariance: npt.NDArray[np.float64],
        observations: npt.NDArray[np.float64],
        alpha: Union[int, npt.NDArray[np.float64]] = 5,
        seed: Union[np.random._generator.Generator, int, None] = None,
    ) -> None:
        """
        Initialize instance.
        """
        # Initialize instance
        super().__init__(
            covariance=covariance, observations=observations, alpha=alpha, seed=seed
        )

        # Ensure self.X3 is initialized to None
        # Is set in prepare_assimilation and used in assimilate_batch
        # Is not used when using assimilate
        self.X3: npt.NDArray[np.float64] = None

    @memory_usage_decorator(
        enabled=ENABLE_MEMORY_LOGGING_FOR_DISTANCE_BASED_LOCALIZATION
    )
    def assimilate(
        self,
        *,
        X: npt.NDArray[np.float64],
        Y: npt.NDArray[np.float64],
        rho: npt.NDArray[np.float64],
        truncation: float = 0.99,
    ):
        """
        Implementation of algorithm described in Appendix B of Emerick's publication.
        Reference: Emerick, Journal of Petroleum Science and Engineering 136(3):219-239
                   DOI:10.1016/j.petrol.2016.01.029  (2016)

        Calculate Ensemble Smoother update with distance-based localization using the
        RHO matrix to define the scaling/tapering function for reduction of observations
        influence on field parameters. This function implements all steps in
        the algorithm.

        Typical workflow:
          Create object of DistanceESMDA
          Loop over batches of field parameter values from a field parameter:
            Run 'assimilate' for each batch of field parameter values

        Args:
            X: Parameter matrix shape=(nparameters, nrealizations)
            Y: Response matrix with predictions of observations,
               shape=(nobservations, nrealizations)
            rho: Localization matrix with scaling factors,
                 shape= (nparameters, nobservations)
            truncation: Define threshold for how many singular values
                 to include from SVD

        Results:
            Posterior (updated) matrix with parameters,
            shape=(nparameters, nrealizations)


        """

        N_n, N_e = Y.shape

        # Subtract the mean of every parameter, see Eqn (B.4)
        M_delta = X - np.mean(X, axis=1, keepdims=True)

        # Subtract the mean of every observation, see Eqn (B.5)
        D_delta = Y - np.mean(Y, axis=1, keepdims=True)

        # See Eqn (B.8)
        # Compute the diagonal of the inverse of S directly, without forming S itself.
        if self.C_D.ndim == 1:
            # If C_D is 1D, it's a vector of variances. S_inv_diag is 1/sqrt(variances).
            S_inv_diag = 1.0 / np.sqrt(self.C_D)
        else:
            # If C_D is 2D, extract its diagonal of variances, then compute S_inv_diag.
            S_inv_diag = 1.0 / np.sqrt(np.diag(self.C_D))

        # See Eqn (B.10)
        U, w, VT = sp.linalg.svd(S_inv_diag[:, np.newaxis] * D_delta)
        idx = singular_values_to_keep(w, truncation=truncation)
        N_r = min(N_n, N_e - 1, idx)  # Number of values in SVD to keep
        U_r, w_r = U[:, :N_r], w[:N_r]

        # See Eqn (B.12)
        # Calculate C_hat_D, the correlation matrix of measurement errors.
        # This is defined as C_hat_D = S^-1 * C_D * S^-1
        if self.C_D.ndim == 1:
            # If C_D is a 1D vector of variances,
            # it represents a diagonal matrix.
            # In this special case,
            # S_inv * C_D * S_inv simplifies to the identity matrix.
            C_hat_D = np.identity(N_n)
        else:  # C_D is a 2D matrix
            # This scales each ROW i of self.C_D by the scalar S_inv_diag[i].
            # This is numerically identical to the matrix multiplication S⁻¹ @ C_D
            C_hat_D_temp = S_inv_diag[:, np.newaxis] * self.C_D

            # This scales each COLUMN j of C_hat_D_temp by the scalar S_inv_diag[j].
            # This is numerically identical to the matrix multiplication (result) @ S⁻¹
            C_hat_D = C_hat_D_temp * S_inv_diag

        U_r_w_inv = U_r / w_r
        # See Eqn (B.13)
        R = (
            self.alpha
            * (N_e - 1)
            * np.linalg.multi_dot([U_r_w_inv.T, C_hat_D, U_r_w_inv])
        )

        # See Eqn (B.14)
        H_r, Z_r = sp.linalg.eigh(R, driver="evr", overwrite_a=True)

        # See Eqn (B.18)
        _X = (S_inv_diag[:, np.newaxis] * U_r) * (1 / w_r) @ Z_r

        # See Eqn (B.19)
        L = np.diag(1.0 / (1.0 + H_r))

        # See Eqn (B.20)
        X1 = L @ _X.T
        # See Eqn (B.21)
        X2 = D_delta.T @ _X
        # See Eqn (B.22)
        X3 = X2 @ X1

        # See Eqn (B.23)
        K_i = M_delta @ X3

        # See Eqn (B.24)
        K_rho_i = rho * K_i

        D = self.perturb_observations(ensemble_size=N_e, alpha=self.alpha)
        # See Eqn (B.25)
        X4 = K_rho_i @ (D - Y)

        # See Eqn (B.26)
        return X + X4

    def prepare_assimilation(
        self,
        Y: npt.NDArray[np.float64],
        truncation: float = 0.99,
        D: npt.NDArray[np.float64] = None,
    ):
        """
        The part of the algorithm that does not depend on the field parameters,
        but only on observations and responses are calculated here.
        No need to re-calculate this multiple times when running a loop to update
        batches of field parameters. This function implements all steps prior
        to the steps involving the field parameters and should be used as a
        preparation step once for each field parameter before running the
        loop over batches of field parameters where the field parameters
        are updated.

        Typical workflow:
          Create object of DistanceESMDA
          Prepare for update of field parameter by running 'prepare_assimilation'
          Loop over batches of field parameter values from a field parameter:
            Run 'assimilate_batch' for each batch of field parameter values


        Args:
            Y: Response matrix with predictions of observations,
               shape=(nobservations, nrealizations)
            D: Matrix of perturbed observations.
               The shape is (nobservations, nrealizations).
               Note that if D is None, the perturbed observations are drawn
               inside this function using assumption of independence
               between observation errors.

        Results:
            Updated internal matrices such that assimilate_batch function
            can be used without having to re-calculate all steps not involving
            the field parameters, but only the observations.

        """

        N_n, N_e = Y.shape

        # Subtract the mean of every observation, see Eqn (B.5)
        D_delta = Y - np.mean(Y, axis=1, keepdims=True)

        # See Eqn (B.8)
        # Compute the diagonal of the inverse of S directly, without forming S itself.
        if self.C_D.ndim == 1:
            # If C_D is 1D, it's a vector of variances. S_inv_diag is 1/sqrt(variances).
            S_inv_diag = 1.0 / np.sqrt(self.C_D)
        else:
            # If C_D is 2D, extract its diagonal of variances, then compute S_inv_diag.
            S_inv_diag = 1.0 / np.sqrt(np.diag(self.C_D))

        # See Eqn (B.10)
        U, w, VT = sp.linalg.svd(S_inv_diag[:, np.newaxis] * D_delta)
        idx = singular_values_to_keep(w, truncation=truncation)
        N_r = min(N_n, N_e - 1, idx)  # Number of values in SVD to keep
        U_r, w_r = U[:, :N_r], w[:N_r]

        # See Eqn (B.12)
        # Calculate C_hat_D, the correlation matrix of measurement errors.
        # This is defined as C_hat_D = S^-1 * C_D * S^-1
        if self.C_D.ndim == 1:
            # If C_D is a 1D vector of variances,
            # it represents a diagonal matrix.
            # In this special case,
            # S_inv * C_D * S_inv simplifies to the identity matrix.
            C_hat_D = np.identity(N_n)
        else:  # C_D is a 2D matrix
            # This scales each ROW i of self.C_D by the scalar S_inv_diag[i].
            # This is numerically identical to the matrix multiplication S⁻¹ @ C_D
            C_hat_D_temp = S_inv_diag[:, np.newaxis] * self.C_D

            # This scales each COLUMN j of C_hat_D_temp by the scalar S_inv_diag[j].
            # This is numerically identical to the matrix multiplication (result) @ S⁻¹
            C_hat_D = C_hat_D_temp * S_inv_diag

        U_r_w_inv = U_r / w_r
        # See Eqn (B.13)
        R = (
            self.alpha
            * (N_e - 1)
            * np.linalg.multi_dot([U_r_w_inv.T, C_hat_D, U_r_w_inv])
        )

        # See Eqn (B.14)
        H_r, Z_r = sp.linalg.eigh(R, driver="evr", overwrite_a=True)

        # See Eqn (B.18)
        _X = (S_inv_diag[:, np.newaxis] * U_r) * (1 / w_r) @ Z_r

        # See Eqn (B.19)
        L = np.diag(1.0 / (1.0 + H_r))

        # See Eqn (B.20)
        X1 = L @ _X.T
        # See Eqn (B.21)
        X2 = D_delta.T @ _X

        # The matrices X3 and D with perturbed observations is saved
        # and reused in assimilation of each batch of field parameters

        # See Eqn (B.22)
        self.X3 = X2 @ X1

        # Observations with added perturbations
        if D is None:
            self.D = self.perturb_observations(ensemble_size=N_e, alpha=self.alpha)
        else:
            assert self.C_D.shape[0] == D.shape[0]
            self.D = D
        return

    @memory_usage_decorator(
        enabled=ENABLE_MEMORY_LOGGING_FOR_DISTANCE_BASED_LOCALIZATION
    )
    def assimilate_batch(
        self,
        *,
        X_batch: npt.NDArray[np.float64],
        Y: npt.NDArray[np.float64],
        rho_batch: npt.NDArray[np.float64],
        D: npt.NDArray[np.float64] = None,
        truncation: float = 0.99,
    ):
        """
        Implementation of algorithm described in Appendix B of Emerick's publication.
        Reference: Emerick, Journal of Petroleum Science and Engineering 136(3):219-239
                   DOI:10.1016/j.petrol.2016.01.029  (2016)

        Calculate Ensemble Smoother update with distance-based localization using the
        RHO matrix to define the scaling/tapering function for reduction of observations
        influence on field parameters. This function implements the update steps
        of the field parameters and requires that the initial steps in the algorithm
        is already run by using the function 'prepare_assimilation'

        Args:
            X: Parameter matrix shape=(nparameters, nrealizations)
            Y: Response matrix with predictions of observations,
               shape=(nobservations, nrealizations)
            rho_batch: Localization matrix with scaling factors,
                 shape= (nparameters, nobservations)
            D: Perturbed observations, shape=(nobservations,nrealizations)
               Optional, default is to simulate the perturbations internally
               within this class

        Results:
            Posterior (updated) matrix with parameters,
            shape=(nparameters, nrealizations)


        """
        if self.X3 is None:
            # Need to run preparation step
            assert Y.shape[0] == self.observations.shape[0]
            assert Y.shape[1] == X_batch.shape[1]
            self.prepare_assimilation(Y=Y, truncation=truncation, D=D)

        # Subtract the mean of every parameter, see Eqn (B.4)
        M_delta = X_batch - np.mean(X_batch, axis=1, keepdims=True)

        # See Eqn (B.23)
        K_i = M_delta @ self.X3

        # See Eqn (B.24)
        K_rho_i = rho_batch * K_i

        # See Eqn (B.25)
        X4 = K_rho_i @ (self.D - Y)

        # See Eqn (B.26)
        return X_batch + X4

    @memory_usage_decorator(
        enabled=ENABLE_MEMORY_LOGGING_FOR_DISTANCE_BASED_LOCALIZATION
    )
    def update_params_1D(
        self,
        X_prior: npt.NDArray[np.float64],
        Y: npt.NDArray[np.float64],
        rho_1D: npt.NDArray[np.float64],
    ) -> npt.NDArray[np.float64]:
        """
        Calculate posterior update with distance-based ESMDA for one 1D parameter
        The RHO matrix is specified as input.
        Result is posterior parameter matrix

        Args:
            X_prior: Matrix with prior realizations of all field parameters,
                    shape=(nparameters, nrealizations)
            Y: Matrix with response values for each observations for each realization,
                    shape=(nobservations, nrealizations)
            rho_1D: RHO matrix elements for a 1D grid with size nx,
                    shape=(nx,nobservations). If None is input, rho is set to 1
                    for every matrix element and there will be no localization.
        Results:
            X_post: Posterior ensemble of field parameters,
            shape=(nx, nrealizations)
        """
        if Y is None or Y.shape[0] == 0:
            # No observation, no update
            return X_prior

        nobs = Y.shape[0]
        assert self.observations.shape[0] == nobs

        # Get dimensions from rho matrix
        nx, _ = rho_1D.shape
        assert nx == X_prior.shape[0]
        assert nobs == rho_1D.shape[1]
        nreal = X_prior.shape[1]
        assert Y.shape[1] == nreal

        return self.assimilate(X=X_prior, Y=Y, rho=rho_1D)

    @memory_usage_decorator(
        enabled=ENABLE_MEMORY_LOGGING_FOR_DISTANCE_BASED_LOCALIZATION
    )
    def update_params_2D(
        self,
        X_prior: npt.NDArray[np.float64],
        Y: npt.NDArray[np.float64],
        rho_2D: npt.NDArray[np.float64],
    ) -> npt.NDArray[np.float64]:
        """
        Calculate posterior update with distance-based ESMDA for one 2D parameter
        The RHO matrix is specified as input. Result is posterior parameter matrix.
        Number of parameters for the field is nx*ny,
        dimensions of the rho_2D in addition to number of observations.
        Args:
            X_prior: Matrix with prior realizations of all field parameters,
                    shape=(nparameters, nrealizations)
            Y: Matrix with response values for each observations for each realization,
                    shape=(nobservations, nrealizations)
            rho_2D: RHO matrix elements for a 2D grid with size (nx,,ny)
                    shape=(nx,ny, nobservations). If None is input, rho is set to 1
                    for every matrix element and there will be no localization.
                    nx and ny are the dimensions of the 2D field.
        Results:
            X_post: Posterior ensemble of field parameters,
                    shape=(nparameters, nrealizations)
        """
        # No update if no observations or responses
        if Y is None or Y.shape[0] == 0:
            # No update of the field parameters
            # Check if it necessary to make a copy or can we only return X_prior?
            return X_prior.copy()

        nx, ny, nobs = rho_2D.shape
        nparam = nx * ny
        nreal = X_prior.shape[1]
        assert X_prior.shape[0] == nparam, (
            f"Mismatch between X_prior dimension {X_prior.shape[0]} and nparam {nparam}"
        )

        assert nobs == Y.shape[0], (
            "Mismatch between Y matrix dimension for number of observation and rho_2D"
        )
        assert Y.shape[1] == nreal, (
            f"Mismatch between X_prior dimension {Y.shape[1]} and nreal {nreal}"
        )
        assert self.observations.shape[0] == nobs

        rho = rho_2D.reshape(nparam, nobs)
        return self.assimilate_batch(X_batch=X_prior, Y=Y, rho_batch=rho)

    @memory_usage_decorator(
        enabled=ENABLE_MEMORY_LOGGING_FOR_DISTANCE_BASED_LOCALIZATION
    )
    def update_params_3D(
        self,
        X_prior: npt.NDArray[np.float64],
        Y: npt.NDArray[np.float64],
        rho_2D_slice: npt.NDArray[np.float64],
        nz: int,
        min_nbatch: int = 1,
        no_update: bool = False,
    ) -> npt.NDArray[np.float64]:
        """
        Calculate posterior update with distance-based ESMDA for one 3D parameter.
        The RHO for one layer of the 3D field parameter is input.
        This is copied to all other layers of RHO in each batch of grid parameter
        layers since only lateral distance is used when calculating distances.
        Result is posterior parameter matrices of field parameters for one field.
        Number of parameters is nx*ny*nz where nx, ny is dimension of the rho_2D_slice
        in addition to number of observations.

        Args:
            X_prior: Matrix with prior realizations of all field parameters,
                    shape=(nparameters, nrealizations)
            Y: Matrix with response values for each observations for each realization,
                    shape=(nobservations, nrealizations)
            rho_2D_slice: RHO matrix elements for one layer (slice) of a 3D grid
                    with size (nx, ny), shape=(nx,ny,nobservations). If this variable
                    is None, then rho is set to 1 for all elements which means
                    that no localization is used.
            nz:     Number of grid layers for the 3D field parameter.
            min_nbatch: Minimum number of batches the field parameter is split into.
                    Default is 1. Usually number of batches will be calculated based
                    on available memory and the size of the field parameters,
                    number of observations and realizations. The actual number of
                    batches will be max(min_nbatch, min_number_of_batches_required).
            no_update: Is set to True only when testing the batch loop.
                    Should be False to get any update.

        Results:
            X_post: Posterior ensemble of field parameters,
            shape=(nx*ny*nz, nrealizations)
        """
        # No update if no observations or responses
        if Y is None or Y.shape[0] == 0:
            # No update of the field parameters
            # Check if it is it necessary to return a copy here?
            return X_prior.copy()

        nx, ny, nobs = rho_2D_slice.shape

        nparam_per_layer = nx * ny
        nparam = nparam_per_layer * nz
        nreal = X_prior.shape[1]
        assert X_prior.shape[0] == nparam, (
            f"Mismatch between X_prior dimension {X_prior.shape[0]} and nparam {nparam}"
        )
        assert nobs == Y.shape[0], (
            f"Mismatch between number of observations in Rho_2D {nobs} and in Y matrix"
        )
        assert Y.shape[1] == nreal, (
            f"Mismatch between X_prior dimension {Y.shape[1]} and nreal {nreal}"
        )

        X_prior_3D = X_prior.reshape(nx, ny, nz, nreal)

        # Check memory constraints and calculate how many grid layers of
        # field parameters is possible to update on one batch
        max_nlayers_per_batch = (
            calc_max_number_of_layers_per_batch_for_distance_localization(
                nx, ny, nz, nobs, nreal, bytes_per_float=8
            )
        )  # Use float64
        nlayer_per_batch = min(max_nlayers_per_batch, nz)
        nbatch = int(nz / nlayer_per_batch)

        # Number of batches is defined by available memory and wanted number of batches
        # Usually one should have as few batches as possible but sufficient to
        # be able to update a batch with the memory available.
        # It is possible to explicit require more batches than the minimum
        # number required by the memory constraint. Main use case for this
        # is probably only for unit testing to avoid having unit tests running
        # slow due to very big size of the field, the number of observations and
        # and number of realizations.
        nbatch = max(min_nbatch, nbatch)
        nlayer_per_batch = int(nz / nbatch)
        if nbatch * nlayer_per_batch == nz:
            logger.debug(f"Number of batches used: {nbatch}")
        else:
            logger.debug(f"Number of batches used: {nbatch + 1}")

        nparam_in_batch = (
            nparam_per_layer * nlayer_per_batch
        )  # For full sized batch of layers

        nlayer_last_batch = nz - nbatch * nlayer_per_batch

        # Initialize the X_post_3D
        X_post_3D = X_prior_3D.copy()
        for batch_number in range(nbatch):
            start_layer_number = batch_number * nlayer_per_batch
            end_layer_number = start_layer_number + nlayer_per_batch

            X_batch = X_prior_3D[:, :, start_layer_number:end_layer_number, :].reshape(
                (nparam_in_batch, nreal)
            )

            # Copy rho calculated from one layer of 3D parameter into all layers for
            # current batch of layers.
            # Size of rho batch: (nx,ny,nlayer_per_batch,nobs)
            rho_3D_batch = np.zeros((nx, ny, nlayer_per_batch, nobs), dtype=np.float64)
            rho_3D_batch[:, :, :, :] = rho_2D_slice[:, :, np.newaxis, :]
            rho_batch = rho_3D_batch.reshape((nparam_in_batch, nobs))

            if no_update:
                # Only used for test purpose
                X_post_batch = X_batch
            else:
                logger.debug(f"Assimilate batch number {batch_number}")
                X_post_batch = self.assimilate_batch(
                    X_batch=X_batch, Y=Y, rho_batch=rho_batch
                )
            X_post_3D[:, :, start_layer_number:end_layer_number, :] = (
                X_post_batch.reshape(nx, ny, nlayer_per_batch, nreal)
            )

        if nlayer_last_batch > 0:
            batch_number = nbatch
            start_layer_number = batch_number * nlayer_per_batch
            end_layer_number = start_layer_number + nlayer_last_batch
            nparam_in_last_batch = nparam_per_layer * nlayer_last_batch

            X_batch = X_prior_3D[:, :, start_layer_number:end_layer_number, :].reshape(
                (nparam_in_last_batch, nreal)
            )

            rho_3D_batch = np.zeros((nx, ny, nlayer_last_batch, nobs), dtype=np.float64)
            # Copy rho calculated from one layer of 3D parameter into all layers for
            # current batch of layers
            rho_3D_batch[:, :, :, :] = rho_2D_slice[:, :, np.newaxis, :]
            rho_batch = rho_3D_batch.reshape((nparam_in_last_batch, nobs))
            if no_update:
                # Only used for test purpose
                X_post_batch = X_batch
            else:
                logger.debug(f"Assimilate batch number {batch_number}")

                X_post_batch = self.assimilate_batch(
                    X_batch=X_batch, Y=Y, rho_batch=rho_batch
                )
            X_post_3D[:, :, start_layer_number:end_layer_number, :] = (
                X_post_batch.reshape(nx, ny, nlayer_last_batch, nreal)
            )

        return X_post_3D.reshape(nparam, nreal)

    def update_params(
        self,
        X_prior: npt.NDArray[np.float64],
        Y: npt.NDArray[np.float64],
        rho_input: npt.NDArray[np.float64],
        nz: int = 1,
        min_nbatch: int = 1,
        no_update: bool = False,
    ) -> npt.NDArray[np.float64]:
        """
        Calculate posterior update with distance-based ESMDA.
        Depending on the shape of rho_input it will update an 1D , 2D or 3D field.
        Dimension of the ensemble X_prior is either (nx), (nx,ny) or (nx,ny,nz)
        and rho_input define nx and ny while nz is number of layers in the grid
        for the field parameter in 3D.

        Args:
            X_prior: Matrix with prior realizations of all field parameters,
                    shape=(nparameters, nrealizations)
            Y: Matrix with response values for each observations for each realization,
                    shape=(nobservations, nrealizations)
            rho_input: RHO matrix elements.
                    If RHO_input has shape= (nx,nobs),
                    X_prior is an ensemble of 1D fields
                    If RHO_input has shape= (nx, ny, nobs),
                    X_prior is an ensemble of 2D fields.
                    If Rho_input has shape=(nx,ny,nobs) and nz > 1,
                    X_prior is an ensemble of 3D fields with size (nx, ny, nz).
            nz:     Number of grid layers for the 3D field parameter.
            min_nbatch: Minimum number of batches the field parameter is split into.
                    Default is 1. Usually number of batches will be calculated based
                    on available memory and the size of the field parameters,
                    number of observations and realizations. The actual number of
                    batches will be max(min_nbatch, min_number_of_batches_required).
            no_update: Is set to True only when testing the batch loop.
                    Should be False to get any update.

        Results:
            X_post: Posterior ensemble of field parameters,
            shape=(nx*ny*nz, nrealizations)
        """
        if rho_input.ndim == 1:
            return self.update_params_1D(X_prior=X_prior, Y=Y, rho_1D=rho_input)
        if rho_input.ndim == 2 and nz == 1:
            return self.update_params_2D(X_prior=X_prior, Y=Y, rho_2D=rho_input)
        return self.update_params_3D(
            X_prior=X_prior,
            Y=Y,
            rho_2D_slice=rho_input,
            nz=nz,
            min_nbatch=min_nbatch,
            no_update=no_update,
        )


if __name__ == "__main__":
    import pytest

    pytest.main(
        args=[
            __file__,
            "--doctest-modules",
            "-v",
        ]
    )
