import numpy as np
from scipy.signal import csd, coherence, argrelmin
from scipy.fft import fft, fftfreq

def detect_antiresonances_xfft(signal_main, signal_ref=None, fs=1.0,
                               nperseg=None, threshold_db=-6.0,
                               min_freq=0.0, max_freq=None, mode='cross'):
    """
    Detect antiresonance frequencies using cross-spectral (xFFT) analysis.

    Parameters
    ----------
    signal_main : array_like
        The main signal (e.g., output of a system, loss function, J-coupling).
    signal_ref : array_like, optional
        Reference signal (e.g., input excitation, s(t) oscillation). If None
        and mode='cross', the reference is taken as the main signal shifted
        (self‑coherence not recommended). For mode='single', ignored.
    fs : float, optional
        Sampling frequency (Hz). Default 1.0.
    nperseg : int, optional
        Segment length for Welch’s method. Default = min(256, len(signal)//2).
    threshold_db : float, optional
        Minimum dip depth (dB) below local average to qualify as antiresonance.
        Default -6.0 (i.e., at least 6 dB drop).
    min_freq, max_freq : float, optional
        Frequency range to search (Hz). Defaults to Nyquist limit.
    mode : {'cross', 'single'}, optional
        'cross' uses transfer function from reference to main (requires signal_ref).
        'single' uses power spectrum of main signal alone (dips in amplitude).

    Returns
    -------
    results : dict
        - 'frequencies' : ndarray of antiresonance frequencies (Hz)
        - 'depths_db'   : Depth of each dip below local average (dB)
        - 'coherences'  : Coherence at those frequencies (only for 'cross' mode)
        - 'transfer_magnitude' : full frequency array and corresponding |H(f)|
                                 or power spectrum (for 'single' mode)
    """
    # ----- Input validation -----
    signal_main = np.asarray(signal_main)
    n = len(signal_main)
    if nperseg is None:
        nperseg = min(256, n // 2)
    if nperseg < 8:
        nperseg = n // 4 if n > 16 else n // 2

    # ----- Mode selection -----
    if mode == 'cross':
        if signal_ref is None:
            raise ValueError("For mode='cross', signal_ref must be provided.")
        signal_ref = np.asarray(signal_ref)
        if len(signal_ref) != n:
            raise ValueError("Main and reference signals must have same length.")

        # Compute cross‑spectral density and auto‑spectrum
        f, Pxx = csd(signal_ref, signal_ref, fs=fs, nperseg=nperseg)
        _, Pxy = csd(signal_ref, signal_main, fs=fs, nperseg=nperseg)
        # Transfer function magnitude |H(f)|
        H_mag = np.abs(Pxy) / (np.abs(Pxx) + 1e-12)

        # Coherence (optional, for confidence)
        _, Coh = coherence(signal_ref, signal_main, fs=fs, nperseg=nperseg)

    elif mode == 'single':
        # Use power spectrum of the main signal alone
        f = fftfreq(n, d=1/fs)[:n//2]
        spectrum = np.abs(fft(signal_main))[:n//2]
        H_mag = spectrum / (np.max(spectrum) + 1e-12)   # normalised magnitude
        Coh = np.ones_like(f)   # no coherence in single mode

    else:
        raise ValueError("mode must be 'cross' or 'single'.")

    # Limit frequency range
    if max_freq is None:
        max_freq = fs / 2.0
    freq_mask = (f >= min_freq) & (f <= max_freq)
    f = f[freq_mask]
    H_mag = H_mag[freq_mask]
    Coh = Coh[freq_mask]

    # Convert to dB scale (avoid log(0))
    H_db = 20 * np.log10(H_mag + 1e-12)

    # Find local minima (antiresonance candidates)
    local_min_idx = argrelmin(H_db)[0]
    antires_freqs = []
    antires_depths = []
    antires_coherence = []

    for idx in local_min_idx:
        # Compute local average around the dip (excluding the dip point)
        left = max(0, idx - 5)
        right = min(len(H_db) - 1, idx + 5)
        neighbor_vals = np.concatenate([H_db[left:idx], H_db[idx+1:right+1]])
        if len(neighbor_vals) == 0:
            continue
        local_avg = np.mean(neighbor_vals)
        dip_depth = local_avg - H_db[idx]   # positive dB
        if dip_depth >= -threshold_db:      # threshold_db negative, e.g. -6 dB
            antires_freqs.append(f[idx])
            antires_depths.append(dip_depth)
            antires_coherence.append(Coh[idx])

    # Prepare output
    results = {
        'frequencies': np.array(antires_freqs),
        'depths_db': np.array(antires_depths),
        'coherences': np.array(antires_coherence),
        'transfer_magnitude': (f, H_mag),
        'mode': mode
    }
    return results


# ===================== EXAMPLE USAGE =====================
if __name__ == '__main__':
    import matplotlib.pyplot as plt

    # Create synthetic signals with an antiresonance at 0.3 Hz
    fs = 10.0
    t = np.linspace(0, 100, 1000, endpoint=False)
    # Reference: sine sweep or oscillation
    ref = np.sin(2 * np.pi * 0.5 * t) + 0.3 * np.sin(2 * np.pi * 0.9 * t)
    # Main signal: transfer function with a notch at 0.3 Hz (antiresonance)
    from scipy.signal import lfilter, butter
    b_notch, a_notch = butter(4, [0.28, 0.32], btype='bandstop', fs=fs)
    main = lfilter(b_notch, a_notch, ref) + 0.02 * np.random.randn(len(ref))

    # Detect antiresonances
    res = detect_antiresonances_xfft(main, ref, fs=fs, threshold_db=-6.0,
                                     min_freq=0.1, max_freq=1.0, mode='cross')

    print("Detected antiresonances at (Hz):", res['frequencies'])
    print("Depths (dB):", res['depths_db'])
    print("Coherences:", res['coherences'])

    # Plot transfer function
    f, H = res['transfer_magnitude']
    plt.figure(figsize=(10, 4))
    plt.semilogy(f, H, label='|H(f)|')
    for fr in res['frequencies']:
        plt.axvline(fr, color='red', linestyle='--', alpha=0.7,
                    label=f"Antiresonance {fr:.2f} Hz" if fr == res['frequencies'][0] else "")
    plt.xlabel('Frequency (Hz)')
    plt.ylabel('|H(f)| (linear)')
    plt.title('Transfer function with detected antiresonances')
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.show()