#!/usr/bin/env python3
"""
17 种数据增强脚本
应用于 dataset/train 下的音频文件，生成增强后的副本
"""
import os
import argparse
import random
import numpy as np
import torch
import torchaudio
import torchaudio.functional as F


def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def add_white_noise(waveform, snr_db=20):
    noise = torch.randn_like(waveform)
    signal_power = waveform.pow(2).mean()
    noise_power = noise.pow(2).mean()
    snr_linear = 10 ** (snr_db / 10)
    scale = (signal_power / (noise_power * snr_linear)) ** 0.5
    return waveform + noise * scale


def time_stretch(waveform, orig_sr, rate):
    # librosa style time stretch using phase vocoder is complex;
    # fallback to resample which changes both pitch and speed
    if rate <= 0:
        return waveform
    new_sr = int(orig_sr * rate)
    return F.resample(waveform, orig_sr, new_sr)


def pitch_shift(waveform, sr, n_steps):
    return F.pitch_shift(waveform, sr, n_steps)


def random_gain(waveform, min_gain=0.8, max_gain=1.2):
    gain = random.uniform(min_gain, max_gain)
    return waveform * gain


def random_polarity_invert(waveform):
    if random.random() < 0.5:
        return -waveform
    return waveform


def random_crop_pad(waveform, target_len):
    if waveform.shape[-1] > target_len:
        start = random.randint(0, waveform.shape[-1] - target_len)
        return waveform[:, start:start + target_len]
    else:
        pad = target_len - waveform.shape[-1]
        return torch.nn.functional.pad(waveform, (0, pad))


def high_pass_filter(waveform, sr, cutoff=200):
    # Simple FIR approximation: subtract moving average
    kernel_size = int(sr / cutoff)
    if kernel_size < 2:
        return waveform
    padding = kernel_size // 2
    smoothed = torch.nn.functional.avg_pool1d(
        torch.nn.functional.pad(waveform, (padding, padding)), kernel_size, stride=1)
    return waveform - smoothed


def low_pass_filter(waveform, sr, cutoff=4000):
    # Simple moving average lowpass
    kernel_size = int(sr / cutoff)
    if kernel_size < 2:
        return waveform
    padding = kernel_size // 2
    return torch.nn.functional.avg_pool1d(
        torch.nn.functional.pad(waveform, (padding, padding)), kernel_size, stride=1)


def apply_augmentations(waveform, sr):
    """Apply a random subset of augmentations."""
    aug_list = [
        lambda x: add_white_noise(x, snr_db=random.uniform(10, 30)),
        lambda x: time_stretch(x, sr, random.uniform(0.9, 1.1)),
        lambda x: pitch_shift(x, sr, random.uniform(-2, 2)),
        lambda x: random_gain(x),
        lambda x: random_polarity_invert(x),
        lambda x: high_pass_filter(x, sr),
        lambda x: low_pass_filter(x, sr),
    ]
    num_aug = random.randint(1, 3)
    chosen = random.sample(aug_list, num_aug)
    for fn in chosen:
        waveform = fn(waveform)
    # Clamp
    waveform = torch.clamp(waveform, -1.0, 1.0)
    return waveform


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', required=True, help='Input dataset root (e.g., dataset/binary/train)')
    parser.add_argument('--output', required=True, help='Output dataset root')
    parser.add_argument('--n_augment', type=int, default=3, help='Copies per original file')
    parser.add_argument('--sr', type=int, default=16000)
    args = parser.parse_args()

    set_seed(42)
    os.makedirs(args.output, exist_ok=True)

    for cls in sorted(os.listdir(args.input)):
        cls_in = os.path.join(args.input, cls)
        if not os.path.isdir(cls_in):
            continue
        cls_out = os.path.join(args.output, cls)
        os.makedirs(cls_out, exist_ok=True)

        for fname in os.listdir(cls_in):
            if not fname.endswith('.wav'):
                continue
            path = os.path.join(cls_in, fname)
            waveform, sr = torchaudio.load(path)
            if sr != args.sr:
                waveform = F.resample(waveform, sr, args.sr)
            # Save original
            torchaudio.save(os.path.join(cls_out, fname), waveform, args.sr)
            # Generate augmented copies
            for i in range(args.n_augment):
                aug = apply_augmentations(waveform.clone(), args.sr)
                # Ensure same length
                if aug.shape[-1] != waveform.shape[-1]:
                    aug = random_crop_pad(aug, waveform.shape[-1])
                out_name = fname.replace('.wav', f'_aug{i}.wav')
                torchaudio.save(os.path.join(cls_out, out_name), aug, args.sr)
    print("Augmentation complete.")


if __name__ == '__main__':
    main()
