#!/usr/bin/env python3
"""
二分类 CNN-GRU 训练脚本
支持 --resume 从 checkpoint 恢复
使用 WeightedRandomSampler + Mixup
"""
import argparse
import os
import json
import random
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
from torchaudio.transforms import MelSpectrogram
import torchaudio.functional as F
import soundfile as sf


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


class AudioDataset(Dataset):
    def __init__(self, root_dir, n_mels=64, n_fft=2048, hop_length=512, sr=16000, augment=False):
        self.samples = []
        self.labels = []
        self.classes = sorted(os.listdir(root_dir))
        self.class_to_idx = {c: i for i, c in enumerate(self.classes)}
        self.sr = sr
        self.augment = augment
        self.mel = MelSpectrogram(sample_rate=sr, n_fft=n_fft, hop_length=hop_length,
                                   n_mels=n_mels, f_min=0, f_max=8000, center=False)
        for cls in self.classes:
            cls_dir = os.path.join(root_dir, cls)
            if not os.path.isdir(cls_dir):
                continue
            for fname in os.listdir(cls_dir):
                if fname.endswith('.wav'):
                    self.samples.append(os.path.join(cls_dir, fname))
                    self.labels.append(self.class_to_idx[cls])

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        path = self.samples[idx]
        waveform, sr = sf.read(path, dtype='float32')
        if waveform.ndim == 1:
            waveform = waveform[np.newaxis, :]
        else:
            waveform = waveform.T  # [channels, samples]
        waveform = torch.from_numpy(waveform)
        if sr != self.sr:
            waveform = F.resample(waveform, sr, self.sr)
        if waveform.shape[0] > 1:
            waveform = waveform.mean(dim=0, keepdim=True)
        # Pad/truncate to 2 seconds
        target_len = self.sr * 2
        if waveform.shape[1] < target_len:
            waveform = torch.nn.functional.pad(waveform, (0, target_len - waveform.shape[1]))
        else:
            waveform = waveform[:, :target_len]

        if self.augment:
            waveform = self.apply_augment(waveform)

        mel = self.mel(waveform)  # [1, n_mels, T]
        mel = (mel + 1e-10).log10() * 10.0
        # Normalize per mel
        mel = (mel - mel.mean(dim=-1, keepdim=True)) / (mel.std(dim=-1, keepdim=True) + 1e-8)
        # Pad/truncate time to 63
        T = mel.shape[-1]
        if T < 63:
            mel = torch.nn.functional.pad(mel, (0, 63 - T))
        else:
            mel = mel[:, :, :63]
        return mel.squeeze(0).transpose(0, 1), self.labels[idx]  # [T, n_mels]

    def apply_augment(self, waveform):
        if random.random() < 0.5:
            gain = random.uniform(0.8, 1.2)
            waveform = waveform * gain
        if random.random() < 0.3:
            noise = torch.randn_like(waveform) * 0.005
            waveform = waveform + noise
        return waveform


class CNNGRU(nn.Module):
    def __init__(self, n_mels=64, num_classes=2):
        super().__init__()
        self.cnn = nn.Sequential(
            nn.Conv1d(n_mels, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool1d(2),
            nn.Conv1d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool1d(2),
            nn.Conv1d(128, 256, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool1d(16)
        )
        self.gru = nn.GRU(256, 128, batch_first=True, bidirectional=True)
        self.fc = nn.Sequential(
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64, num_classes)
        )

    def forward(self, x):
        # x: [B, T, n_mels]
        x = x.transpose(1, 2)  # [B, n_mels, T]
        x = self.cnn(x)  # [B, 256, 16]
        x = x.transpose(1, 2)  # [B, 16, 256]
        x, _ = self.gru(x)
        x = x.mean(dim=1)
        return self.fc(x)


def mixup_data(x, y, alpha=0.4):
    if alpha > 0:
        lam = np.random.beta(alpha, alpha)
    else:
        lam = 1
    batch_size = x.size(0)
    index = torch.randperm(batch_size).to(x.device)
    mixed_x = lam * x + (1 - lam) * x[index]
    y_a, y_b = y, y[index]
    return mixed_x, y_a, y_b, lam


def train_epoch(model, loader, criterion, optimizer, device, use_mixup=True):
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        if use_mixup:
            x, y_a, y_b, lam = mixup_data(x, y)
            outputs = model(x)
            loss = lam * criterion(outputs, y_a) + (1 - lam) * criterion(outputs, y_b)
        else:
            outputs = model(x)
            loss = criterion(outputs, y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total += y.size(0)
        if not use_mixup:
            correct += predicted.eq(y).sum().item()
        else:
            correct += (lam * predicted.eq(y_a).sum().item()
                        + (1 - lam) * predicted.eq(y_b).sum().item())
    return total_loss / len(loader), correct / total


def eval_epoch(model, loader, criterion, device):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            outputs = model(x)
            _, predicted = outputs.max(1)
            total += y.size(0)
            correct += predicted.eq(y).sum().item()
    return correct / total


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', default='dataset/binary/train')
    parser.add_argument('--val_data', default='dataset/binary/val')
    parser.add_argument('--epochs', type=int, default=30)
    parser.add_argument('--batch_size', type=int, default=16)
    parser.add_argument('--lr', type=float, default=1e-3)
    parser.add_argument('--resume', default='')
    parser.add_argument('--output', default='train_output')
    args = parser.parse_args()

    set_seed(42)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    train_ds = AudioDataset(args.data, augment=True)
    val_ds = AudioDataset(args.val_data, augment=False) if os.path.isdir(args.val_data) else None

    # Weighted sampler
    class_counts = np.bincount(train_ds.labels)
    class_weights = 1.0 / class_counts
    sample_weights = [class_weights[l] for l in train_ds.labels]
    sampler = WeightedRandomSampler(sample_weights, len(sample_weights), replacement=True)

    train_loader = DataLoader(train_ds, batch_size=args.batch_size, sampler=sampler)
    val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False) if val_ds else None

    model = CNNGRU(num_classes=len(train_ds.classes)).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=args.lr)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)

    start_epoch = 0
    if args.resume and os.path.isfile(args.resume):
        ckpt = torch.load(args.resume, map_location=device)
        model.load_state_dict(ckpt['model'])
        optimizer.load_state_dict(ckpt['optimizer'])
        start_epoch = ckpt['epoch'] + 1
        print(f"Resumed from epoch {start_epoch}")

    os.makedirs(args.output, exist_ok=True)
    best_acc = 0.0
    for epoch in range(start_epoch, args.epochs):
        train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)
        val_acc = eval_epoch(model, val_loader, criterion, device) if val_loader else 0.0
        scheduler.step()
        print(f"Epoch {epoch+1}/{args.epochs} | Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Acc: {val_acc:.4f}")
        if val_acc > best_acc:
            best_acc = val_acc
            torch.save({
                'epoch': epoch,
                'model': model.state_dict(),
                'optimizer': optimizer.state_dict(),
                'acc': val_acc
            }, os.path.join(args.output, 'best_model.pth'))

    # Save label map
    label_map = {str(i): c for i, c in enumerate(train_ds.classes)}
    with open(os.path.join(args.output, 'label_map.json'), 'w') as f:
        json.dump(label_map, f, indent=2)
    print(f"Training complete. Best val acc: {best_acc:.4f}")


if __name__ == '__main__':
    import torchaudio
    main()
