#!/usr/bin/env python3
"""
Synthetic Digital Organism (SDO) - 44.1kHz 7-Formant Resonator Vocal Synthesizer Engine
Features:
 - 44.1 kHz Studio Quality Audio PCM output.
 - Rosenberg Glottal Flow Pulse vocal fold excitation model.
 - 7-Formant IIR Biquad Resonators (F1-F7) covering 200 Hz to 9500 Hz.
 - Nasal Zero Anti-Resonator notch filter (Zn) for nasal acoustic coupling.
 - Zero-static silence gate with filter memory reset.
 - Slew-rate parameter smoothing for pop-free coarticulation transitions.
"""

import sys
import os
import signal

if hasattr(signal, 'SIGPIPE'):
    signal.signal(signal.SIGPIPE, signal.SIG_DFL)

import json
import threading
import numpy as np
import sounddevice as sd
from scipy.signal import lfilter

SAMPLE_RATE = 48000
BLOCK_SIZE = int(SAMPLE_RATE / 48) # ~1000 samples for 48 FPS low-latency articulatory sync

# Initial state for 10 vocal parameters:
# [pitch, volume, f1, f2, f3, f4, f5, f6, f7, noise_level]
current_params = [0.5, 0.0, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.0]
param_lock = threading.Lock()

last_audio_frame = None
frame_lock = threading.Lock()

def cascade_resonator(freq, bw, fs):
    """Klatt Cascade Resonator with Unity DC gain."""
    freq = min(max(20.0, float(freq)), float(fs) * 0.45)
    r = np.exp(-np.pi * bw / fs)
    c = 2.0 * r * np.cos(2.0 * np.pi * freq / fs)
    a = [1.0, -c, r * r]
    gain = 1.0 - c + r * r
    b = [gain]
    return b, a

def cascade_notch(freq, bw, fs):
    """Anti-resonator (nasal zero notch)."""
    freq = min(max(20.0, float(freq)), float(fs) * 0.45)
    r = np.exp(-np.pi * bw / fs)
    c = 2.0 * r * np.cos(2.0 * np.pi * freq / fs)
    gain = 1.0 - c + r * r
    if abs(gain) < 1e-6: gain = 1.0
    b = [1.0 / gain, -c / gain, (r * r) / gain]
    a = [1.0]
    return b, a

def lfilter_block(b, a, x, zi=None):
    if zi is None:
        zi = np.zeros(max(len(a), len(b)) - 1)
    y, zi = lfilter(b, a, x, zi=zi)
    return y, zi

# Internal state for smooth parameter interpolation (10 parameters)
smoothed_params = np.array([0.5, 0.0, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.0], dtype=np.float64)

# Phase & filter state memories for 7 formants + 1 nasal zero + 1 noise shaper + lip radiation
phase = 0.0
lfo_phase = 0.0
zi_f1 = np.zeros(2)
zi_f2 = np.zeros(2)
zi_f3 = np.zeros(2)
zi_f4 = np.zeros(2)
zi_f5 = np.zeros(2)
zi_f6 = np.zeros(2)
zi_f7 = np.zeros(2)
zi_zn = np.zeros(2)
zi_noise = np.zeros(2)
zi_lip = np.zeros(1)

def create_mel_filterbank(num_filters=128, n_fft=1024, fs=48000, f_min=20.0, f_max=16000.0):
    """Creates a 128-channel logarithmic/Mel tonotopic cochlear filterbank covering 20Hz-16kHz (matching av_driver)."""
    def hz_to_mel(hz):
        return 2595.0 * np.log10(1.0 + hz / 700.0)
    def mel_to_hz(mel):
        return 700.0 * (10.0**(mel / 2595.0) - 1.0)

    mel_min = hz_to_mel(f_min)
    mel_max = hz_to_mel(f_max)
    mel_points = np.linspace(mel_min, mel_max, num_filters + 2)
    hz_points = mel_to_hz(mel_points)
    
    n_freqs = n_fft // 2 + 1
    fft_freqs = np.linspace(0, fs / 2, n_freqs)
    
    fb = np.zeros((num_filters, n_freqs), dtype=np.float32)
    for m in range(1, num_filters + 1):
        f_left = hz_points[m - 1]
        f_center = hz_points[m]
        f_right = hz_points[m + 1]
        
        # Rising slope
        up_mask = (fft_freqs >= f_left) & (fft_freqs <= f_center)
        if f_center > f_left:
            fb[m - 1, up_mask] = (fft_freqs[up_mask] - f_left) / (f_center - f_left)
            
        # Falling slope
        down_mask = (fft_freqs >= f_center) & (fft_freqs <= f_right)
        if f_right > f_center:
            fb[m - 1, down_mask] = (f_right - fft_freqs[down_mask]) / (f_right - f_center)
            
    return fb

vocal_mel_filterbank = None
vocal_hann_window = None

def lf_glottal_pulse(phase_arr, T_o=0.45, T_c=0.18):
    """Liljencrants-Fant (LF) inspired glottal flow pulse with natural glottal opening and exponential closure."""
    pulse = np.zeros_like(phase_arr, dtype=np.float64)
    # Open phase: smooth cosine rise
    m_open = phase_arr < T_o
    pulse[m_open] = 0.5 * (1.0 - np.cos(np.pi * phase_arr[m_open] / T_o))
    # Closing phase: steep closure with exponential return
    m_close = (phase_arr >= T_o) & (phase_arr < (T_o + T_c))
    rel_phase = (phase_arr[m_close] - T_o) / T_c
    pulse[m_close] = np.cos(np.pi * rel_phase * 0.5) * np.exp(-3.0 * rel_phase)
    return pulse - 0.20

active_sample_rate = 48000
vol_env = 0.0
vocal_tract_scale = 1.0

def audio_callback(outdata, frames, time_info, status):
    global phase, lfo_phase, zi_f1, zi_f2, zi_f3, zi_f4, zi_f5, zi_f6, zi_f7, zi_zn, zi_noise, zi_lip, smoothed_params, active_sample_rate, vol_env, last_audio_frame
    if status and not status.output_underflow:
        print(f"[VocalDriver] Audio status: {status}", file=sys.stderr)
        
    with param_lock:
        target_params = np.array(current_params, dtype=np.float64)

    fs = active_sample_rate if active_sample_rate > 0 else 48000

    # 1. Adaptive Parameter Slew-Rate Coarticulation
    delta_noise = abs(target_params[9] - smoothed_params[9])
    delta_f12 = abs(target_params[2] - smoothed_params[2]) + abs(target_params[3] - smoothed_params[3])
    if delta_noise > 0.15 or delta_f12 > 0.35:
        slew_rate = 0.85 # Fast release for plosive bursts and fricatives
    else:
        slew_rate = 0.45 # Smooth biological vocal fold/tongue glide kinematics

    smoothed_params = (1.0 - slew_rate) * smoothed_params + slew_rate * target_params
    p_pitch, p_vol, p_f1, p_f2, p_f3, p_f4, p_f5, p_f6, p_f7, p_noise = smoothed_params

    # 2. Pitch Micro-Vibrato & Affective Intonation Contouring (5.0 Hz LFO with ~1.2% vibrato depth)
    lfo_phase = (lfo_phase + frames * 5.0 / fs) % 1.0
    vibrato = 1.0 + 0.012 * np.sin(2.0 * np.pi * lfo_phase)
    pitch = (90.0 + (p_pitch * 140.0)) * vibrato  # Natural human F0: 90Hz to 230Hz
    vol = p_vol * 1.0  # Full clear amplitude

    # 3. 7-Formant Frequency Mapping (Hz) with Vocal Tract Length Normalization (VTLN)
    vt_scale = vocal_tract_scale
    f1 = (200.0 + (p_f1 * 800.0)) * vt_scale      # F1: 200 - 1000 Hz (Jaw height / Vowel openness)
    f2 = (600.0 + (p_f2 * 1900.0)) * vt_scale     # F2: 600 - 2500 Hz (Tongue advancement)
    f3 = (1500.0 + (p_f3 * 2000.0)) * vt_scale    # F3: 1500 - 3500 Hz (Tongue tip / Lip aperture)
    f4 = (3500.0 + (p_f4 * 700.0)) * vt_scale     # F4: 3500 - 4200 Hz (Head voice brilliance)
    f5 = (4500.0 + (p_f5 * 700.0)) * vt_scale     # F5: 4500 - 5200 Hz (Nasal cavity / Epiglottal warmth)
    f6 = (5500.0 + (p_f6 * 700.0)) * vt_scale     # F6: 5500 - 6200 Hz (Laryngeal tube length)
    f7 = (6800.0 + (p_f7 * 2700.0)) * vt_scale    # F7: 6800 - 9500 Hz (Dental & alveolar sibilance)

    # Formant Bandwidths (Klatt standard)
    bw1 = 60.0 + 0.05 * f1
    bw2 = 90.0 + 0.04 * f2
    bw3 = 120.0 + 0.03 * f3
    bw4 = 200.0
    bw5 = 250.0
    bw6 = 300.0
    bw7 = 400.0

    # 4. Glottal Excitation Waveform
    t = np.arange(frames) / fs
    sample_phases = (phase + t * pitch) % 1.0
    phase = (phase + frames * pitch / fs) % 1.0

    glottal_source = lf_glottal_pulse(sample_phases)
    white_noise = np.random.normal(0, 0.20, frames)

    # 5. Klatt Cascade Formant Resonators (Series Acoustic Tube)
    b1, a1 = cascade_resonator(f1, bw1, fs)
    b2, a2 = cascade_resonator(f2, bw2, fs)
    b3, a3 = cascade_resonator(f3, bw3, fs)
    b4, a4 = cascade_resonator(f4, bw4, fs)
    b5, a5 = cascade_resonator(f5, bw5, fs)
    b6, a6 = cascade_resonator(f6, bw6, fs)
    b7, a7 = cascade_resonator(f7, bw7, fs)
    b_zn, a_zn = cascade_notch(1500.0, 200.0, fs) # Nasal anti-zero

    # Execute Series Cascade: Glottal -> Nasal Zero -> F1 -> F2 -> F3 -> F4 -> F5 -> F6 -> F7
    x, zi_zn = lfilter_block(b_zn, a_zn, glottal_source, zi_zn)
    x, zi_f1 = lfilter_block(b1, a1, x, zi_f1)
    x, zi_f2 = lfilter_block(b2, a2, x, zi_f2)
    x, zi_f3 = lfilter_block(b3, a3, x, zi_f3)
    x, zi_f4 = lfilter_block(b4, a4, x, zi_f4)
    x, zi_f5 = lfilter_block(b5, a5, x, zi_f5)
    x, zi_f6 = lfilter_block(b6, a6, x, zi_f6)
    x, zi_f7 = lfilter_block(b7, a7, x, zi_f7)

    # Lip Radiation Filter (+6 dB/octave differentiation): y[n] = x[n] - 0.96 x[n-1]
    voiced_out, zi_lip = lfilter_block([1.0, -0.96], [1.0], x, zi_lip)

    # Parallel Fricative & Sibilance Branch (Tuned High-Frequency Acoustic Turbulence for /s/, /sh/, /f/, /th/)
    fric_freq = min(max(1800.0, float(f7 if p_f7 > 0.4 else (f3 * 1.5 + 800.0))), float(fs) * 0.45)
    fric_bw = 1400.0
    b_fric, a_fric = cascade_resonator(fric_freq, fric_bw, fs)
    
    # High-pass pre-emphasis on noise source to prevent low-frequency rumble
    noise_source = white_noise.copy()
    if len(noise_source) > 1:
        noise_source[1:] -= 0.6 * noise_source[:-1]
    fric_out, zi_noise = lfilter_block(b_fric, a_fric, noise_source, zi_noise)

    # 6. Smooth Cosine Volume Envelope Gate (Click-Free Transition)
    target_env = 1.0 if p_vol >= 0.01 else 0.0
    vol_env = 0.82 * vol_env + 0.18 * target_env

    # Mix Voiced Cascade + Parallel Fricative Noise (Balanced for full voiced vowels)
    total_voiced = voiced_out * 1.4 * (1.0 - p_noise * 0.5)
    total_fric = fric_out * (p_noise * 0.8)
    
    mixed = np.tanh((total_voiced + total_fric) * vol * vol_env * 3.0) * 0.95

    with frame_lock:
        last_audio_frame = mixed.copy()

    outdata[:] = mixed.reshape(-1, 1)

def reporting_loop():
    global vocal_mel_filterbank, vocal_hann_window
    import time
    while True:
        frame_to_process = None
        with frame_lock:
            global last_audio_frame
            if last_audio_frame is not None:
                frame_to_process = last_audio_frame.copy()
        
        if frame_to_process is not None:
            n_samples = len(frame_to_process)
            fs = active_sample_rate if active_sample_rate > 0 else 48000
            
            if vocal_hann_window is None or len(vocal_hann_window) != n_samples:
                vocal_hann_window = np.hanning(n_samples).astype(np.float32)
            if vocal_mel_filterbank is None or vocal_mel_filterbank.shape[1] != (n_samples // 2 + 1):
                vocal_mel_filterbank = create_mel_filterbank(128, n_fft=n_samples, fs=fs, f_min=20.0, f_max=16000.0)

            windowed_frame = frame_to_process * vocal_hann_window
            fft_vals = np.abs(np.fft.rfft(windowed_frame)).astype(np.float32)
            tonotopic_bins = np.dot(vocal_mel_filterbank, fft_vals)

            max_val = float(np.max(tonotopic_bins))
            if max_val > 1e-5:
                norm_fft = (tonotopic_bins / max_val).astype(np.float32)
            else:
                norm_fft = tonotopic_bins.astype(np.float32)
            
            try:
                print(json.dumps({"type": "vocal_fft", "bins": norm_fft.tolist()}), flush=True)
            except (BrokenPipeError, IOError, OSError):
                os._exit(0)
            except Exception as e:
                try:
                    print(json.dumps({"status": "error", "message": str(e)}), file=sys.stderr, flush=True)
                except:
                    os._exit(0)

        time.sleep(1.0 / 48.0)

def main():
    stream = None
    candidate_outputs = [(0, None, "System Default Sink")]
    for i, dev in enumerate(sd.query_devices()):
        if dev['max_output_channels'] > 0:
            name_lower = dev['name'].lower()
            is_hdmi = 'hdmi' in name_lower or 'displayport' in name_lower
            is_speaker = ('speaker' in name_lower or 'headphone' in name_lower or 'analog' in name_lower or 'pulse' in name_lower or 'pipewire' in name_lower or 'default' in name_lower) and not is_hdmi
            prio = 1 if is_speaker else (3 if is_hdmi else 2)
            candidate_outputs.append((prio, i, dev['name']))

    candidate_outputs.sort(key=lambda x: (x[0], x[1] if x[1] is not None else -1))
    for prio, dev_id, dev_name in candidate_outputs:
        for sr in [48000, 44100, 16000]:
            try:
                block = int(sr / 48)
                test_stream = sd.OutputStream(device=dev_id, samplerate=sr, channels=1, dtype='float32', blocksize=block, callback=audio_callback)
                test_stream.start()
                stream = test_stream
                global active_sample_rate
                active_sample_rate = int(sr)
                sys.stderr.write(f"[VocalDriver] Physical Speaker active: Device {dev_id} ({dev_name}) @ {sr}Hz\n")
                break
            except Exception:
                stream = None
        if stream is not None:
            break
            
    if stream is None:
        try:
            stream = sd.OutputStream(samplerate=48000, channels=1, dtype='float32', callback=audio_callback)
            stream.start()
        except Exception as e:
            sys.stderr.write(f"[VocalDriver] Fallback speaker stream init failed: {e}\n")
    
    print(json.dumps({"status": "ready"}), flush=True)
    
    reporter = threading.Thread(target=reporting_loop)
    reporter.daemon = True
    reporter.start()
    
    try:
        for line in sys.stdin:
            line = line.strip()
            if not line:
                continue
            try:
                cmd = json.loads(line)
                if cmd.get('type') == 'vocal_dof':
                    dof = cmd.get('dof')
                    val = cmd.get('value', 0.0)
                    with param_lock:
                        global current_params
                        norm_val = val if (val >= 0.0 and val <= 1.0) else (val + 1.0) / 2.0
                        if dof == 'vocal_pitch': current_params[0] = norm_val
                        elif dof == 'vocal_volume': current_params[1] = max(0.0, min(1.0, val))
                        elif dof == 'vocal_formant_1': current_params[2] = norm_val
                        elif dof == 'vocal_formant_2': current_params[3] = norm_val
                        elif dof == 'vocal_formant_3': current_params[4] = norm_val
                        elif dof == 'vocal_formant_4': current_params[5] = norm_val
                        elif dof == 'vocal_formant_5': current_params[6] = norm_val
                        elif dof == 'vocal_formant_6': current_params[7] = norm_val
                        elif dof == 'vocal_formant_7': current_params[8] = norm_val
                        elif dof == 'vocal_noise': current_params[9] = max(0.0, min(1.0, val))
                elif cmd.get('type') == 'vocal':
                    vec = cmd.get('params', [])
                    if len(vec) >= 6:
                        with param_lock:
                            # Support both 6-param legacy and 10-param 7-formant vectors
                            scaled = [v if (v >= 0.0 and v <= 1.0) else (v + 1.0) / 2.0 for v in vec]
                            if len(scaled) >= 10:
                                current_params = scaled[:10]
                            else:
                                current_params[0] = scaled[0] # pitch
                                current_params[1] = max(0.0, vec[1]) # volume
                                current_params[2] = scaled[2] # f1
                                current_params[3] = scaled[3] # f2
                                current_params[4] = scaled[4] # f3
                                current_params[9] = max(0.0, scaled[5]) # noise
                elif cmd.get('type') == 'vocal_maturation':
                    global vocal_tract_scale
                    vocal_tract_scale = float(cmd.get('scale', 1.0))
            except (BrokenPipeError, IOError, OSError):
                sys.exit(0)
            except Exception as e:
                try:
                    print(json.dumps({"status": "error", "message": str(e)}), file=sys.stderr, flush=True)
                except:
                    sys.exit(0)
    except (KeyboardInterrupt, BrokenPipeError, IOError, OSError):
        pass
    except Exception as e:
        try:
            print(json.dumps({"status": "error", "message": str(e)}), file=sys.stderr, flush=True)
        except:
            pass
    finally:
        if stream is not None:
            try:
                stream.stop()
                stream.close()
            except Exception:
                pass

if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(0)
