# -*- coding: utf-8 -*-

import copy
from typing import List
from typing import Optional

import numpy as np
import torch
from transformers.audio_utils import spectrogram
from transformers.audio_utils import window_function
from transformers.models.whisper.feature_extraction_whisper import WhisperFeatureExtractor


class MiniCPMAAudioProcessor(WhisperFeatureExtractor):
    """
    在 WhisperFeatureExtractor 基础上：
    - 支持 dynamic_log_norm（原版 max-8dB，可调 dynamic_range_db）
    - 或固定对数下限 log_floor_db（例如 -10dB）
        - 这是由于我们需要做 streaming 方案，在这个方案下，没办法进行动态设置
        - 这个 可以在中途修改，通过 set_dynamic_log_norm
    两条路径（torch / numpy）保持一致的裁剪与缩放顺序：
        log10 -> (动态/固定下限裁剪) -> (+4)/4
    """

    def __init__(
        self,
        *args,
        dynamic_log_norm: bool = True,
        dynamic_range_db: float = 8.0,
        log_floor_db: float = -10.0,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.dynamic_log_norm = bool(dynamic_log_norm)
        self.dynamic_range_db = float(dynamic_range_db)
        self.log_floor_db = float(log_floor_db)

    def set_spac_log_norm(
        self,
        dynamic_range_db: Optional[float] = None,
        log_floor_db: Optional[float] = None,
        *,
        inplace: bool = True,
    ) -> "MiniCPMAAudioProcessor":
        """
        热更新动态/固定下限策略。

        Args:
            enabled: True=使用动态阈值(max - dynamic_range_db)，False=使用固定下限 log_floor_db。
                    传 None 表示保持不变。
            dynamic_range_db: 动态范围（dB），仅在 enabled=True 时生效。None 表示保持不变。
            log_floor_db: 固定对数下限（dB，通常<=0），仅在 enabled=False 时使用。None 表示保持不变。
            inplace: True 直接修改当前实例；False 返回一个浅拷贝并在其上修改。

        Returns:
            self 或新的实例（当 inplace=False 时）。
        """

        target = self if inplace else copy.copy(self)

        if dynamic_range_db is not None:
            val = float(dynamic_range_db)
            if val < 0:
                raise ValueError("dynamic_range_db must be >= 0.")
            target.dynamic_log_norm = True  # 给出显式提示：设置了该值就走动态模式
            target.dynamic_range_db = val

        if log_floor_db is not None:
            val = float(log_floor_db)
            # 一般 log10(mel) 最大不超过 ~0dB，floor 设为<=0更合理；这里做宽松校验
            if val > 0:
                raise ValueError("log_floor_db should be <= 0 (log10 scale).")
            target.dynamic_log_norm = False  # 设置了该值就走固定下限模式
            target.log_floor_db = val

        return target

    # ------------ NumPy 实现：覆写 ----------
    def _np_extract_fbank_features(self, waveform_batch: np.ndarray, device: str) -> np.ndarray:
        """
        与上游一致的 NumPy 版本，但将 max-8dB 替换为可配置的 动态/固定 下限裁剪。
        """
        if device != "cpu":
            raise ValueError(
                f"Got device `{device}` for feature extraction, but feature extraction on CUDA accelerator "
                "devices requires torch. Set device='cpu' or install torch."
            )

        log_spec_batch: List[np.ndarray] = []
        for waveform in waveform_batch:
            # 生成 log10 Mel
            log_spec = spectrogram(
                waveform,
                window_function(self.n_fft, "hann"),
                frame_length=self.n_fft,
                hop_length=self.hop_length,
                power=2.0,
                dither=self.dither,
                mel_filters=self.mel_filters,
                log_mel="log10",
            )
            # 与上游保持一致：去掉最后一帧
            log_spec = log_spec[:, :-1]

            # 动态/固定裁剪
            if self.dynamic_log_norm:
                threshold = log_spec.max() - self.dynamic_range_db
                log_spec = np.maximum(log_spec, threshold)
            else:
                log_spec = np.maximum(log_spec, self.log_floor_db)

            # 与 Whisper 一致的线性缩放
            log_spec = (log_spec + 4.0) / 4.0

            log_spec_batch.append(log_spec)

        return np.array(log_spec_batch)

    # ------------ PyTorch 实现：覆写 ----------
    def _torch_extract_fbank_features(self, waveform: np.ndarray, device: str = "cpu") -> np.ndarray:
        if torch is None:
            raise RuntimeError("PyTorch 未安装，无法在 GPU 上计算 STFT。")

        waveform = torch.from_numpy(waveform).to(device, torch.float32)
        window = torch.hann_window(self.n_fft, device=device)

        if self.dither != 0.0:
            waveform = waveform + self.dither * torch.randn_like(waveform)

        stft = torch.stft(waveform, n_fft=self.n_fft, hop_length=self.hop_length, window=window, return_complex=True)
        magnitudes = stft[..., :-1].abs() ** 2

        mel_filters = torch.from_numpy(self.mel_filters).to(device, torch.float32)  # [n_mels, 1+n_fft//2]
        mel_spec = mel_filters.T @ magnitudes  # [..., n_mels, T]

        log_spec = torch.clamp(mel_spec, min=1e-10).log10()  # ≤ 0

        if self.dynamic_log_norm:
            if waveform.dim() == 2:
                max_val_t = log_spec.max(dim=2, keepdim=True)[0]  # over T
                max_val_bt = max_val_t.max(dim=1, keepdim=True)[0]  # over mel
                threshold = max_val_bt - self.dynamic_range_db
                log_spec = torch.maximum(log_spec, threshold)
            else:
                threshold = log_spec.max() - self.dynamic_range_db
                log_spec = torch.maximum(log_spec, threshold)
        else:
            floor_tensor = torch.tensor(self.log_floor_db, dtype=log_spec.dtype, device=log_spec.device)
            log_spec = torch.maximum(log_spec, floor_tensor)

        log_spec = (log_spec + 4.0) / 4.0

        if device != "cpu":
            log_spec = log_spec.detach().cpu()
        return log_spec.numpy()

    def process(self, *args, **kwargs):
        """Alias of __call__ for convenience."""
        return self.__call__(*args, **kwargs)