physics_audio/run_real_audio.py

#!/usr/bin/env python3
"""
run_real_audio.py
Real audio extension for the physics_audio / mlpa Stiefel manifold model.

Usage:
    python run_real_audio.py --audio path/to/guitar_note.wav --duration 3.0
    python run_real_audio.py --audio_dir path/to/notes/ --analyze_batch
    python run_real_audio.py --audio note.wav --jump_test --low_patience 300
    python run_real_audio.py --live_mic --duration 5 --fit_after
    python run_real_audio.py --stft_benchmark
"""

import argparse
import csv
import json
import librosa
import numpy as np
import soundfile as sf
import torch
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from pathlib import Path
from typing import Optional

from training_evaluation.config import (
    K_MODES, DIM, REAL_AUDIO_MAX_STEPS, REAL_AUDIO_STFT_WEIGHT, REAL_AUDIO_SR,
    JUMP_TEST_LOW_PATIENCE, JUMP_TEST_MIN_STEP, JUMP_TEST_FORCE_EVERY, JUMP_TEST_PLATEAU_AT,
    LIVE_MIC_BLOCKSIZE, LIVE_MIC_DEFAULT_SECONDS, GPU_STFT_N_FFT, GPU_STFT_HOP,
    fr_invariant_weight,
    fr_invariant_coupling,
    fr_invariant_speed,
    fr_invariant_inharm,
    fr_invariant_modal,
    fr_replace_mse_priors,
    REAL_AUDIO_FR_MODE_WEIGHT,
)
from training_evaluation.model import StiefelDampedCoupledInharmGR
from training_evaluation.audio_utils import (
    extract_partials_piptrack,
    extract_partials_streaming,
    StreamingPartialTracker,
    estimate_physical_params,
    partials_to_manifold_data,
    initialize_model_from_physics,
    build_prior_targets,
    modal_synthesis,
)
from training_evaluation.streaming_stft import StreamingGPUSTFT
from training_evaluation.batch_analysis import analyze_batch_results
from training_evaluation.trainer import run_single_seed

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
AUDIO_EXTENSIONS = {'.wav', '.flac', '.ogg', '.mp3', '.aiff', '.aif'}


def load_audio(audio_path: str, sr: int = REAL_AUDIO_SR, duration: float = None):
    y, sr = librosa.load(audio_path, sr=sr, duration=duration, mono=True)
    y, _ = librosa.effects.trim(y, top_db=25)
    y = librosa.util.normalize(y)
    return y, sr


def build_jump_test_config(args) -> Optional[dict]:
    if not getattr(args, 'jump_test', False):
        return None
    cfg = {'verbose': True}
    if args.low_patience is not None:
        cfg['low_patience'] = args.low_patience
    else:
        cfg['low_patience'] = JUMP_TEST_LOW_PATIENCE
    if args.min_step_for_jump is not None:
        cfg['min_step_for_jump'] = args.min_step_for_jump
    else:
        cfg['min_step_for_jump'] = JUMP_TEST_MIN_STEP
    if args.force_jump_every is not None:
        cfg['force_jump_every'] = args.force_jump_every
    elif getattr(args, 'force_jumps', False):
        cfg['force_jump_every'] = JUMP_TEST_FORCE_EVERY
    if args.artificial_plateau_at is not None:
        cfg['artificial_plateau_at'] = args.artificial_plateau_at
    elif getattr(args, 'artificial_plateau', False):
        cfg['artificial_plateau_at'] = JUMP_TEST_PLATEAU_AT
    cfg['pop_size'] = getattr(args, 'jump_pop_size', None) or 6
    cfg['rollout_horizon'] = getattr(args, 'jump_rollout_horizon', None) or 800
    return cfg


def save_spectrogram_comparison(y_target, y_synth, sr, output_path, title="Real vs Synthesized"):
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    for ax, y, label in zip(axes[:2], [y_target, y_synth], ["Target", "Synthesized"]):
        S = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=128)
        S_db = librosa.power_to_db(S, ref=np.max)
        img = librosa.display.specshow(S_db, sr=sr, x_axis='time', y_axis='mel', ax=ax)
        ax.set_title(label)
        fig.colorbar(img, ax=ax, format='%+2.0f dB')
    min_len = min(len(y_target), len(y_synth))
    residual = y_target[:min_len] - y_synth[:min_len]
    S_res = librosa.feature.melspectrogram(y=residual, sr=sr, n_mels=128)
    S_res_db = librosa.power_to_db(S_res, ref=np.max)
    img = librosa.display.specshow(S_res_db, sr=sr, x_axis='time', y_axis='mel', ax=axes[2])
    axes[2].set_title("Residual")
    fig.colorbar(img, ax=axes[2], format='%+2.0f dB')
    fig.suptitle(title)
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    plt.close()


def save_partial_plot(partial_freqs, partial_amps, times, output_path, title_suffix=""):
    fig, axes = plt.subplots(2, 1, figsize=(12, 6), sharex=True)
    for k in range(partial_freqs.shape[0]):
        axes[0].plot(times, partial_freqs[k], label=f"mode {k+1}")
        axes[1].plot(times, partial_amps[k], alpha=0.8)
    axes[0].set_ylabel("Frequency (Hz)")
    axes[0].legend(ncol=4, fontsize=8)
    axes[0].set_title(f"Tracked partial frequencies{title_suffix}")
    axes[1].set_ylabel("Magnitude")
    axes[1].set_xlabel("Time (s)")
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    plt.close()


def discover_audio_files(path: Path) -> list[Path]:
    if path.is_file():
        return [path]
    return sorted(p for p in path.rglob('*') if p.is_file() and p.suffix.lower() in AUDIO_EXTENSIONS)


def process_single_note(
    audio_path: Path,
    output_dir: Path,
    duration: float,
    max_steps: int,
    stft_weight: float,
    seed: int,
    streaming: bool = False,
    jump_test: Optional[dict] = None,
    use_gpu_stft: bool = False,
    fr_invariant_weight_override: float | None = None,
    fr_invariant_coupling_override: float | None = None,
    fr_invariant_speed_override: float | None = None,
    fr_invariant_inharm_override: float | None = None,
    fr_invariant_modal_override: float | None = None,
    fr_mode_weight_override: float | None = None,
    fr_augment_priors: bool = False,
) -> dict:
    output_dir.mkdir(parents=True, exist_ok=True)
    note_name = audio_path.stem

    print(f"\n{'='*60}")
    print(f"Processing: {audio_path}")

    y, sr = load_audio(str(audio_path), duration=duration)
    audio_duration = len(y) / sr
    print(f"Audio loaded: {audio_duration:.2f}s @ {sr} Hz")

    tracker_label = "streaming" if streaming else "piptrack + continuity"
    print(f"Extracting partials ({tracker_label})...")
    if streaming:
        partial_freqs, partial_amps, times, _ = extract_partials_streaming(y, sr, n_partials=K_MODES)
    else:
        partial_freqs, partial_amps, times, _ = extract_partials_piptrack(y, sr, n_partials=K_MODES)

    if use_gpu_stft:
        gpu_stft = StreamingGPUSTFT(n_fft=GPU_STFT_N_FFT, hop_length=GPU_STFT_HOP, sr=sr)
        chunk_t = torch.tensor(y, device=device, dtype=torch.float32)
        for i in range(0, len(chunk_t), GPU_STFT_HOP):
            gpu_stft.push(chunk_t[i:i + GPU_STFT_HOP])
        bench = gpu_stft.benchmark(n_chunks=50, chunk_samples=GPU_STFT_HOP)
        print(f"  GPU STFT: hop_latency={bench['hop_latency_ms']:.2f}ms mean_push={bench['mean_push_ms']:.3f}ms")

    save_partial_plot(partial_freqs, partial_amps, times, str(output_dir / "partials_tracked.png"),
                      title_suffix=f" ({tracker_label})")

    damping_rates, f0_est, b_est = estimate_physical_params(partial_freqs, partial_amps, times)
    print(f"Estimated f0 ≈ {f0_est:.1f} Hz, B ≈ {b_est:.6f}")

    data_points, model_times = partials_to_manifold_data(partial_freqs, partial_amps, times, f0_est)
    model = StiefelDampedCoupledInharmGR(DIM, K_MODES, data_points[0])
    model = initialize_model_from_physics(model, damping_rates, f0_est, b_est, data_points)

    prior_targets = build_prior_targets(damping_rates, f0_est, b_est, partial_amps=partial_amps)
    target_waveform = torch.tensor(y, device=device, dtype=torch.float32)

    active_fr_inv = fr_invariant_weight if fr_invariant_weight_override is None else fr_invariant_weight_override
    active_fr_coup = (
        fr_invariant_coupling if fr_invariant_coupling_override is None else fr_invariant_coupling_override
    )
    active_fr_speed = (
        fr_invariant_speed if fr_invariant_speed_override is None else fr_invariant_speed_override
    )
    active_fr_inharm = (
        fr_invariant_inharm if fr_invariant_inharm_override is None else fr_invariant_inharm_override
    )
    active_fr_modal = (
        fr_invariant_modal if fr_invariant_modal_override is None else fr_invariant_modal_override
    )
    active_fr_mode = (
        REAL_AUDIO_FR_MODE_WEIGHT if fr_mode_weight_override is None else fr_mode_weight_override
    )
    jump_label = " + jump_test" if jump_test else ""
    print(
        f"Running optimization ({max_steps} steps, STFT={stft_weight}, "
        f"fr_invariant={active_fr_inv}, fr_coupling={active_fr_coup}, "
        f"fr_speed={active_fr_speed}, fr_inharm={active_fr_inharm}, "
        f"fr_modal={active_fr_modal}, fr_mode={active_fr_mode}{jump_label})..."
    )
    result = run_single_seed(
        seed=seed,
        real_audio_data=data_points,
        real_audio_times=model_times,
        target_waveform=target_waveform,
        audio_sr=sr,
        audio_duration=audio_duration,
        prior_targets=prior_targets,
        max_steps=max_steps,
        stft_weight=stft_weight,
        fr_invariant_weight=active_fr_inv,
        fr_invariant_coupling=active_fr_coup,
        fr_invariant_speed=active_fr_speed,
        fr_invariant_inharm=active_fr_inharm,
        fr_invariant_modal=active_fr_modal,
        fr_mode_weight=active_fr_mode,
        fr_replace_mse_priors=not fr_augment_priors,
        preinitialized_model=model,
        jump_test=jump_test,
    )

    learned_freqs = result['full_freq'].cpu().numpy()
    learned_damps = result['damping_rates'].cpu().numpy()
    coupling_strength = result['coupling_strength']
    coupling_skew = result['coupling_skew'].cpu().numpy()

    print(f"  Recon MSE: {result['total_recon_mse_pred']:.6f} | Jumps: {result['jumps']} | "
          f"Coupling: {coupling_strength:.4f}")

    y_synth = modal_synthesis(learned_freqs, learned_damps, sr, audio_duration,
                              coupling_strength=coupling_strength, coupling_skew=coupling_skew)
    sf.write(str(output_dir / "synthesized.wav"), y_synth, sr)
    save_spectrogram_comparison(y, y_synth, sr, str(output_dir / "spectrogram_comparison.png"),
                                title=f"{note_name}: Real vs Synthesized")

    summary = {
        'file': str(audio_path), 'name': note_name, 'duration_s': audio_duration,
        'f0_est_hz': f0_est, 'b_est': b_est,
        'recon_mse': result['total_recon_mse_pred'],
        'damping_rmse': result['damping_rmse'], 'damping_corr': result['damping_corr'],
        'coupling_strength': coupling_strength, 'jumps': result['jumps'],
        'wall_time_s': result['wall_time'],
        'learned_freqs': learned_freqs.tolist(), 'learned_damps': learned_damps.tolist(),
        'tracker': tracker_label,
    }
    with open(output_dir / "summary.json", 'w') as f:
        json.dump(summary, f, indent=2)
    return summary


def run_batch(audio_dir, output_dir, duration, max_steps, stft_weight, seed, streaming,
              jump_test, use_gpu_stft, analyze_batch, n_clusters,
              fr_invariant_weight_override=None, fr_invariant_coupling_override=None,
              fr_invariant_speed_override=None, fr_invariant_inharm_override=None,
              fr_invariant_modal_override=None, fr_mode_weight_override=None,
              fr_augment_priors: bool = False):
    files = discover_audio_files(audio_dir)
    if not files:
        raise FileNotFoundError(f"No audio files in {audio_dir}")
    print(f"Batch: {len(files)} file(s)")
    results = []
    for i, p in enumerate(files):
        results.append(process_single_note(
            p, output_dir / p.stem, duration, max_steps, stft_weight,
            seed=seed + i, streaming=streaming, jump_test=jump_test, use_gpu_stft=use_gpu_stft,
            fr_invariant_weight_override=fr_invariant_weight_override,
            fr_invariant_coupling_override=fr_invariant_coupling_override,
            fr_invariant_speed_override=fr_invariant_speed_override,
            fr_invariant_inharm_override=fr_invariant_inharm_override,
            fr_invariant_modal_override=fr_invariant_modal_override,
            fr_mode_weight_override=fr_mode_weight_override,
            fr_augment_priors=fr_augment_priors,
        ))
    csv_path = output_dir / "batch_summary.csv"
    fields = ['name', 'file', 'duration_s', 'f0_est_hz', 'b_est', 'recon_mse',
              'damping_rmse', 'damping_corr', 'coupling_strength', 'jumps', 'wall_time_s', 'tracker']
    with open(csv_path, 'w', newline='') as f:
        w = csv.DictWriter(f, fieldnames=fields)
        w.writeheader()
        for r in results:
            w.writerow({k: r[k] for k in fields})
    print(f"Batch CSV: {csv_path}")
    if analyze_batch:
        analyze_batch_results(results, output_dir, n_clusters=n_clusters)
    return results


def run_stft_benchmark():
    print("GPU streaming STFT benchmark")
    for dev_name in (['cuda'] if torch.cuda.is_available() else []) + ['cpu']:
        dev = torch.device(dev_name)
        stft = StreamingGPUSTFT(n_fft=GPU_STFT_N_FFT, hop_length=GPU_STFT_HOP, dev=dev)
        bench = stft.benchmark(n_chunks=500, chunk_samples=GPU_STFT_HOP)
        print(f"  [{dev_name}] hop={bench['hop_latency_ms']:.2f}ms "
              f"mean={bench['mean_push_ms']:.3f}ms max={bench['max_push_ms']:.3f}ms "
              f"sub10ms={bench['under_10ms']}")


def run_live_mic(duration, sr, output_dir, fit_after, max_steps, stft_weight, streaming):
    try:
        import sounddevice as sd
    except ImportError as e:
        raise ImportError("Install sounddevice: pip install sounddevice") from e

    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    tracker = StreamingPartialTracker(sr=sr, n_partials=K_MODES)
    gpu_stft = StreamingGPUSTFT(n_fft=GPU_STFT_N_FFT, hop_length=GPU_STFT_HOP, sr=sr)
    recorded = []

    print(f"Live mic: {duration}s @ {sr}Hz (blocksize={LIVE_MIC_BLOCKSIZE})")
    print("Speak or play now...")

    def callback(indata, frames, time_info, status):
        if status:
            print(f"  [sd] {status}")
        chunk = indata[:, 0].astype(np.float32)
        recorded.append(chunk.copy())
        state = tracker.process_chunk(chunk)
        gpu_stft.push(torch.tensor(chunk, device=device, dtype=torch.float32))
        if state is not None:
            print(f"  t={state['time']:.2f}s f0={state['freqs'][0]:.1f}Hz "
                  f"gpu_stft={gpu_stft._last_frame_ms:.2f}ms", flush=True)

    with sd.InputStream(samplerate=sr, channels=1, blocksize=LIVE_MIC_BLOCKSIZE, callback=callback):
        sd.sleep(int(duration * 1000))

    if not recorded:
        print("No audio captured.")
        return

    y = np.concatenate(recorded)
    y, _ = librosa.effects.trim(y, top_db=30)
    if y.size == 0:
        print("Captured silence only.")
        return
    y = librosa.util.normalize(y)

    wav_path = output_dir / "live_recording.wav"
    sf.write(str(wav_path), y, sr)
    print(f"Saved recording: {wav_path} ({len(y)/sr:.2f}s)")

    partial_freqs, partial_amps, times, _ = tracker.finalize()
    save_partial_plot(partial_freqs, partial_amps, times,
                      str(output_dir / "live_partials.png"), title_suffix=" (live)")

    bench = gpu_stft.benchmark(n_chunks=100, chunk_samples=GPU_STFT_HOP)
    print(f"GPU STFT benchmark: mean={bench['mean_push_ms']:.3f}ms sub10ms={bench['under_10ms']}")

    if fit_after:
        print("Fitting model on live recording...")
        process_single_note(
            wav_path, output_dir / "live_fit", len(y) / sr,
            max_steps, stft_weight, seed=0, streaming=streaming,
        )


def demo_streaming_live(audio_path, duration=3.0):
    y, sr = load_audio(audio_path, duration=duration)
    tracker = StreamingPartialTracker(sr=sr, n_partials=K_MODES)
    gpu_stft = StreamingGPUSTFT(n_fft=GPU_STFT_N_FFT, hop_length=GPU_STFT_HOP, sr=sr)
    chunk_size = tracker.chunk_size
    for start in range(0, len(y), chunk_size):
        chunk = y[start:start + chunk_size]
        state = tracker.process_chunk(chunk)
        gpu_stft.push(torch.tensor(chunk, device=device, dtype=torch.float32))
        if state:
            print(f"  t={state['time']:.3f}s f0={state['freqs'][0]:.1f}Hz "
                  f"gpu_stft={gpu_stft._last_frame_ms:.2f}ms")
    freqs, _, times = tracker.get_trajectories()
    print(f"Done: {times.shape[0]} frames, f0 [{freqs[0].min():.1f}, {freqs[0].max():.1f}] Hz")


def main():
    parser = argparse.ArgumentParser(description="Real-audio fitting pipeline (Phase 4)")
    src = parser.add_mutually_exclusive_group(required=True)
    src.add_argument("--audio", type=str, help="Single audio file")
    src.add_argument("--audio_dir", type=str, help="Folder for batch processing")
    src.add_argument("--live_mic", action="store_true", help="Record from microphone")
    src.add_argument("--stft_benchmark", action="store_true", help="Benchmark GPU streaming STFT")

    parser.add_argument("--duration", type=float, default=4.0)
    parser.add_argument("--output_dir", type=str, default="real_audio_results")
    parser.add_argument("--max_steps", type=int, default=REAL_AUDIO_MAX_STEPS)
    parser.add_argument("--stft_weight", type=float, default=REAL_AUDIO_STFT_WEIGHT)
    parser.add_argument(
        "--fr_invariant_weight",
        type=float,
        default=None,
        help=f"Damping double-coset prior weight (default config: {fr_invariant_weight})",
    )
    parser.add_argument(
        "--fr_invariant_coupling",
        type=float,
        default=None,
        help=f"Coupling skew singular-value invariant weight (default config: {fr_invariant_coupling})",
    )
    parser.add_argument(
        "--fr_invariant_speed",
        type=float,
        default=None,
        help=f"Speed profile Fisher-Rao invariant weight (default config: {fr_invariant_speed})",
    )
    parser.add_argument(
        "--fr_invariant_inharm",
        type=float,
        default=None,
        help=f"Inharmonicity profile Fisher-Rao invariant weight (default config: {fr_invariant_inharm})",
    )
    parser.add_argument(
        "--fr_invariant_modal",
        type=float,
        default=None,
        help=f"Piptrack modal amp pair invariant weight (default config: {fr_invariant_modal})",
    )
    parser.add_argument(
        "--fr_mode_weight",
        type=float,
        default=None,
        help=f"Simplex Fisher-Rao modal loss weight (default real-audio: {REAL_AUDIO_FR_MODE_WEIGHT})",
    )
    parser.add_argument(
        "--fr_augment_priors",
        action="store_true",
        help="Keep MSE prior terms alongside invariants (disable Phase 4 replacement)",
    )
    parser.add_argument("--no_stft", action="store_true")
    parser.add_argument("--streaming", action="store_true")
    parser.add_argument("--gpu_stft", action="store_true", help="Enable GPU streaming STFT path")
    parser.add_argument("--stream_demo", action="store_true")
    parser.add_argument("--seed", type=int, default=0)

    # Jump testing
    parser.add_argument("--jump_test", action="store_true", help="Enable jump-test mode")
    parser.add_argument("--force_jumps", action="store_true", help="Force periodic jumps")
    parser.add_argument("--low_patience", type=int, default=None, help="Override stagnation patience")
    parser.add_argument("--min_step_for_jump", type=int, default=None)
    parser.add_argument("--force_jump_every", type=int, default=None, help="Force jump every N steps")
    parser.add_argument("--artificial_plateau", action="store_true", help="Freeze LR to simulate plateau")
    parser.add_argument("--artificial_plateau_at", type=int, default=None)
    parser.add_argument("--jump_pop_size", type=int, default=None, help="Smaller pop for jump tests")
    parser.add_argument("--jump_rollout_horizon", type=int, default=None)

    # Batch analysis
    parser.add_argument("--analyze_batch", action="store_true", help="Cluster notes + key detection")
    parser.add_argument("--n_clusters", type=int, default=None)

    # Live mic
    parser.add_argument("--fit_after", action="store_true", help="Fit model after live recording")

    args = parser.parse_args()
    stft_weight = 0.0 if args.no_stft else args.stft_weight
    jump_test = build_jump_test_config(args)
    output_dir = Path(args.output_dir)

    if args.stft_benchmark:
        run_stft_benchmark()
        return

    if args.live_mic:
        run_live_mic(args.duration or LIVE_MIC_DEFAULT_SECONDS, REAL_AUDIO_SR, output_dir,
                     args.fit_after, args.max_steps, stft_weight, args.streaming)
        return

    if args.stream_demo:
        if not args.audio:
            parser.error("--stream_demo requires --audio")
        demo_streaming_live(args.audio, args.duration)
        return

    if args.audio_dir:
        run_batch(Path(args.audio_dir), output_dir, args.duration, args.max_steps,
                  stft_weight, args.seed, args.streaming, jump_test, args.gpu_stft,
                  args.analyze_batch, args.n_clusters, args.fr_invariant_weight,
                  args.fr_invariant_coupling, args.fr_invariant_speed, args.fr_invariant_inharm,
                  args.fr_invariant_modal, args.fr_mode_weight, args.fr_augment_priors)
    else:
        process_single_note(Path(args.audio), output_dir, args.duration, args.max_steps,
                            stft_weight, args.seed, args.streaming, jump_test, args.gpu_stft,
                            fr_invariant_weight_override=args.fr_invariant_weight,
                            fr_invariant_coupling_override=args.fr_invariant_coupling,
                            fr_invariant_speed_override=args.fr_invariant_speed,
                            fr_invariant_inharm_override=args.fr_invariant_inharm,
                            fr_invariant_modal_override=args.fr_invariant_modal,
                            fr_mode_weight_override=args.fr_mode_weight,
                            fr_augment_priors=args.fr_augment_priors)


if __name__ == "__main__":
    main()