import random
import unicodedata

import torch
from torch import Tensor
from torch.utils.data import Dataset, DataLoader
import torchaudio
from torchaudio.transforms import FrequencyMasking, MelSpectrogram, AmplitudeToDB, Resample, TimeMasking, TimeStretch
from pathlib import Path
import torchaudio.functional as F
import pandas as pd
from typing import Dict, List, TypedDict
from handle.text_normalizer import collapse_spaces, normalize_extended_uyghur_characters
from tokenizer import ASRTokenizer


# 单个样本的数据结构（Dataset.__getitem__ 返回）
class BatchItem(TypedDict):
    mel_spec: Tensor          # [n_mels, time] Mel频谱
    target_ids: Tensor        # [seq_len] 目标文本的token IDs
    target_text: str          # 原始文本
    audio_path: str           # 音频文件路径


# 批量数据的数据结构（collate_fn 返回，DataLoader 输出）
class Batch(TypedDict):
    mel_specs: Tensor         # [batch, n_mels, time] padding后的Mel频谱
    targets: Tensor           # [batch, max_len] padding后的目标IDs
    mel_lengths: Tensor       # [batch] 每个样本的实际Mel长度
    target_lengths: Tensor    # [batch] 每个样本的实际目标长度
    target_texts: List[str]   # [batch] 原始文本列表
    audio_paths: List[str]    # [batch] 音频路径列表

class TsvFormat(TypedDict):
    client_id: str
    path: str
    sentence: str
    up_votes: int
    down_votes: int
    age: str
    gender: str
    locale: str

class CommonVoiceDataset(Dataset[BatchItem]):
    def __init__(
        self,
        tsv_path: Path,
        audio_dir: Path,
        tokenizer: ASRTokenizer,
        sample_rate: int = 16000,
        n_mels: int = 80,
        max_audio_len: int = 480000,  # 30秒 @ 16kHz
        augment: bool = True,  # 是否启用数据增强
        augment_prob: float = 0.5,  # 数据增强的概率
    ) -> None:
        super().__init__()
    
        self.audio_dir = Path(audio_dir)
        self.tokenizer = tokenizer
        self.sample_rate = sample_rate
        self.max_audio_len = max_audio_len
        self.augment = augment
        self.augment_prob = augment_prob

        self.data: pd.DataFrame = pd.read_csv(tsv_path, sep='\t')

        valid_indices = []
        for index, row in self.data.iterrows():
            audio_path: Path = self.audio_dir / row['path']
            if audio_path.exists():
                valid_indices.append(index)

        self.data = self.data.loc[valid_indices].reset_index(drop=True)

        # Mel频谱转换
        self.mel_transform = MelSpectrogram(
            sample_rate=sample_rate,
            n_fft=400,
            win_length=400,
            hop_length=80,
            n_mels=n_mels,
            f_min=0,
            f_max=8000,
            power=2.0
        )
        self.amplitude_to_db = AmplitudeToDB()

        # SpecAugment 转换
        self.time_masking = TimeMasking(time_mask_param=30)  # 遮蔽最多30帧
        self.freq_masking = FrequencyMasking(freq_mask_param=15)  # 遮蔽最多15个频率

    def __len__(self):
        return len(self.data)
    
    def _load_audio(self, audio_path: Path) -> Tensor:
        waveform, sample_rate = torchaudio.load_with_torchcodec(audio_path)

        if sample_rate != self.sample_rate:
            waveform = Resample(sample_rate, self.sample_rate)(waveform)

        if waveform.shape[0] > 1:
            waveform = waveform.mean(dim=0, keepdim=True)

        max_val = waveform.abs().max()
        if max_val > 0:
            waveform = waveform / max_val
        
        if waveform.shape[1] > self.max_audio_len:
            waveform = waveform[:, :self.max_audio_len]

        return waveform
    
    def _extract_features(self, waveform: Tensor) -> Tensor:
        """ 提取 Mel 频谱特征 """
        mel_spec: Tensor = self.mel_transform(waveform)
        log_mel_spec: Tensor = self.amplitude_to_db(mel_spec)
        return log_mel_spec.squeeze(0) # [n_mels, time]
    
    def _augment_waveform(self, waveform: Tensor) -> Tensor:
        if not self.augment or random.random() > self.augment_prob:
            return waveform
        
        # 1. voice Stretch/Compress 
        if random.random() < 0.5:
            waveform = self._voice_stretch_or_compress(waveform=waveform)
        
        if random.random() < 0.3:
            waveform = self._drop_frames(waveform)

        if random.random() < 0.4:
            waveform = self._add_noise(waveform)
        
        return waveform
    
    def _voice_stretch_or_compress(self, waveform: Tensor) -> Tensor:
        speed_factor = random.uniform(0.6, 1.4) # (Speed Change: 0.6x - 1.4x)
        spec = torch.stft(
            waveform.squeeze(0),
            n_fft=400,
            hop_length=160,
            window=torch.hann_window(400).to(waveform.device),
            return_complex=True
        )
        
        # 时间拉伸（不改变音高）
        stretch = TimeStretch(
            hop_length=160,
            n_freq=201,
            fixed_rate=speed_factor
        )
        stretched_spec = stretch(spec)
        
        # 转回波形
        waveform_stretched = torch.istft(
            stretched_spec,
            n_fft=400,
            hop_length=160,
            window=torch.hann_window(400).to(waveform.device)
        ).unsqueeze(0)
    
        return waveform_stretched
    
    def _drop_frames(self, waveform: Tensor) -> Tensor:
        audio_len = waveform.shape[1]
        drop_ratio = random.uniform(0.05, 0.15)
        drop_len = int(audio_len * drop_ratio)

        if audio_len > drop_len:
            start_pos = random.randint(0, audio_len - drop_len)
            # clean
            waveform = torch.cat([waveform[:, :start_pos], waveform[:, start_pos + drop_len:]], dim=1)

        return waveform
    
    def _add_noise(self, waveform: Tensor, snr_db: float = None) -> Tensor:
        if snr_db is None:
            snr_db = random.uniform(15, 25)
        
        signal_power = waveform.norm(p=2)
        
        snr_linear = 10 ** (snr_db / 10)
        noise_power = signal_power / snr_linear
        
        # Generate Gaussian noise
        noise = torch.randn_like(waveform) * noise_power / waveform.shape[1] ** 0.5
        
        noisy_waveform: Tensor = waveform + noise
        
        max_val = noisy_waveform.abs().max()
        if max_val > 0:
            noisy_waveform = noisy_waveform / max_val
        
        return noisy_waveform
    
    def __getitem__(self, index) -> BatchItem:
        row: TsvFormat = self.data.iloc[index]
        audio_path: Path = self.audio_dir / row['path']
        # text: str = unicodedata.normalize('NFC', row['sentence'].strip()) 
        text: str = normalize_extended_uyghur_characters(collapse_spaces(row['sentence'].strip()))

        waveform = self._load_audio(audio_path=audio_path)
        
        # waveform = self._augment_waveform(waveform)

        mel_spec = self._extract_features(waveform=waveform)

        return BatchItem(
            mel_spec=mel_spec,
            target_ids=torch.tensor(self.tokenizer.encode(text), dtype=torch.long),
            target_text=text,
            audio_path=str(audio_path)
        )
    
def collate_fn(items: List[BatchItem]) -> Batch:
    max_mel_len = max(item['mel_spec'].shape[1] for item in items)
    max_target_len = max(len(item['target_ids']) for item in items)

    batch_size = len(items)
    n_mels = items[0]['mel_spec'].shape[0]

    mel_specs = torch.zeros(batch_size, n_mels, max_mel_len)
    targets = torch.zeros(batch_size, max_target_len, dtype=torch.long)
    mel_lengths = torch.zeros(batch_size, dtype=torch.long)
    target_lengths = torch.zeros(batch_size, dtype=torch.long)

    target_texts = []
    audio_paths = []
    

    for i, item in enumerate(items):
        mel_len = item['mel_spec'].shape[1]
        target_len = len(item['target_ids'])

        mel_specs[i, :, :mel_len] = item['mel_spec']
        targets[i, :target_len] = item['target_ids']
        mel_lengths[i] = mel_len
        target_lengths[i] = target_len

        target_texts.append(item['target_text'])
        audio_paths.append(item['audio_path'])

    return Batch(
        mel_specs=mel_specs,
        targets=targets,
        mel_lengths=mel_lengths,
        target_lengths=target_lengths,
        target_texts=target_texts,
        audio_paths=audio_paths
    )


def create_dataloader(tsv_path: Path, audio_dir: Path, tokenizer: ASRTokenizer, batch_size: int = 8, shuffle: bool = True, augment: bool = True) -> DataLoader:
    dataset = CommonVoiceDataset(tsv_path=tsv_path, audio_dir=audio_dir, tokenizer=tokenizer, augment=augment)

    return DataLoader(dataset=dataset, batch_size=batch_size, shuffle=shuffle, collate_fn=collate_fn, pin_memory=True, num_workers=8, prefetch_factor=8, persistent_workers=True)


# ============ 测试代码 ============
if __name__ == "__main__":
    from pathlib import Path
    
    workspace_dir = Path(__file__).parent.parent
    
    # 初始化tokenizer
    tokenizer = ASRTokenizer(workspace_dir / 'config' / 'asr_vocab.json')
    
    # 创建数据加载器
    dataloader = create_dataloader(
        tsv_path=workspace_dir / 'data'  / 'ug' / 'train.tsv',
        audio_dir=workspace_dir / 'data'  / 'ug' / 'clips',
        tokenizer=tokenizer,
        batch_size=2,
        shuffle=True,
    )
    
    # 测试加载一个batch
    print("测试数据加载:")
    for batch in dataloader:
        batch: Batch
        print(f"Mel specs shape: {batch['mel_specs'].shape}")
        print(f"Targets shape: {batch['targets'].shape}")
        print(f"Mel lengths: {batch['mel_lengths']}")
        print(f"Target lengths: {batch['target_lengths']}")
        print(f"Target texts: {batch['target_texts']}")
        print(f"Audio paths: {batch['audio_paths']}")
        break