import numpy as np

from .utils import apply_sliding_window


def mutual_information(signal=None, window_size=None, overlap=None, windows=None, bins=50):
    """
    Calculate the mutual information between consecutive samples for each window in the signal.

    :param signal: The input signal as a 1D array or list (if windows not provided).
    :type signal: array-like
    :param window_size: The number of samples in each window (if windows not provided).
    :type window_size: int
    :param overlap: The number of samples to overlap between consecutive windows (if windows not provided).
    :type overlap: int
    :param windows: Optional precomputed list of windows.
    :type windows: list
    :param bins: Number of bins for the 2D histogram.
    :type bins: int
    :return: 1D array of mutual information values for each window.
    :rtype: numpy.ndarray

    Reference: https://en.wikipedia.org/wiki/Mutual_information
    """

    def mi_func(x):
        # Ensure at least 2 samples for a 2D histogram
        if len(x) < 2:
            return 0.0
        hist2d, _, _ = np.histogram2d(x[:-1], x[1:], bins=bins)
        hist2d += 1e-12  # Avoid zeros in the histogram
        pxy = hist2d / hist2d.sum()
        px = pxy.sum(axis=1)
        py = pxy.sum(axis=0)
        # Calculate mutual information, avoiding log of zero by construction
        return np.sum(pxy * np.log2(pxy / (px[:, None] * py[None, :])))

    return apply_sliding_window(mi_func, windows=windows, signal=signal, window_size=window_size, overlap=overlap)


def symbolic_dynamics_entropy(signal=None, window_size=None, overlap=None, windows=None, n_symbols=2):
    """
    Calculate symbolic dynamics entropy for each window in the signal by partitioning
    the window's values into discrete symbols and computing the Shannon entropy of
    the resulting symbol distribution.

    :param signal: The input signal as a 1D array or list (if windows not provided).
    :type signal: array-like
    :param window_size: The number of samples in each window (if windows not provided).
    :type window_size: int
    :param overlap: The number of samples to overlap between consecutive windows (if windows not provided).
    :type overlap: int
    :param windows: Optional precomputed list of windows.
    :type windows: list
    :param n_symbols: The number of discrete symbols to partition the signal's range.
                      Default is 2.
    :type n_symbols: int
    :return: 1D array of symbolic dynamics entropy values for each window.
    :rtype: numpy.ndarray

    Reference: https://en.wikipedia.org/wiki/Symbolic_dynamics
    """

    def sym_dyn_entropy(x):
        # Partition the data into n_symbols bins.
        hist, _ = np.histogram(x, bins=n_symbols)
        # Normalize counts to obtain a probability distribution.
        probs = hist / np.sum(hist)
        # Filter out zero probabilities to avoid log2(0).
        probs = probs[probs > 0]
        # Compute the Shannon entropy.
        return -np.sum(probs * np.log2(probs))

    return apply_sliding_window(sym_dyn_entropy, windows=windows, signal=signal, window_size=window_size,
                                overlap=overlap)


def conditional_entropy(signal=None, window_size=None, overlap=None, windows=None, m=2):
    """
    Calculate the conditional entropy for each window.
    Here we define conditional entropy as the negative average log-probability
    (so that it is nonnegative).

    :param signal: The input signal as a 1D array or list (if windows not provided).
    :type signal: array-like
    :param window_size: The number of samples in each window (if windows not provided).
    :type window_size: int
    :param overlap: The number of samples to overlap between consecutive windows (if windows not provided).
    :type overlap: int
    :param windows: Optional precomputed list of windows.
    :type windows: list
    :param m: Embedding dimension.
    :type m: int
    :return: 1D array of conditional entropy values for each window.
    :rtype: numpy.ndarray

    Reference: https://en.wikipedia.org/wiki/Conditional_entropy
    """

    def conditional_entropy_func(x):
        N = len(x)
        if N <= m:
            return 0.0
        x_embedded = np.array([x[i:i + m] for i in range(N - m + 1)])
        std_x = np.std(x)
        if std_x == 0:
            return 0.0
        C = np.sum(
            np.abs(x_embedded[:, None] - x_embedded[None, :]).max(axis=2)
            <= 0.2 * std_x,
            axis=0
        ) / (N - m + 1)
        # Return the negative average log to ensure nonnegativity.
        return -np.sum(np.log(C)) / (N - m + 1)
    return apply_sliding_window(conditional_entropy_func, windows=windows,
                                signal=signal, window_size=window_size, overlap=overlap)
