scripts/update_doa_analysis.py

"""
Compute TDOA-based DoA for original 4-channel audio files
and update results.json with actual spatial estimates
"""

import json
from pathlib import Path
import soundfile as sf

from .doa_tdoa import compute_tdoa_based_doa_for_mixture

def analyze_4ch_doa(input_4ch_file, output_dir):
    """
    Analyze a 4-channel audio file for Direction of Arrival

    Args:
        input_4ch_file: Path to 4-channel WAV
        output_dir: Output directory where results.json exists
    """
    # Load 4-channel audio (downsample for faster TDOA computation)
    audio_4ch, sr = sf.read(str(input_4ch_file))

    # Downsample to 16 kHz for faster processing while maintaining accuracy
    downsample_factor = sr // 16000
    if downsample_factor > 1:
        audio_4ch = audio_4ch[::downsample_factor]
        sr_analysis = 16000
    else:
        sr_analysis = sr

    print(f"\nAnalyzing {input_4ch_file}")
    print(f"  Original shape: {audio_4ch.shape}")
    print(f"  Analysis sample rate: {sr_analysis} Hz (downsampled by {downsample_factor}x)")

    # Compute TDOA-based DoA
    print("\nComputing TDOA-based Direction of Arrival...")
    doa_result = compute_tdoa_based_doa_for_mixture(audio_4ch, sr_analysis)

    print(f"  Left-Right angle: {doa_result['lr_angle']:.1f}° (xcorr: {doa_result['lr_xcorr']:.3f})")
    print(f"  Front-Rear angle: {doa_result['fb_angle']:.1f}° (xcorr: {doa_result['fb_xcorr']:.3f})")

    # Update results.json with DoA information
    results_file = Path(output_dir) / 'results.json'

    if results_file.exists():
        with open(results_file) as f:
            results = json.load(f)

        # Add TDOA DoA analysis
        results['tdoa_analysis'] = {
            'input_4ch_file': str(input_4ch_file),
            'method': 'TDOA (Time Difference of Arrival) cross-correlation',
            'microphone_pairs': {
                'front_pair_lr': {
                    'channels': [0, 2],
                    'labels': ['LF', 'RF'],
                    'lr_angle_degrees': float(doa_result['lr_angle']),
                    'xcorr_confidence': float(doa_result['lr_xcorr']),
                    'interpretation': 'Negative = left, Positive = right'
                },
                'front_rear_pair': {
                    'channels': [0, 1],
                    'labels': ['LF', 'LR'],
                    'fb_angle_degrees': float(doa_result['fb_angle']),
                    'xcorr_confidence': float(doa_result['fb_xcorr']),
                    'interpretation': 'Negative = front, Positive = rear'
                }
            },
            'compass_angles': {
                '0_degrees': 'Front (center)',
                '90_degrees': 'Right',
                '180_degrees': 'Rear',
                '270_degrees': 'Left',
                'note': 'Sound source direction relative to listener'
            }
        }

        # Update direction_of_arrival field with actual result
        results['processing_methods']['direction_of_arrival'] = 'TDOA cross-correlation (4-channel analysis)'

        # Save updated results
        with open(results_file, 'w') as f:
            json.dump(results, f, indent=2)

        print(f"\n✓ Updated {results_file}")
        print("  Added TDOA analysis with spatial estimates")

        return doa_result
    else:
        print(f"✗ {results_file} not found")
        return None


if __name__ == '__main__':
    # Analyze both files
    files = [
        ('data/example_mixture.wav', 'output/analysis_example_frankenstein'),
        ('data/mixture.wav', 'output/analysis_frankenstein'),
    ]

    print("=" * 60)
    print("TDOA-Based Direction of Arrival Analysis")
    print("=" * 60)

    for input_file, output_dir in files:
        analyze_4ch_doa(input_file, output_dir)

    print("\n" + "=" * 60)
    print("✓ TDOA analysis complete")
    print("=" * 60)