# coding: utf-8
"""GCRN 依赖极简的音频预处理：16kHz 单声道 WAV <-> STFT [1,2,401,161]。

与 GCRN.AXERA 的 board_inference/audio.py 保持一致（hamming sym 窗 + numpy rfft）。
"""
from __future__ import annotations

import wave
from pathlib import Path

import numpy as np

SAMPLE_RATE = 16000
CHUNK_SAMPLES = 4 * SAMPLE_RATE
N_FFT = 320
HOP_LENGTH = 160
WIN_LENGTH = 320


def read_wav(path: str | Path) -> tuple[np.ndarray, int]:
    with wave.open(str(path), "rb") as source:
        if source.getnchannels() != 1 or source.getsampwidth() != 2:
            raise ValueError(f"expected mono 16-bit PCM WAV: {path}")
        sample_rate = source.getframerate()
        samples = np.frombuffer(
            source.readframes(source.getnframes()), dtype=np.int16
        ).copy()
    return samples, sample_rate


def write_wav(path: str | Path, samples: np.ndarray, sample_rate: int) -> None:
    values = np.asarray(samples, dtype=np.float32).reshape(-1)
    pcm = np.clip(values, -1.0, 1.0)
    pcm = np.rint(pcm * 32767.0).astype(np.int16)
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with wave.open(str(path), "wb") as output:
        output.setnchannels(1)
        output.setsampwidth(2)
        output.setframerate(sample_rate)
        output.writeframes(pcm.tobytes())


def _symmetric_hamming(size: int) -> np.ndarray:
    index = np.arange(size, dtype=np.float32)
    return (0.54 - 0.46 * np.cos(2.0 * np.pi * index / (size - 1))).astype(
        np.float32
    )


def chunks(samples: np.ndarray) -> list[tuple[np.ndarray, int]]:
    """int16 PCM -> 4 秒 float32 chunk 序列 [(chunk, valid_samples)]。"""
    values = np.asarray(samples, dtype=np.float32).reshape(-1) / 32768.0
    if values.size == 0:
        raise ValueError("input WAV is empty")
    result = []
    for start in range(0, values.size, CHUNK_SAMPLES):
        end = min(start + CHUNK_SAMPLES, values.size)
        valid = end - start
        chunk = np.zeros(CHUNK_SAMPLES, dtype=np.float32)
        chunk[:valid] = values[start:end]
        result.append((chunk, valid))
    return result


def stft(chunk: np.ndarray) -> np.ndarray:
    """4 秒波形 -> [1, 2, 401, 161]（batch, real/imag, time, freq）。"""
    signal = np.asarray(chunk, dtype=np.float32).reshape(-1)
    if signal.size != CHUNK_SAMPLES:
        raise ValueError(f"expected {CHUNK_SAMPLES} samples, got {signal.size}")
    pad = N_FFT // 2
    signal = np.pad(signal, (pad, pad), mode="constant")
    frame_count = 1 + (signal.size - WIN_LENGTH) // HOP_LENGTH
    frames = np.lib.stride_tricks.sliding_window_view(signal, WIN_LENGTH)[
        ::HOP_LENGTH
    ][:frame_count]
    spectrum = np.fft.rfft(frames * _symmetric_hamming(WIN_LENGTH)[None, :], axis=1)
    value = np.stack((spectrum.real, spectrum.imag), axis=0)
    return value.astype(np.float32)[None, ...]


def istft(value: np.ndarray) -> np.ndarray:
    """[1,2,401,161] 增强谱 -> 64000 个 float32 样本。"""
    output = np.asarray(value, dtype=np.float32)
    if output.shape != (1, 2, 401, 161):
        raise ValueError(f"unexpected GCRN output shape: {output.shape}")
    value = output[0]
    spectrum = value[0].astype(np.float64) + 1j * value[1].astype(np.float64)
    frames = np.fft.irfft(spectrum, n=N_FFT, axis=1)
    window = _symmetric_hamming(WIN_LENGTH).astype(np.float64)
    frames *= window[None, :]

    output_length = N_FFT + HOP_LENGTH * (frames.shape[0] - 1)
    waveform = np.zeros(output_length, dtype=np.float64)
    window_sum = np.zeros(output_length, dtype=np.float64)
    for index, frame in enumerate(frames):
        start = index * HOP_LENGTH
        waveform[start : start + N_FFT] += frame
        window_sum[start : start + N_FFT] += window * window
    valid = window_sum > np.finfo(np.float64).eps
    waveform[valid] /= window_sum[valid]
    pad = N_FFT // 2
    return waveform[pad:-pad].astype(np.float32)