import os
import numpy as np
import scipy.signal as signal
import scipy.cluster.vq as vq
from scipy.io import wavfile

# --- CONFIGURATION ---
FS = 16000
FRAME_MS = 25
FRAME_LEN = int(FS * FRAME_MS / 1000)
ORDER = 18
CODEBOOK_SIZE = 64
AUDIO_FILE = "gold_standard_voice.wav"
OUTPUT_FILE = "auto_test_output.wav"

def v_uv_discriminator(frame):
    """V5 Time-Lag Autocorrelation Pitch Tracker."""
    energy = np.mean(frame**2)
    if energy < 1e-7: 
        return False, 0.0, energy
    
    b, a = signal.butter(2, 600 / (FS / 2), btype='low')
    lpf_frame = signal.lfilter(b, a, frame)
    corr = np.correlate(lpf_frame, lpf_frame, mode='full')
    corr = corr[len(corr)//2:]
    corr = corr / (corr[0] + 1e-9)
    
    min_lag = int(FS / 400)
    max_lag = int(FS / 50)
    
    if max_lag >= len(corr): 
        return False, 0.0, energy
        
    valid_corr = corr[min_lag:max_lag]
    peak_idx = np.argmax(valid_corr)
    true_lag = min_lag + peak_idx
    
    if corr[true_lag] > 0.45:
        return True, FS / true_lag, energy
    return False, 0.0, energy

def extract_poles_raw(frame, order=ORDER):
    """
    Extracts poles using Universal Autocorrelation and Tikhonov Regularization.
    This solves the 18th-order Matrix Ill-Conditioning problem.
    """
    windowed = frame * np.hamming(len(frame))
    corr = np.correlate(windowed, windowed, mode='full')
    proc_signal = corr[len(corr)//2 : len(corr)//2 + FRAME_LEN//2]
    
    # === NORMALIZATION FIX ===
    # Scale the autocorrelation to prevent Tikhonov from crushing quiet frames
    proc_signal = proc_signal / (np.max(np.abs(proc_signal)) + 1e-9)
    
    n = len(proc_signal)
    if n <= order:
        return np.array([1.0] + [0.0] * order)

    y = proc_signal[order:]
    X = np.zeros((n - order, order))
    for i in range(order):
        X[:, i] = proc_signal[order - i - 1 : n - i - 1]
    
    try:
        # === THE FIX: TIKHONOV REGULARIZATION (Ridge Regression) ===
        # Prevents matrix singularities and explosions at high orders
        R = X.T @ X + 0.001 * np.eye(order)
        p = X.T @ -y
        a_coeffs = np.linalg.inv(R) @ p
        
        a = np.concatenate(([1.0], a_coeffs))
        
        # Stability check and Surgical Cellophane Crusher
        roots = np.roots(a)
        for i in range(len(roots)):
            r = roots[i]
            freq_hz = (np.abs(np.angle(r)) / (2 * np.pi)) * FS
            radius = np.abs(r)
            if radius > 0.99:
                r = r / (radius + 1e-6) * 0.99
            if freq_hz > 3500:
                r = r * 0.70  
            else:
                r = r * 0.98
            roots[i] = r
            
        a = np.poly(roots).real
                
        return a
    except np.linalg.LinAlgError:
        return np.array([1.0] + [0.0] * order)

def synthesize_frame(f0, energy, poles, is_voiced, zi, phase_offset):
    """
    Synthesizes the frame with Mixed Excitation, Continuous Phase, 
    and Post-Filter Gain Scaling.
    """
    if is_voiced and f0 > 50:
        period = int(FS / f0)
        source = np.zeros(FRAME_LEN)
        indices = np.arange(phase_offset, FRAME_LEN, period)
        if len(indices) > 0:
            source[indices] = 1.0
            phase_offset = indices[-1] + period - FRAME_LEN
        else:
            phase_offset = max(0, phase_offset - FRAME_LEN)
        
        b_glot, a_glot = signal.butter(2, 3000 / (FS / 2), btype='low')
        source = signal.lfilter(b_glot, a_glot, source)
        source += np.random.normal(0, 1, FRAME_LEN) * 0.15
    else:
        source = np.random.normal(0, 1, FRAME_LEN)
        phase_offset = 0

    if len(zi) != len(poles) - 1:
        zi = np.zeros(len(poles) - 1)
    
    synth, zf = signal.lfilter([1.0], poles, source, zi=zi)
    
    synth_energy = np.mean(synth**2) + 1e-9
    scaling_factor = np.sqrt(energy / synth_energy)
    if not is_voiced:
        scaling_factor *= 0.5 
        
    synth = synth * scaling_factor
    return synth, zf, phase_offset

def calc_lsd(target_frame, synth_frame):
    """Calculates the Log-Spectral Distance (LSD) in dB."""
    win = np.hamming(FRAME_LEN)
    H_target = np.fft.rfft(target_frame * win)
    H_synth = np.fft.rfft(synth_frame * win)
    
    # Use a floor of 1e-3 (-60dB) to prevent deep phase-nulls from causing mathematically explosive distances
    log_target = 20 * np.log10(np.maximum(np.abs(H_target), 1e-3))
    log_synth = 20 * np.log10(np.maximum(np.abs(H_synth), 1e-3))
    
    return np.sqrt(np.mean((log_target - log_synth)**2))

def main():
    print("--- Blackout Comms: Automated VQ Benchmark ---")
    
    if not os.path.exists(AUDIO_FILE):
        print(f"Error: '{AUDIO_FILE}' not found. Please run a Codebook Builder script first to record audio.")
        return
        
    sr, audio = wavfile.read(AUDIO_FILE)
    if audio.dtype != np.float32:
        audio = audio.astype(np.float32)
        if np.max(np.abs(audio)) > 1.0:
            audio /= 32768.0
            
    print(f"1. Loaded {len(audio)/FS:.2f} seconds of Gold Standard Audio.")
    audio = np.append(audio[0], audio[1:] - 0.97 * audio[:-1])

    # === PIPELINE STAGE 1: CODEBOOK GENERATION ===
    print("2. Extracting Regularized Poles for Codebook...")
    all_poles = []
    for i in range(0, len(audio) - FRAME_LEN, FRAME_LEN):
        frame = audio[i : i + FRAME_LEN]
        if np.mean(frame**2) < 1e-7: continue
        a = extract_poles_raw(frame)
        all_poles.append(a)
        
    print(f"   Generating K-Medoids Spectrogram Clusters (Size: {CODEBOOK_SIZE})...")
    freq_responses = []
    for a in all_poles:
        _, h = signal.freqz([1.0], a, worN=128)
        freq_responses.append(np.log10(np.abs(h) + 1e-9))
        
    data = np.array(freq_responses)
    # Use k-means++ initialization to prevent empty clusters and guarantee 64 distinct shapes
    centroids, labels = vq.kmeans2(data, CODEBOOK_SIZE, minit='++')
    
    stable_medoids = []
    for i in range(CODEBOOK_SIZE):
        cluster_indices = np.where(labels == i)[0]
        if len(cluster_indices) == 0: continue
        best_idx = cluster_indices[0]
        min_dist = float('inf')
        for idx in cluster_indices:
            dist = np.linalg.norm(data[idx] - centroids[i])
            if dist < min_dist:
                min_dist = dist
                best_idx = idx
        stable_medoids.append(all_poles[best_idx])
    
    codebook = np.array(stable_medoids)
    print(f"   Codebook Built: {len(codebook)} stable entries.")
    
    # Pre-compute codebook log-spectrograms for fast matching
    codebook_logs = []
    for idx in range(len(codebook)):
        _, h_code = signal.freqz([1.0], codebook[idx], worN=128)
        codebook_logs.append(np.log10(np.abs(h_code) + 1e-9))

    # === PIPELINE STAGE 2: CLOSED-LOOP AUTO-ENCODE ===
    print("3. Running Auto-Encode & Benchmark...")
    reconstructed_audio = np.array([])
    filter_memory = np.zeros(ORDER)
    glottal_phase = 0
    
    lsd_scores = []

    for i in range(0, len(audio) - FRAME_LEN, FRAME_LEN):
        frame = audio[i : i + FRAME_LEN]
        
        is_voiced, f0, energy = v_uv_discriminator(frame)
        target_poles = extract_poles_raw(frame)
        
        # Spectral Match
        _, h_target = signal.freqz([1.0], target_poles, worN=128)
        log_target = np.log10(np.abs(h_target) + 1e-9)
        log_target_norm = log_target - np.mean(log_target)
        
        best_dist = float('inf')
        winning_index = 0
        
        for idx in range(len(codebook)):
            log_code_norm = codebook_logs[idx] - np.mean(codebook_logs[idx])
            dist = np.mean((log_target_norm - log_code_norm)**2)
            if dist < best_dist:
                best_dist = dist
                winning_index = idx
                
        winning_poles = codebook[winning_index]
        
        synth_frame, filter_memory, glottal_phase = synthesize_frame(
            f0, energy, winning_poles, is_voiced, filter_memory, glottal_phase
        )
        
        # Automated Benchmark: Frame-by-Frame LSD
        if energy > 1e-6: # Only benchmark active speech frames
            lsd = calc_lsd(frame, synth_frame)
            lsd_scores.append(lsd)
        
        reconstructed_audio = np.append(reconstructed_audio, synth_frame)

    # === PIPELINE STAGE 3: OUTPUT ===
    mean_lsd = np.mean(lsd_scores)
    print("\n--- BENCHMARK RESULTS ---")
    print(f"Mean Log-Spectral Distance (LSD): {mean_lsd:.2f} dB")
    
    if mean_lsd < 1.5:
        print("Verdict: EXCELLENT (Near-transparent parametric reconstruction)")
    elif mean_lsd < 3.0:
        print("Verdict: ACCEPTABLE (Highly intelligible communications quality)")
    else:
        print("Verdict: POOR (Significant spectral distortion present)")

    print(f"\n4. Saving Output Audio to '{OUTPUT_FILE}'...")
    reconstructed_audio = signal.lfilter([1.0], [1.0, -0.97], reconstructed_audio)
    
    # Normalize before saving to 16-bit PCM
    reconstructed_audio = reconstructed_audio / np.max(np.abs(reconstructed_audio))
    wavfile.write(OUTPUT_FILE, FS, np.int16(reconstructed_audio * 32767))
    print("Done.")

if __name__ == "__main__":
    main()