# import re
import matplotlib.pyplot as plt
import nmrglue as ng
import numpy as np

from ..libs.logger import get_logger

logger = get_logger(__name__)


def interleavefid(fid):
    """
    Interleave the real and imaginary parts of a 1D or 2D FID signal.

    Args:
        fid (numpy.ndarray): A numpy array representing the FID signal. Can be 1D or 2D.

    Returns:
        numpy.ndarray: A numpy array with interleaved real and imaginary parts of the FID signal.
    """
    if len(fid.shape) == 1:
        newfid = np.zeros(len(fid) * 2, dtype=float)
        newfid[::2] = fid.real
        newfid[1::2] = fid.imag

    elif len(fid.shape) == 2:
        row, col = fid.shape
        newfid = np.zeros((row * 2, col))
        newfid[::2, :] = fid.real
        newfid[1::2, :] = fid.imag

    else:
        raise TypeError(
            f"This function requires a 1D or 2D array, but received an array with shape {fid.shape}."
        )

    return np.squeeze(newfid)


def uninterleave(fid):
    """
    Reconstruct a complex FID signal from an interleaved array.

    Args:
        fid (numpy.ndarray): A numpy array with interleaved real and imaginary parts
            of the FID signal. Can be 1D or 2D.

    Returns:
        numpy.ndarray: A numpy array representing the complex FID signal.
    """
    if len(fid.shape) == 1:
        fid2 = fid[::2] + fid[1::2] * 1j
    elif len(fid.shape) == 2:
        fid2 = fid[::2, :] + fid[1::2, :] * 1j
    else:
        raise TypeError("This function requires 1D or 2D array!" % fid.shape)
    return fid2


def equation6(x, ak=75, fk=0, dk=50, phi=0, g=0):
    """
    Generate a FID signal according to equation 6 of Vanhamme, L. et al,
    J Magn Reson 1997, 129 (1), 35-43.

    Args:
        x (1D array): The time axis for which the FID signal is generated.
        ak (float): Amplitude.
        fk (float): Frequency.
        dk (float): Damping factor.
        phi (float): Phase.
        g (float): Lineshape parameter. g=0 for Lorentzian; g=1 for Gaussian.

    Returns:
        1D array: An interleaved array containing the real and imaginary parts of the generated FID signal.
    """
    fid = (
        ak
        * np.exp(1j * phi)
        * np.exp((-dk * (1 - g + g * x) * x))
        * np.exp((1j * 2 * np.pi * fk * x))
    )
    return interleavefid(fid)


def multieq6(params, x, fid=None, return_mat=False):
    """
    Combine multiple FID signals according to equation 6 of Vanhamme, L. et al, J Magn Reson 1997,
    129 (1), 35-43.

    This function iterates over subsets of parameters in ``params``, each subset used to generate
    an FID signal via ``equation6``. The function can return either a matrix of individual FID signals
    or their sum. An existing FID signal can be subtracted from the sum if provided.

    Args:
        params (dict): Dictionary or similar structure with FID signal parameters.
        x (1D array): The time axis for which the FID signal is generated.
        fid (1D array, optional): Existing FID signal to subtract from the sum of generated FID signals.
        return_mat (bool): If True, returns a matrix of individual FID signals; otherwise, returns their sum.

    Returns:
        1D or 2D array: Array of individual FID signals or their sum, optionally with ``fid`` subtracted.
    """
    xInit = np.array(list(params.valuesdict().values()))
    paranum = 5
    assert np.mod(len(xInit), paranum) == 0  # assert 5 parameters

    fidmat = []
    for counter in range(
        len(xInit) // paranum
    ):  # Each peak must have identical number of parameters
        args = xInit[counter * paranum : (counter + 1) * paranum]
        equation_result = equation6(x, *args)

        if return_mat:
            fidmat.append(uninterleave(equation_result))
        else:
            fidmat.append(equation_result)

    fids = np.array(fidmat)

    if return_mat:
        return fids
    else:
        fittedfid = np.sum(fids, axis=0)
        if fid is not None:
            return fittedfid - interleavefid(fid)
        return fittedfid


def Jac6(params, x, fid=None):
    """
    Calculate the Jacobian matrix for the FID signals generated by ``multieq6``.

    This function computes the partial derivatives of FID signals with respect to the
    parameters used to generate them, which is useful for optimization and analysis.

    Args:
        params (dict): Dictionary or similar structure with FID signal parameters.
        x (1D array): The time axis for which the FID signal is generated.
        fid (1D array): Not used in this function but included for interface consistency.

    Returns:
        numpy.ndarray: The interleaved Jacobian matrix of the FID signals with respect to the parameters.
    """
    poptall = list(params.valuesdict().values())
    ak = np.array(poptall[::5])
    dk = np.array(poptall[2::5])
    g = np.array(poptall[4::5])

    # if len(g[g > 1]) > 0:
    #     print("Warning, g>1", g)
    # if len(g[g < 0]):
    #     print("warning! g<0", g)
    g[g > 1] = 1.0
    g[g < 0] = 0.0

    inputfids = multieq6(params, x, return_mat=True).T
    assert len(inputfids.shape) == 2
    pts, npeak = inputfids.shape
    x2 = np.tile(x, (npeak, 1)).T

    d_freq = inputfids * 2j * np.pi * x2
    d_damp = -1 * inputfids * x2 * (g * x2 - g + 1)
    d_amp = inputfids * 1 / ak
    d_ph = inputfids * 1j
    d_g = -1 * dk * (x2 - 1) * x2 * inputfids

    jacobian = np.zeros((pts, npeak * 5), dtype=complex)

    jacobian[:, 0::5] = d_amp
    jacobian[:, 1::5] = d_freq
    jacobian[:, 2::5] = d_damp
    jacobian[:, 3::5] = d_ph
    jacobian[:, 4::5] = d_g

    return interleavefid(jacobian)


def Jac6c(params, x, fid=None):
    """
    Calculate the Jacobian matrix for the FID signals generated by ``multieq6``, excluding `g`.

    This function is a variant of ``Jac6`` and computes the partial derivatives of FID signals
    with respect to the parameters used to generate them, excluding the lineshape parameter ``g``.

    Args:
        params (dict): Dictionary or similar structure with FID signal parameters.
        x (1D array): The time axis for which the FID signal is generated.
        fid (1D array): Not used in this function but included for interface consistency.

    Returns:
        numpy.ndarray: The interleaved Jacobian matrix of the FID signals with respect to the parameters,
                       excluding the lineshape parameter ``g``.
    """
    poptall = list(params.valuesdict().values())
    ak = np.array(poptall[::5])
    dk = np.array(poptall[2::5])  # noqa F841
    g = np.array(poptall[4::5])

    # if len(g[g > 1]) > 0:
    #     print("Warning, g>1", g)
    # if len(g[g < 0]):
    #     print("warning! g<0", g)
    g[g > 1] = 1.0
    g[g < 0] = 0.0

    inputfids = multieq6(params, x, return_mat=True).T
    assert len(inputfids.shape) == 2
    pts, npeak = inputfids.shape
    x2 = np.tile(x, (npeak, 1)).T

    d_freq = inputfids * 2j * np.pi * x2
    d_damp = -1 * inputfids * x2 * (g * x2 - g + 1)
    d_amp = inputfids * 1 / ak
    d_ph = inputfids * 1j

    jacobian = np.zeros((pts, npeak * 4), dtype=complex)

    jacobian[:, 0::4] = d_amp
    jacobian[:, 1::4] = d_freq
    jacobian[:, 2::4] = d_damp
    jacobian[:, 3::4] = d_ph
    # jacobian[:, 4::5] = d_g
    return interleavefid(jacobian)


def fft_params(timeaxis, params, fid=False, return_mat=False):
    """
    Perform Fast Fourier Transform (FFT) on the FID signals generated by ``multieq6``.

    This function optionally returns the raw FID signals or their FFT spectra depending on
    the provided arguments. It is designed to work with lmfit-style parameter objects.

    Args:
        timeaxis (1D array): The time axis for which the FID signals are generated.
        params (lmfit.Parameters): Object containing the FID signal parameters.
        fid (bool): If True, returns the uninterleaved FID signals without applying FFT.
        return_mat (bool): If True, returns the raw FID signals as generated by ``multieq6``.

    Returns:
        numpy.ndarray: Depending on the arguments, either raw FID signals, uninterleaved FID signals, or the FFT spectra of the FID signals.
    """
    # Return raw FID signals if return_mat is True
    # parmas is the lmfit Parameters() style
    if return_mat:
        return multieq6(params, timeaxis, return_mat=return_mat)
    if fid:
        return uninterleave(multieq6(params, timeaxis, return_mat=return_mat))
    # spec = np.fft.fftshift(np.fft.fft((uninterleave(multieq6(params, timeaxis)))))
    spec = ng.proc_base.fft((uninterleave(multieq6(params, timeaxis))))
    return spec


def remove_zero_padding(fid, threshold=1e-10, window_size=5):
    """
    Detect and remove the zero-filling-like constant tail from an FID signal using derivative analysis.

    This function analyzes an FID signal using derivative analysis to identify and remove
    constant-value tail sections. These sections may have zero or near-zero values but
    share the characteristic of being constant.

    Args:
        fid (numpy.ndarray): Input FID signal.
        threshold (float): Threshold value for detecting constant regions using signal
            differentiation. Default: 1e-10
        window_size (int): Number of consecutive points needed to confirm a constant
            region. Default: 5

    Returns:
        int: Index where the signal was truncated.
    """
    differences = np.diff(fid)
    constant_start = np.where(np.abs(differences) < threshold)[0]
    cutoff_idx = len(fid)

    for i in range(len(constant_start) - window_size):
        if np.all(np.diff(constant_start[i : i + window_size]) == 1):
            cutoff_idx = constant_start[i]
            break

    return cutoff_idx
    # return fid[:cutoff_idx], cutoff_idx


def process_fid(fid, deadtime=0.0, sw=10000, lb=5.0, ifphase=False, ifplot=False):
    """
    Processes an FID signal, applies exponential linbroadening,
    performs FFT, and optionally phases and plots the spectrum.

    Args:
        fid (ndarray): The FID signal array.
        deadtime (float, optional): FID dead time in seconds. Defaults to 0.0.
        sw (int, optional): Spectral width in Hz. Defaults to 10000.
        lb (float, optional): Line broadening factor in Hz. Defaults to 5.0
        ifphase (bool, optional): If True, apply phase correction. Defaults to False.
        ifplot (bool, optional): If True, plot the spectrum. Defaults to False.

    Returns:
        1D numpy array: The processed spectrum.
    """
    fidpt = len(fid)
    dwelltime = 1 / sw  # noqa F841  #place holder
    # timeaxis = np.arange(0, dwelltime * fidpt, dwelltime) + deadtime
    Hz = np.linspace(-sw / 2, sw / 2, fidpt)
    fid2 = ng.proc_base.em(fid, lb / sw)  # Exponential apodization
    spec = ng.proc_base.fft(fid2)

    if ifphase:
        p0 = np.mean(np.angle(fid2[0:10]))
        p1 = np.exp(-1j * (p0 + 2 * np.pi * Hz * deadtime))
        spec = spec * p1  # Phasing the spectrum

    if ifplot:
        import matplotlib.pyplot as plt

        plt.plot(Hz, spec.real, "b-")  # Plotting the real part of the spectrum
        plt.xlabel("Hz")
        plt.show()

    return spec


def Compare_to_OXSA(inputfid, resultfid):
    dataNormSq = np.linalg.norm(inputfid - np.mean(inputfid)) ** 2
    resNormSq = np.sum(np.abs((resultfid - inputfid)) ** 2)
    relativeNorm = resNormSq / dataNormSq
    logger.info("Norm of residual = %3.3f" % resNormSq)
    logger.info("Norm of the data = %3.3f" % dataNormSq)
    logger.info("resNormSq / dataNormSq = %3.3f" % relativeNorm)
    return resNormSq, relativeNorm


def fidSNR(fid, indsignal=(0, 10), pts_noise=200):
    """
    Calculates the signal-to-noise ratio (SNR) of a FID signal.

    Here, the SNR is measured by dividing the mean signal amplitude within a
    specified region (``indsignal``) by the standard deviation of the noise, determined
    from the last points (``pts_noise``) of the FID signal.


    Args:
        fid (numpy.ndarray): The FID from which the SNR is to be calculated.
        indsignal (tuple, optional): A tuple including start and end indices
        (s1, s2) of the signal region used to calculate the mean signal amplitude.
        Defaults to (0, 10). Sometimes the first bad points may be excluded from
        the SNR calculation.
        pts_noise (int, optional): The number of noisy points in an FID. There
        should be no signal in this region.

    Returns:
        float: The calculated signal-to-noise ratio (SNR) of the provided FID signal.
    """
    s1, s2 = indsignal
    signal_p = np.mean(np.abs(fid[s1:s2]))
    noise_p = np.std(fid[-pts_noise:])
    return signal_p / noise_p


def add_noise_FID(purefid, snr_target, indsignal=(0, 10), pts_noise=200):
    """
    Adds Gaussian white noise to a pure FID signal to achieve a target SNR

    Args:
        purefid (numpy.ndarray): simulated pure FID signal
        snr_target (float): The target SNR for noisy FID
        indsignal (tuple, optional) and pts_noise (int, optional): See ``fidSNR``

    Returns:
        numpy.ndarray: The FID signal array with added Gaussian white noise, aiming for the specified target SNR.
    """
    s1, s2 = indsignal
    signal_p = np.mean(np.abs(purefid[s1:s2]))
    noise_std = signal_p / snr_target
    noise = np.random.normal(0, noise_std, purefid.shape)
    noisy_fid = purefid + noise
    return noisy_fid


def simulate_fid(
    params,
    MHz=120.0,
    sw=10000.0,
    deadtime=0.0,
    fid_len=1024,
    snr_target=50,
    indsignal=(0, 10),
    pts_noise=200,
    preview=False,
):
    """
    A function that simulates an FID from an lmfit Parameter object and optionally adds Gaussian white noise to achieve a target SNR.

    Args:
        params (lmfit.Parameters): Object containing the FID signal parameters.
        MHz (float): The field strength in MHz.
        sw (float): The spectral width in Hz.
        deadtime (float): The dead time or begin time in seconds before the FID signal starts.
        fid_len (int, optional): The number of points of FID. Defaults to 1024.
        snr_target (float, optional): The target SNR for the simulated FID signal.
          If None, a simulated FID without noise is returned. Defaults to 50.
        indsignal (tuple, optional): A tuple including start and end indices
          (s1, s2) of the signal region used to calculate the mean signal amplitude.
          Defaults to (0, 10). Sometimes the first bad points may be excluded from
          the SNR calculation.
        pts_noise (int, optional): The number of noisy points in an FID. There
          should be no signal in this region.
        preview (bool, optional): If True, plots the FTed simulated FID signal.
          Defaults to False.

    Returns:
        numpy.ndarray: The simulated FID signal, optionally with added noise to achieve the target SNR.
    """

    sw = float(sw)
    MHz = float(MHz)
    deadtime = float(deadtime)
    dwelltime = 1.0 / sw  # noqa F841  #place holder
    timeaxis = np.arange(0, dwelltime * fid_len, dwelltime) + deadtime  # timeaxis
    fidsim = uninterleave(multieq6(x=timeaxis, params=params))
    if snr_target is not None:
        fidsim = add_noise_FID(fidsim, snr_target, indsignal, pts_noise)
    if preview:
        Hz = np.linspace(-sw / 2, sw / 2, fid_len)
        if snr_target is None:
            label = "Pure FID"
            plt.title("Simulated FID")
        else:
            label = "SNR=%2.2f" % fidSNR(
                fid=fidsim, indsignal=indsignal, pts_noise=pts_noise
            )
            plt.title("Simulated FID with an SNR of %2.2f" % snr_target)
        plt.plot(Hz, np.real(ng.proc_base.fft(fidsim)), label=label)
        plt.legend()
        plt.xlabel("Hz")
        plt.show()
    return fidsim
