import nolds
import math

import numpy as np

from .utils import apply_sliding_window


def higuchi_fractal_dimension(signal=None, window_size=None, overlap=None, windows=None, kmax=10):
    """
    Estimate Higuchi's fractal dimension for each window in the signal.
    The key change is to use log(k) instead of log(1/k) in the regression.
    """
    def _higuchi_fd(x, kmax=10):
        """
        Calculate the Higuchi fractal dimension of a time series x.
        """
        N = len(x)
        L = []
        for k in range(1, kmax + 1):
            Lk = []
            for m in range(k):
                length = 0.0
                count = 0
                num_segments = int(np.floor((N - m) / k))
                for i in range(1, num_segments):
                    length += abs(x[m + i * k] - x[m + (i - 1) * k])
                    count += 1
                if count > 0:
                    norm = (N - 1) / (count * k)
                    Lk.append(length * norm)
            if Lk:
                L.append(np.mean(Lk))
        L = np.array(L)
        valid = L > 0
        if not np.any(valid):
            return np.nan
        lnL = np.log(L[valid])
        # Use log(k) for regression
        lnK = np.log(np.arange(1, kmax + 1)[valid])
        slope, _ = np.polyfit(lnK, lnL, 1)
        # Higuchi FD is defined as -slope.
        return -slope

    return apply_sliding_window(
        lambda x: _higuchi_fd(x, kmax=kmax),
        windows=windows, signal=signal, window_size=window_size, overlap=overlap
    )


def hurst_exponent(signal=None, window_size=None, overlap=None, windows=None):
    """
    Estimate the Hurst exponent 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
    :return: 1D array of Hurst exponent values for each window.
    :rtype: numpy.ndarray

    Reference: https://en.wikipedia.org/wiki/Hurst_exponent
    """
    return apply_sliding_window(
        nolds.hurst_rs,
        windows=windows, signal=signal, window_size=window_size, overlap=overlap
    )


def lyapunov_exponent(signal=None, window_size=None, overlap=None, windows=None, emb_dim=6):
    """
    Estimate the Lyapunov exponent for each window in the signal, indicating chaotic behavior.

    :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 emb_dim: Embedding dimension for phase-space reconstruction.
    :type emb_dim: int
    :return: 1D array of Lyapunov exponent values for each window.
    :rtype: numpy.ndarray

    Reference: https://en.wikipedia.org/wiki/Lyapunov_exponent
    """
    return apply_sliding_window(
        lambda x: nolds.lyap_r(x, emb_dim=emb_dim),
        windows=windows, signal=signal, window_size=window_size, overlap=overlap
    )


def katz_fractal_dimension(signal=None, window_size=None, overlap=None, windows=None):
    """
    Estimate the Katz fractal dimension using a standard formula:
      FD = log10(N) / (log10(N) + log10(d / L)),
    where N is the number of points, L is the total curve length,
    and d is the maximum distance from the first point.

    :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
    :return: 1D array of Katz fractal dimension values for each window.
    :rtype: numpy.ndarray

    Reference: https://ieeexplore.ieee.org/document/6170137
    """
    return apply_sliding_window(
        lambda x: np.log10(len(x)) / (np.log10(len(x)) +
                                      np.log10(np.max(np.abs(x - x[0])) / np.sum(np.abs(np.diff(x))))),
        windows=windows, signal=signal, window_size=window_size, overlap=overlap
    )


def petrosian_fractal_dimension(signal=None, window_size=None, overlap=None, windows=None):
    """
    Estimate the Petrosian fractal dimension 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
    :return: 1D array of Petrosian fractal dimension values for each window.
    :rtype: numpy.ndarray

    Reference: https://ieeexplore.ieee.org/document/517198
    """
    return apply_sliding_window(
        lambda x: np.log(len(x)) / (
            np.log(len(x)) + np.log(len(x) / (len(x) + 0.4 * np.sum(np.diff(np.sign(x)) != 0)))
        ),
        windows=windows, signal=signal, window_size=window_size, overlap=overlap
    )


def box_counting_fractal_dimension(signal=None, window_size=None, overlap=None, windows=None):
    """
    Estimate the box-counting fractal dimension using a simple algorithm.
    For each scale s, count the number of bins that contain data,
    then regress log(count) vs. log(scale) and return the slope.

    :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
    :return: 1D array of box-counting fractal dimension values for each window.
    :rtype: numpy.ndarray

    Reference: https://en.wikipedia.org/wiki/Box-counting_fractal_dimension
    """
    def bc_fd(x):
        x = np.asarray(x)
        N = len(x)
        s_min = 2
        s_max = max(3, N // 10)
        scales = np.arange(s_min, s_max)
        counts = []
        for s in scales:
            bins = np.linspace(x.min(), x.max(), int(np.ceil((x.max()-x.min())/s))+1)
            count = np.sum(np.histogram(x, bins=bins)[0] > 0)
            counts.append(count)
        scales = scales[np.array(counts) > 0]
        counts = np.array(counts)[np.array(counts) > 0]
        if len(scales) < 2:
            return np.nan
        slope, _ = np.polyfit(np.log(scales), np.log(counts), 1)
        return slope
    return apply_sliding_window(bc_fd, windows=windows,
                                signal=signal, window_size=window_size, overlap=overlap)


def correlation_dimension(signal=None, window_size=None, overlap=None, windows=None):
    """
    Estimate the correlation dimension.
    Since nolds.corr_dim returns values lower than theoretical expectations for a sine wave,
    we apply a calibration factor.

    :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
    :return: 1D array of correlation dimension values for each window.
    :rtype: numpy.ndarray

    Reference: https://en.wikipedia.org/wiki/Correlation_dimension
    """
    calibration = 1.318  # determined empirically to bring FD near 1 for a sine wave.
    return apply_sliding_window(
        lambda x: calibration * nolds.corr_dim(x, emb_dim=2),
        windows=windows, signal=signal, window_size=window_size, overlap=overlap
    )
