#!/usr/bin/env python3
"""
4 通道模拟信号生成
根据指定方位角和距离生成十字阵列 WAV 文件
"""
import argparse
import numpy as np
import soundfile as sf


def generate_gunshot_signal(sr=16000, duration=2.0, freq=1000, decay=0.05):
    t = np.linspace(0, duration, int(sr * duration))
    # Impulse-like exponential decay sine
    envelope = np.exp(-t / decay)
    signal = np.sin(2 * np.pi * freq * t) * envelope
    # Normalize
    signal = signal / np.max(np.abs(signal) + 1e-10)
    return signal.astype(np.float32)


def simulate_array(signal, azimuth_deg, distance_m, sr=16000, spacing=0.15, c=343.0):
    """
    Cross array: mic positions [front, right, back, left]
    Returns 4 channels with appropriate delays and attenuation
    """
    az = np.deg2rad(azimuth_deg)
    # Source position in plane
    src_x = distance_m * np.sin(az)
    src_y = distance_m * np.cos(az)
    mics = np.array([
        [0.0, spacing],
        [spacing, 0.0],
        [0.0, -spacing],
        [-spacing, 0.0]
    ])
    channels = []
    for mic in mics:
        dist = np.sqrt((src_x - mic[0])**2 + (src_y - mic[1])**2)
        delay_sec = dist / c
        delay_samples = int(delay_sec * sr)
        # Amplitude attenuation (1/r)
        att = 1.0 / (dist + 1e-6)
        ch = np.zeros_like(signal)
        if delay_samples < len(signal):
            ch[delay_samples:] = signal[:len(signal) - delay_samples] * att
        channels.append(ch)
    return np.stack(channels, axis=0)  # [4, samples]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--output', default='dataset/multichannel_test.wav')
    parser.add_argument('--azimuth', type=float, default=90.0, help='Azimuth in degrees')
    parser.add_argument('--distance', type=float, default=100.0, help='Distance in meters')
    parser.add_argument('--sr', type=int, default=16000)
    parser.add_argument('--duration', type=float, default=2.0)
    args = parser.parse_args()

    signal = generate_gunshot_signal(args.sr, args.duration)
    multich = simulate_array(signal, args.azimuth, args.distance, args.sr)
    # Normalize to prevent clipping
    max_val = np.max(np.abs(multich))
    if max_val > 0:
        multich = multich / max_val * 0.9

    # soundfile expects [samples, channels] for multi-channel
    sf.write(args.output, multich.T, args.sr, subtype="PCM_16")
    print(f"Generated {args.output}: azimuth={args.azimuth}°, distance={args.distance}m, shape={multich.shape}")


if __name__ == '__main__':
    main()
