import numpy as np
import matplotlib.pyplot as plt
from scipy.fft import fft, fftfreq

# ----------------------------------------------------------------------
# 1. Geometric mean from a signal's Fourier components
# ----------------------------------------------------------------------
def geometric_mean_from_signal(signal, sample_rate, num_components=None, eps=1e-12):
    """
    Decompose signal into sinusoids via FFT, then recombine using geometric mean.
    
    Parameters:
        signal : np.ndarray
            Input 1D signal (any length).
        sample_rate : float
            Sampling rate (Hz).
        num_components : int, optional
            Number of strongest frequency components to use.
            If None, use all positive frequencies (except DC and Nyquist).
        eps : float
            Small constant for numerical stability.
    
    Returns:
        t : np.ndarray
            Time axis (same length as signal).
        geom_signal : np.ndarray
            Geometric‑mean re‑synthesised signal.
        freqs_used : np.ndarray
            Frequencies used (Hz).
        amps_used : np.ndarray
            Amplitudes used.
        phases_used : np.ndarray
            Phases used (radians).
    """
    N = len(signal)
    t = np.arange(N) / sample_rate
    
    # Compute FFT
    yf = fft(signal)
    # Positive frequencies only (0 to Nyquist)
    n_pos = N // 2 + 1 if N % 2 == 0 else (N + 1) // 2
    freqs = fftfreq(N, 1/sample_rate)[:n_pos]
    amps = np.abs(yf[:n_pos]) * 2 / N   # scale to original amplitude
    amps[0] /= 2  # DC component correction
    phases = np.angle(yf[:n_pos])
    
    # Exclude DC (freq=0) because cos(0) is constant – geometric mean would be trivial
    mask = freqs > 0
    freqs_pos = freqs[mask]
    amps_pos = amps[mask]
    phases_pos = phases[mask]
    
    # Sort by amplitude descending and optionally limit number of components
    idx_sorted = np.argsort(amps_pos)[::-1]
    if num_components is not None:
        idx_sorted = idx_sorted[:num_components]
    freqs_used = freqs_pos[idx_sorted]
    amps_used = amps_pos[idx_sorted]
    phases_used = phases_pos[idx_sorted]
    
    print(f"Using {len(freqs_used)} frequency components (strongest).")
    
    # Build matrix of sinusoids: rows = time, columns = components
    sinusoids = np.zeros((N, len(freqs_used)))
    for i, (f, A, phi) in enumerate(zip(freqs_used, amps_used, phases_used)):
        sinusoids[:, i] = A * np.cos(2 * np.pi * f * t + phi)
    
    # Geometric mean across components (axis=1)
    abs_prod = np.prod(np.abs(sinusoids) + eps, axis=1)
    sign_prod = np.sign(np.prod(sinusoids, axis=1))
    geom_signal = sign_prod * (abs_prod ** (1.0 / len(freqs_used)))
    
    return t, geom_signal, freqs_used, amps_used, phases_used

# ----------------------------------------------------------------------
# 2. Helper: plot original vs geometric mean signal and spectra
# ----------------------------------------------------------------------
def compare_original_and_geom(signal, t, geom_signal, sample_rate, title="Geometric Mean Reconstruction"):
    fig, axes = plt.subplots(2, 2, figsize=(12, 6))
    
    # Time domain: original
    axes[0,0].plot(t, signal)
    axes[0,0].set_title("Original Signal (Time)")
    axes[0,0].set_xlabel("Time (s)")
    axes[0,0].set_ylabel("Amplitude")
    axes[0,0].grid(True)
    
    # Time domain: geometric mean
    axes[0,1].plot(t, geom_signal)
    axes[0,1].set_title(f"{title} (Time)")
    axes[0,1].set_xlabel("Time (s)")
    axes[0,1].set_ylabel("Amplitude")
    axes[0,1].grid(True)
    
    # Frequency spectra
    n = len(signal)
    yf_orig = fft(signal)
    yf_geom = fft(geom_signal)
    xf = fftfreq(n, t[1]-t[0])[:n//2]
    
    axes[1,0].plot(xf, 2.0/n * np.abs(yf_orig[:n//2]))
    axes[1,0].set_title("Original Spectrum")
    axes[1,0].set_xlabel("Frequency (Hz)")
    axes[1,0].set_ylabel("Magnitude")
    axes[1,0].grid(True)
    
    axes[1,1].plot(xf, 2.0/n * np.abs(yf_geom[:n//2]))
    axes[1,1].set_title(f"{title} Spectrum")
    axes[1,1].set_xlabel("Frequency (Hz)")
    axes[1,1].set_ylabel("Magnitude")
    axes[1,1].grid(True)
    
    plt.tight_layout()
    plt.show()

# ----------------------------------------------------------------------
# 3. Examples using custom signals
# ----------------------------------------------------------------------
if __name__ == "__main__":
    sample_rate = 1000  # Hz
    duration = 2.0      # seconds
    t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
    
    # --- Example A: Synthetic signal (sum of two sinusoids) ---
    print("Example A: Input = sum of 5Hz and 8Hz sinusoids")
    signal_A = np.cos(2 * np.pi * 5 * t) + 0.7 * np.cos(2 * np.pi * 8 * t)
    tA, geomA, freqsA, ampsA, phasesA = geometric_mean_from_signal(signal_A, sample_rate, num_components=4)
    compare_original_and_geom(signal_A, tA, geomA, sample_rate, "Geometric Mean from 5+8Hz")
    
    # --- Example B: Realistic – chirp signal ---
    print("\nExample B: Input = linear chirp from 2Hz to 10Hz")
    signal_B = np.cos(2 * np.pi * (2 + 4*t) * t)  # frequency increases linearly
    tB, geomB, _, _, _ = geometric_mean_from_signal(signal_B, sample_rate, num_components=10)
    compare_original_and_geom(signal_B, tB, geomB, sample_rate, "Geometric Mean from Chirp")
    
    # --- Example C: Noisy signal ---
    print("\nExample C: Input = 3Hz sine + Gaussian noise")
    signal_C = np.cos(2 * np.pi * 3 * t) + 0.5 * np.random.randn(len(t))
    tC, geomC, _, _, _ = geometric_mean_from_signal(signal_C, sample_rate, num_components=6)
    compare_original_and_geom(signal_C, tC, geomC, sample_rate, "Geometric Mean from Noisy 3Hz")
    
    # --- Example D: Your own signal – load from file or generate ---
    # Uncomment and modify to use your own data:
    # import scipy.io.wavfile as wav
    # fs, data = wav.read('your_audio.wav')
    # if data.ndim > 1: data = data[:,0]  # mono
    # t_own, geom_own, _, _, _ = geometric_mean_from_signal(data, fs, num_components=50)
    # compare_original_and_geom(data, t_own, geom_own, fs, "Geometric Mean from Your Signal")