import random

import librosa
import torch
import torchaudio
from torch import Tensor, no_grad, device
from torchaudio.transforms import MelSpectrogram, AmplitudeToDB, Resample, TimeStretch
from pathlib import Path
import torchaudio.functional as F

from tokenizer import ASRTokenizer
from model import ASRModel

CONFIG = {
    # 模型配置
    'input_dim': 640,
    'num_heads': 8,
    'ffn_dim': 2048,
    'num_layers': 8,
    'dropout': 0.1,
}

class ASRInference:
    def __init__(self, model_path: Path, vocab_path: Path, device: device, augment: bool = True, augment_prob: float = 0.5) -> None:
        self.device = device
        self.augment: bool = augment
        self.augment_prob: float = augment_prob
        self.tokenizer = ASRTokenizer(vocab_path=vocab_path)
        self.model = ASRModel(
            vocab_size=self.tokenizer.vocab_size(),
            input_dim=CONFIG['input_dim'],
            num_heads=CONFIG['num_heads'],
            ffn_dim=CONFIG['ffn_dim'],
            num_layers=CONFIG['num_layers'],
            dropout=CONFIG['dropout'],
        ).to(device)

        self.model.load_state_dict(torch.load(model_path, map_location=device)['model_state_dict'])
        self.model.eval()

        print(f"params params: {self.model.get_num_params():,}",)

        self.sample_rate = 16000
        self.mel_transform = MelSpectrogram(
            sample_rate=self.sample_rate,
            n_fft=400,
            win_length=400,
            hop_length=160,
            n_mels=80,
            f_min=0,
            f_max=8000,
            power=2.0,
        )
        self.amplitude_to_db = AmplitudeToDB()

    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)
        
        waveform = waveform / (waveform.abs().max() + 1e-8)
        return waveform
    
    def _extract_features(self, waveform: Tensor) -> Tensor:
        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) -> Tensor:
        snr_db = random.uniform(10, 20)
        signal_power = torch.mean(waveform ** 2)
        
        snr_linear = 10 ** (snr_db / 10)
        noise_power = signal_power / snr_linear
        
        noise = torch.randn_like(waveform) * torch.sqrt(noise_power)
        return waveform + noise
    
    def transcribe(self, audio_path: Path) -> str:
        waveform = self._load_audio(audio_path=audio_path)
        waveform = self._augment_waveform(waveform=waveform)
        mel_spec = self._extract_features(waveform=waveform)

        mel_spec = mel_spec.unsqueeze(0).to(self.device) # [1, n_mels, time]
        mel_length = torch.tensor([mel_spec.shape[2]], dtype=torch.long, device=self.device)

        with no_grad():
            log_probs, _ = self.model(mel_specs=mel_spec, mel_lengths=mel_length) # [1, T, vocab]

        text = self.tokenizer.ctc_greedy_decode(log_probs=log_probs[0])
        return text
        
    def transcribe_batch(self, audio_paths: list[Path]) -> list[str]:
        results = []

        for audio_path in audio_paths:
            text = self.transcribe(audio_path=audio_path)
            results.append(text)

        return results


def main():
    workspace_dir = Path(__file__).parent.parent
    device = torch.device('cuda:0')

    # checkpoint = sorted(workspace_dir.glob('.checkpoints/checkpoint_epoch_*.pt'), key=lambda p: int(p.stem.split('_')[1]))[-1]
    checkpoint = workspace_dir / ".checkpoints/checkpoint_step_9500.pt"
    print(f"Load Checkpoint: {checkpoint}")

    inference = ASRInference(model_path=checkpoint, vocab_path=workspace_dir / 'config/uig_vocab.json' , device=device)

    audio_path = workspace_dir / 'data/test/F001_001.wav'
    print(f"\n转录音频: {audio_path}")
    text = inference.transcribe(audio_path=audio_path)
    print(f"\n识别结果: {text}")

if __name__ == "__main__":
    main()
# workspace_dir = Path(__file__).parent.parent
# checkpoint = sorted(workspace_dir.glob('.checkpoints/checkpoint_step_*.pt'), key=lambda p: int(p.stem.split('_')[-1]))[-1]

# print(checkpoint)