import warnings

import numpy as np
import pandas as pd
import sklearn.preprocessing

NONE_REPLACEMENT = -32768


def information_mutual_conditional(
    x,
    y,
    z,
    cartesian_product=False,
    base=2,
    fill_value=-1,
    estimator="ML",
    alphabet_x=None,
    alphabet_y=None,
    Alphabet_Z=None,
    keep_dims=False,
):
    x, fill_value_X = _sanitise_array_input(x, fill_value)
    y, fill_value_Y = _sanitise_array_input(y, fill_value)
    z, fill_value_Z = _sanitise_array_input(z, fill_value)
    if alphabet_x is not None:
        alphabet_x, fill_value_Alphabet_X = _sanitise_array_input(
            alphabet_x, fill_value
        )
        alphabet_x, _ = _autocreate_alphabet(alphabet_x, fill_value_Alphabet_X)
    else:
        alphabet_x, fill_value_Alphabet_X = _autocreate_alphabet(x, fill_value_X)
    if alphabet_y is not None:
        alphabet_y, fill_value_Alphabet_Y = _sanitise_array_input(
            alphabet_y, fill_value
        )
        alphabet_y, _ = _autocreate_alphabet(alphabet_y, fill_value_Alphabet_Y)
    else:
        alphabet_y, fill_value_Alphabet_Y = _autocreate_alphabet(y, fill_value_Y)
    if Alphabet_Z is not None:
        Alphabet_Z, fill_value_Alphabet_Z = _sanitise_array_input(
            Alphabet_Z, fill_value
        )
        Alphabet_Z, _ = _autocreate_alphabet(Alphabet_Z, fill_value_Alphabet_Z)
    else:
        Alphabet_Z, fill_value_Alphabet_Z = _autocreate_alphabet(z, fill_value_Z)

    if x.size == 0:
        raise ValueError("arg X contains no elements")
    if y.size == 0:
        raise ValueError("arg Y contains no elements")
    if z.size == 0:
        raise ValueError("arg Z contains no elements")
    if np.any(_isnan(x)):
        raise ValueError("arg X contains NaN values")
    if np.any(_isnan(y)):
        raise ValueError("arg Y contains NaN values")
    if np.any(_isnan(z)):
        raise ValueError("arg Z contains NaN values")
    if alphabet_x.size == 0:
        raise ValueError("arg Alphabet_X contains no elements")
    if np.any(_isnan(alphabet_x)):
        raise ValueError("arg Alphabet_X contains NaN values")
    if alphabet_y.size == 0:
        raise ValueError("arg Alphabet_Y contains no elements")
    if np.any(_isnan(alphabet_y)):
        raise ValueError("arg Alphabet_Y contains NaN values")
    if Alphabet_Z.size == 0:
        raise ValueError("arg Alphabet_Z contains no elements")
    if np.any(_isnan(Alphabet_Z)):
        raise ValueError("arg Alphabet_Z contains NaN values")
    if _isnan(fill_value_X):
        raise ValueError("fill value for arg X is NaN")
    if _isnan(fill_value_Y):
        raise ValueError("fill value for arg Y is NaN")
    if _isnan(fill_value_Z):
        raise ValueError("fill value for arg Z is NaN")
    if x.shape[:-1] != alphabet_x.shape[:-1]:
        raise ValueError("leading dimensions of args X and Alphabet_X do not " "match")
    if y.shape[:-1] != alphabet_y.shape[:-1]:
        raise ValueError("leading dimensions of args Y and Alphabet_Y do not " "match")
    if z.shape[:-1] != Alphabet_Z.shape[:-1]:
        raise ValueError("leading dimensions of args Z and Alphabet_Z do not " "match")
    if not cartesian_product and (x.shape != y.shape or x.shape != z.shape):
        raise ValueError("dimensions of args X, Y, Z do not match")
    if cartesian_product and (x.shape[-1] != y.shape[-1] or x.shape[-1] != z.shape[-1]):
        raise ValueError("trailing dimensions of args X, Y, Z do not match")
    if not (np.isscalar(base) and np.isreal(base) and base > 0):
        raise ValueError("arg base not a positive real-valued scalar")

    S, fill_value = _map_observations_to_integers(
        (x, alphabet_x, y, alphabet_y, z, Alphabet_Z),
        (
            fill_value_X,
            fill_value_Alphabet_X,
            fill_value_Y,
            fill_value_Alphabet_Y,
            fill_value_Z,
            fill_value_Alphabet_Z,
        ),
    )
    x, alphabet_x, y, alphabet_y, z, Alphabet_Z = S

    if not cartesian_product:
        I = np.empty(x.shape[:-1])
        if I.ndim > 0:
            I[:] = np.NaN
        else:
            I = np.float64(np.NaN)
    else:
        shapeI_Z = z.shape[:-1]
        z = np.reshape(z, (-1, z.shape[-1]))
        Alphabet_Z = np.reshape(Alphabet_Z, (-1, Alphabet_Z.shape[-1]))
        I = []
        for i in range(z.shape[0]):

            def f(X, Y, Alphabet_X, Alphabet_Y):
                return information_mutual_conditional(
                    X,
                    Y,
                    z[i],
                    False,
                    base,
                    fill_value,
                    estimator,
                    Alphabet_X,
                    Alphabet_Y,
                    Alphabet_Z[i],
                )

            I.append(_cartesian_product_apply(x, y, f, alphabet_x, alphabet_y))
        shapeI_XY = I[0].shape
        if len(shapeI_Z) == 0:
            I = np.array(I)[0].reshape(shapeI_XY)
        else:
            I = np.array(I)
            I = np.rollaxis(I, 0, len(I.shape))
            I = I.reshape(np.append(shapeI_XY, shapeI_Z).astype("int"))
        return I

    # Re-shape H, X,Y,Z so that we may handle multi-dimensional arrays
    # equivalently and iterate across 0th axis
    x = np.reshape(x, (-1, x.shape[-1]))
    y = np.reshape(y, (-1, y.shape[-1]))
    z = np.reshape(z, (-1, z.shape[-1]))
    alphabet_x = np.reshape(alphabet_x, (-1, alphabet_x.shape[-1]))
    alphabet_y = np.reshape(alphabet_y, (-1, alphabet_y.shape[-1]))
    Alphabet_Z = np.reshape(Alphabet_Z, (-1, Alphabet_Z.shape[-1]))
    orig_shape_I = I.shape
    I = np.reshape(I, (-1, 1))

    for i in range(x.shape[0]):
        I_ = (
            entropy_joint(
                np.vstack((x[i], z[i])),
                base,
                fill_value,
                estimator,
                _vstack_pad((alphabet_x[i], Alphabet_Z[i]), fill_value),
            )
            + entropy_joint(
                np.vstack((y[i], z[i])),
                base,
                fill_value,
                estimator,
                _vstack_pad((alphabet_y[i], Alphabet_Z[i]), fill_value),
            )
            - entropy_joint(
                np.vstack((x[i], y[i], z[i])),
                base,
                fill_value,
                estimator,
                _vstack_pad((alphabet_x[i], alphabet_y[i], Alphabet_Z[i]), fill_value),
            )
            - entropy_joint(z[i], base, fill_value, estimator, Alphabet_Z[i])
        )
        I[i] = I_

    # Reverse re-shaping
    I = np.reshape(I, orig_shape_I)

    if keep_dims and not cartesian_product:
        I = I[..., np.newaxis]

    return I


def information_mutual(
    X,
    Y=None,
    cartesian_product=False,
    base=2,
    fill_value=-1,
    estimator="ML",
    Alphabet_X=None,
    Alphabet_Y=None,
    keep_dims=False,
):
    H_conditional = entropy_conditional(
        X, Y, cartesian_product, base, fill_value, estimator, Alphabet_X, Alphabet_Y
    )
    H = entropy(X, base, fill_value, estimator, Alphabet_X)

    H = np.reshape(
        H, np.append(H.shape, np.ones(H_conditional.ndim - H.ndim)).astype("int")
    )

    H = H - H_conditional

    if keep_dims and not cartesian_product:
        H = H[..., np.newaxis]

    return H


def entropy_pmf(P, base=2, require_valid_pmf=True, keep_dims=False):
    P, _ = _sanitise_array_input(P)

    if P.size == 0:
        raise ValueError("arg P contains no elements")
    if np.any(_isnan(P)):
        raise ValueError("arg P contains NaN values")
    if np.any(np.logical_or(P < 0, P > 1)):
        raise ValueError("arg P contains values outside unit interval")
    if require_valid_pmf and not np.allclose(np.sum(P, axis=-1), 1):
        raise ValueError("arg P does not sum to unity across last axis")
    if not (np.isscalar(base) and np.isreal(base) and base > 0):
        raise ValueError("arg base not a positive real-valued scalar")

    H = -np.sum(P * np.log2(P + np.spacing(0)), axis=-1)
    H = H / np.log2(base)

    if keep_dims:
        H = H[..., np.newaxis]

    return H


def entropy_joint(
    X, base=2, fill_value=-1, estimator="ML", Alphabet_X=None, keep_dims=False
):
    X, fill_value_X = _sanitise_array_input(X, fill_value)
    if Alphabet_X is not None:
        Alphabet_X, fill_value_Alphabet_X = _sanitise_array_input(
            Alphabet_X, fill_value
        )
        Alphabet_X, _ = _autocreate_alphabet(Alphabet_X, fill_value_Alphabet_X)
    else:
        Alphabet_X, fill_value_Alphabet_X = _autocreate_alphabet(X, fill_value_X)

    if X.size == 0:
        raise ValueError("arg X contains no elements")
    if np.any(_isnan(X)):
        raise ValueError("arg X contains NaN values")
    if Alphabet_X.size == 0:
        raise ValueError("arg Alphabet_X contains no elements")
    if np.any(_isnan(Alphabet_X)):
        raise ValueError("arg Alphabet_X contains NaN values")
    if _isnan(fill_value_X):
        raise ValueError("fill value for arg X is NaN")
    if X.shape[:-1] != Alphabet_X.shape[:-1]:
        raise ValueError("leading dimensions of args X and Alphabet_X do not " "match")
    if not (np.isscalar(base) and np.isreal(base) and base > 0):
        raise ValueError("arg base not a positive real-valued scalar")

    S, fill_value = _map_observations_to_integers(
        (X, Alphabet_X), (fill_value_X, fill_value_Alphabet_X)
    )
    X, Alphabet_X = S

    X = np.reshape(X, (-1, X.shape[-1]))
    Alphabet_X = np.reshape(Alphabet_X, (-1, Alphabet_X.shape[-1]))

    _verify_alphabet_sufficiently_large(X, Alphabet_X, fill_value)

    for i in range(X.shape[0]):
        X = X[:, X[i].argsort(kind="mergesort")]

    B = np.any(X[:, 1:] != X[:, :-1], axis=0)
    I = np.append(np.where(B), X.shape[1] - 1)
    L = np.diff(np.append(-1, I))

    alphabet_X = X[:, I]
    if estimator != "ML":
        n_additional_empty_bins = _determine_number_additional_empty_bins(
            L, alphabet_X, Alphabet_X, fill_value
        )
    else:
        n_additional_empty_bins = 0
    L, _ = _remove_counts_at_fill_value(L, alphabet_X, fill_value)
    if not np.any(L):
        return np.float64(np.NaN)

    # P_0 is the probability mass assigned to each additional empty bin
    P, P_0 = _estimate_probabilities(L, estimator, n_additional_empty_bins)
    H_0 = n_additional_empty_bins * P_0 * -np.log2(P_0 + np.spacing(0)) / np.log2(base)
    H = entropy_pmf(P, base, require_valid_pmf=False) + H_0

    if keep_dims:
        H = H[..., np.newaxis]

    return H


def entropy_conditional(
    X,
    Y=None,
    cartesian_product=False,
    base=2,
    fill_value=-1,
    estimator="ML",
    Alphabet_X=None,
    Alphabet_Y=None,
    keep_dims=False,
):
    if Y is None:
        Y = X
        cartesian_product = True
        Alphabet_Y = Alphabet_X

    X, fill_value_X = _sanitise_array_input(X, fill_value)
    Y, fill_value_Y = _sanitise_array_input(Y, fill_value)
    if Alphabet_X is not None:
        Alphabet_X, fill_value_Alphabet_X = _sanitise_array_input(
            Alphabet_X, fill_value
        )
        Alphabet_X, _ = _autocreate_alphabet(Alphabet_X, fill_value_Alphabet_X)
    else:
        Alphabet_X, fill_value_Alphabet_X = _autocreate_alphabet(X, fill_value_X)
    if Alphabet_Y is not None:
        Alphabet_Y, fill_value_Alphabet_Y = _sanitise_array_input(
            Alphabet_Y, fill_value
        )
        Alphabet_Y, _ = _autocreate_alphabet(Alphabet_Y, fill_value_Alphabet_Y)
    else:
        Alphabet_Y, fill_value_Alphabet_Y = _autocreate_alphabet(Y, fill_value_Y)

    if X.size == 0:
        raise ValueError("arg X contains no elements")
    if Y.size == 0:
        raise ValueError("arg Y contains no elements")
    if np.any(_isnan(X)):
        raise ValueError("arg X contains NaN values")
    if np.any(_isnan(Y)):
        raise ValueError("arg Y contains NaN values")
    if Alphabet_X.size == 0:
        raise ValueError("arg Alphabet_X contains no elements")
    if np.any(_isnan(Alphabet_X)):
        raise ValueError("arg Alphabet_X contains NaN values")
    if Alphabet_Y.size == 0:
        raise ValueError("arg Alphabet_Y contains no elements")
    if np.any(_isnan(Alphabet_Y)):
        raise ValueError("arg Alphabet_Y contains NaN values")
    if _isnan(fill_value_X):
        raise ValueError("fill value for arg X is NaN")
    if _isnan(fill_value_Y):
        raise ValueError("fill value for arg Y is NaN")
    if X.shape[:-1] != Alphabet_X.shape[:-1]:
        raise ValueError("leading dimensions of args X and Alphabet_X do not " "match")
    if Y.shape[:-1] != Alphabet_Y.shape[:-1]:
        raise ValueError("leading dimensions of args Y and Alphabet_Y do not " "match")
    if not cartesian_product and X.shape != Y.shape:
        raise ValueError("dimensions of args X and Y do not match")
    if cartesian_product and X.shape[-1] != Y.shape[-1]:
        raise ValueError("trailing dimensions of args X and Y do not match")
    if not (np.isscalar(base) and np.isreal(base) and base > 0):
        raise ValueError("arg base not a positive real-valued scalar")

    S, fill_value = _map_observations_to_integers(
        (X, Alphabet_X, Y, Alphabet_Y),
        (fill_value_X, fill_value_Alphabet_X, fill_value_Y, fill_value_Alphabet_Y),
    )
    X, Alphabet_X, Y, Alphabet_Y = S

    if not cartesian_product:
        H = np.empty(X.shape[:-1])
        if H.ndim > 0:
            H[:] = np.NaN
        else:
            H = np.float64(np.NaN)
    else:

        def f(X, Y, Alphabet_X, Alphabet_Y):
            return entropy_conditional(
                X, Y, False, base, fill_value, estimator, Alphabet_X, Alphabet_Y
            )

        return _cartesian_product_apply(X, Y, f, Alphabet_X, Alphabet_Y)

    # Re-shape H, X and Y, so that we may handle multi-dimensional arrays
    # equivalently and iterate across 0th axis
    X = np.reshape(X, (-1, X.shape[-1]))
    Y = np.reshape(Y, (-1, Y.shape[-1]))
    Alphabet_X = np.reshape(Alphabet_X, (-1, Alphabet_X.shape[-1]))
    Alphabet_Y = np.reshape(Alphabet_Y, (-1, Alphabet_Y.shape[-1]))
    orig_shape_H = H.shape
    H = np.reshape(H, (-1, 1))

    for i in range(X.shape[0]):
        H[i] = entropy_joint(
            np.vstack((X[i], Y[i])),
            base,
            fill_value,
            estimator,
            _vstack_pad((Alphabet_X[i], Alphabet_Y[i]), fill_value),
        ) - entropy(Y[i], base, fill_value, estimator, Alphabet_Y[i])

    # Reverse re-shaping
    H = np.reshape(H, orig_shape_H)

    if keep_dims and not cartesian_product:
        H = H[..., np.newaxis]

    return H


def entropy(X, base=2, fill_value=-1, estimator="ML", Alphabet_X=None, keep_dims=False):
    X, fill_value_X = _sanitise_array_input(X, fill_value)
    if Alphabet_X is not None:
        Alphabet_X, fill_value_Alphabet_X = _sanitise_array_input(
            Alphabet_X, fill_value
        )
        Alphabet_X, _ = _autocreate_alphabet(Alphabet_X, fill_value_Alphabet_X)
    else:
        Alphabet_X, fill_value_Alphabet_X = _autocreate_alphabet(X, fill_value_X)

    if X.size == 0:
        raise ValueError("arg X contains no elements")
    if np.any(_isnan(X)):
        raise ValueError("arg X contains NaN values")
    if Alphabet_X.size == 0:
        raise ValueError("arg Alphabet_X contains no elements")
    if np.any(_isnan(Alphabet_X)):
        raise ValueError("arg Alphabet_X contains NaN values")
    if _isnan(fill_value_X):
        raise ValueError("fill value for arg X is NaN")
    if X.shape[:-1] != Alphabet_X.shape[:-1]:
        raise ValueError("leading dimensions of args X and Alphabet_X do not " "match")
    if not (np.isscalar(base) and np.isreal(base) and base > 0):
        raise ValueError("arg base not a positive real-valued scalar")

    S, fill_value = _map_observations_to_integers(
        (X, Alphabet_X), (fill_value_X, fill_value_Alphabet_X)
    )
    X, Alphabet_X = S

    H = np.empty(X.shape[:-1])
    if H.ndim > 0:
        H[:] = np.NaN
    else:
        H = np.float64(np.NaN)

    # Re-shape H and X, so that we may handle multi-dimensional arrays
    # equivalently and iterate across 0th axis
    X = np.reshape(X, (-1, X.shape[-1]))
    Alphabet_X = np.reshape(Alphabet_X, (-1, Alphabet_X.shape[-1]))
    orig_shape_H = H.shape
    H = np.reshape(H, (-1, 1))

    _verify_alphabet_sufficiently_large(X, Alphabet_X, fill_value)

    # NB: This is not joint entropy. Elements in each row are sorted
    # independently
    X = np.sort(X, axis=1)

    # Compute symbol run-lengths
    # Compute symbol change indicators
    B = X[:, 1:] != X[:, :-1]
    for i in range(X.shape[0]):
        # Obtain symbol change positions
        I = np.append(np.where(B[i]), X.shape[1] - 1)
        # Compute run lengths
        L = np.diff(np.append(-1, I))

        alphabet_X = X[i, I]
        if estimator != "ML":
            n_additional_empty_bins = _determine_number_additional_empty_bins(
                L, alphabet_X, Alphabet_X[i], fill_value
            )
        else:
            n_additional_empty_bins = 0
        L, _ = _remove_counts_at_fill_value(L, alphabet_X, fill_value)
        if not np.any(L):
            continue

        # P_0 is the probability mass assigned to each additional empty bin
        P, P_0 = _estimate_probabilities(L, estimator, n_additional_empty_bins)
        H_0 = (
            n_additional_empty_bins
            * P_0
            * -np.log2(P_0 + np.spacing(0))
            / np.log2(base)
        )
        H[i] = entropy_pmf(P, base, require_valid_pmf=False) + H_0

    # Reverse re-shaping
    H = np.reshape(H, orig_shape_H)

    if keep_dims:
        H = H[..., np.newaxis]

    return H


def _autocreate_alphabet(X, fill_value):
    Lengths = np.apply_along_axis(lambda x: np.unique(x).size, axis=-1, arr=X)
    max_length = np.max(Lengths)

    def pad_with_fillvalue(x):
        return np.append(x, np.tile(fill_value, int(max_length - x.size)))

    Alphabet = np.apply_along_axis(
        lambda x: pad_with_fillvalue(np.unique(x)), axis=-1, arr=X
    )
    return (Alphabet, fill_value)


def _sanitise_array_input(x, fill_value=-1):
    # Avoid Python 3 issues with numpy arrays containing None elements
    if np.any(np.equal(x, None)) or fill_value is None:
        x = np.array(x, copy=False)
        assert np.all(x != NONE_REPLACEMENT)
        m = np.equal(x, None)
        x = np.where(m, NONE_REPLACEMENT, x)
    if fill_value is None:
        x = np.array(x, copy=False)
        fill_value = NONE_REPLACEMENT

    if isinstance(x, (pd.core.frame.DataFrame, pd.core.series.Series)):
        # Create masked array, honouring Dataframe/Series missing entries
        # NB: We transpose for convenience, so that quantities are computed for
        # each column
        x = np.ma.MaskedArray(x, x.isnull())

    if isinstance(x, np.ma.MaskedArray):
        fill_value = x.fill_value

        if np.any(x == fill_value):
            warnings.warn("Masked array contains data equal to fill value")

        if x.dtype.kind in ("S", "U"):
            kind = x.dtype.kind
            current_dtype_len = int(x.dtype.str.split(kind)[1])
            if current_dtype_len < len(fill_value):
                # Fix numpy's broken array string type behaviour which causes
                # X.filled() placeholder entries to be no longer than
                # non-placeholder entries
                warnings.warn(
                    "Changing numpy array dtype internally to "
                    "accommodate fill_value string length"
                )
                m = x.mask
                x = np.array(x.filled(), dtype=kind + str(len(fill_value)))
                x[m] = fill_value
            else:
                x = x.filled()
        else:
            x = x.filled()
    else:
        x = np.array(x, copy=False)

    if x.dtype.kind not in "biufcmMOSUV":
        raise TypeError("Unsupported array dtype")

    if x.size == 1 and x.ndim == 0:
        x = np.array((x,))

    return x, np.array(fill_value)


def _map_observations_to_integers(Symbol_matrices, Fill_values):
    assert len(Symbol_matrices) == len(Fill_values)
    FILL_VALUE = -1
    if np.any([A.dtype != "int" for A in Symbol_matrices]) or np.any(
        np.array(Fill_values) != FILL_VALUE
    ):
        L = sklearn.preprocessing.LabelEncoder()
        F = [np.atleast_1d(v) for v in Fill_values]
        L.fit(np.concatenate([A.ravel() for A in Symbol_matrices] + F))
        Symbol_matrices = [
            L.transform(A.ravel()).reshape(A.shape) for A in Symbol_matrices
        ]
        Fill_values = [L.transform(np.atleast_1d(f)) for f in Fill_values]

        for A, f in zip(Symbol_matrices, Fill_values):
            assert not np.any(A == FILL_VALUE)
            A[A == f] = FILL_VALUE

    assert np.all([A.dtype == "int" for A in Symbol_matrices])
    return Symbol_matrices, FILL_VALUE


def _isnan(X):
    X = np.array(X, copy=False)
    if X.dtype in ("int", "float"):
        return np.isnan(X)
    else:
        f = np.vectorize(_isnan_element)
        return f(X)


def _isnan_element(x):
    if isinstance(x, type(np.nan)):
        return np.isnan(x)
    else:
        return False


def _determine_number_additional_empty_bins(
    Counts, Alphabet, Full_Alphabet, fill_value
):
    alphabet_sizes = np.sum(np.atleast_2d(Full_Alphabet) != fill_value, axis=-1)
    if np.any(alphabet_sizes != fill_value):
        joint_alphabet_size = np.prod(alphabet_sizes[alphabet_sizes > 0])
        if joint_alphabet_size <= 0:
            raise ValueError(
                "Numerical overflow detected. Joint alphabet " "size too large."
            )
    else:
        joint_alphabet_size = 0
    return joint_alphabet_size - np.sum(
        np.all(np.atleast_2d(Alphabet) != fill_value, axis=0)
    )


def _estimate_probabilities(Counts, estimator, n_additional_empty_bins=0):
    assert np.sum(Counts) > 0
    assert np.all(Counts.astype("int") == Counts)
    assert n_additional_empty_bins >= 0
    Counts = Counts.astype("int")

    if isinstance(estimator, str):
        estimator = estimator.upper().replace(" ", "")

    if np.isreal(estimator) or estimator in ("ML", "PERKS", "MINIMAX"):
        if np.isreal(estimator):
            alpha = estimator
        elif estimator == "PERKS":
            alpha = 1.0 / (Counts.size + n_additional_empty_bins)
        elif estimator == "MINIMAX":
            alpha = np.sqrt(np.sum(Counts)) / (Counts.size + n_additional_empty_bins)
        else:
            alpha = 0
        Theta = (Counts + alpha) / (
            1.0 * np.sum(Counts) + alpha * (Counts.size + n_additional_empty_bins)
        )
        # Theta_0 is the probability mass assigned to each additional empty bin
        if n_additional_empty_bins > 0:
            Theta_0 = alpha / (
                1.0 * np.sum(Counts) + alpha * (Counts.size + n_additional_empty_bins)
            )
        else:
            Theta_0 = 0


def _remove_counts_at_fill_value(Counts, Alphabet, fill_value):
    I = np.any(np.atleast_2d(Alphabet) == fill_value, axis=0)
    if np.any(I):
        Counts = Counts[~I]
        Alphabet = Alphabet.T[~I].T
    return (Counts, Alphabet)


def _cartesian_product_apply(X, Y, function, Alphabet_X=None, Alphabet_Y=None):
    assert X.ndim > 0 and Y.ndim > 0
    assert X.size > 0 and Y.size > 0
    if Alphabet_X is not None or Alphabet_Y is not None:
        assert Alphabet_X.ndim > 0 and Alphabet_Y.ndim > 0
        assert Alphabet_X.size > 0 and Alphabet_Y.size > 0

    H = np.empty(np.append(X.shape[:-1], Y.shape[:-1]).astype("int"))
    if H.ndim > 0:
        H[:] = np.NaN
    else:
        H = np.float64(np.NaN)

    X = np.reshape(X, (-1, X.shape[-1]))
    Y = np.reshape(Y, (-1, Y.shape[-1]))
    if Alphabet_X is not None or Alphabet_Y is not None:
        Alphabet_X = np.reshape(Alphabet_X, (-1, Alphabet_X.shape[-1]))
        Alphabet_Y = np.reshape(Alphabet_Y, (-1, Alphabet_Y.shape[-1]))
    orig_shape_H = H.shape
    H = np.reshape(H, (-1, 1))

    n = 0
    for i in range(X.shape[0]):
        for j in range(Y.shape[0]):
            if Alphabet_X is not None or Alphabet_Y is not None:
                H[n] = function(X[i], Y[j], Alphabet_X[i], Alphabet_Y[j])
            else:
                H[n] = function(X[i], Y[j])
            n = n + 1

    # Reverse re-shaping
    H = np.reshape(H, orig_shape_H)

    return H


def _verify_alphabet_sufficiently_large(X, Alphabet, fill_value):
    assert not np.any(X == np.array(None))
    assert not np.any(Alphabet == np.array(None))
    for i in range(X.shape[0]):
        I = X[i] != fill_value
        J = Alphabet[i] != fill_value
        # NB: This causes issues when both arguments contain None. But it is
        # always called after observations have all been mapped to integers.
        if np.setdiff1d(X[i, I], Alphabet[i, J]).size > 0:
            raise ValueError(
                "provided alphabet does not contain all observed " "outcomes"
            )


def _vstack_pad(arrays, fill_value):
    max_length = max([A.shape[-1] for A in arrays])
    arrays = [
        np.append(
            A,
            np.tile(
                fill_value,
                np.append(A.shape[:-1], max_length - A.shape[-1]).astype(int),
            ),
        )
        for A in arrays
    ]
    return np.vstack(arrays)
