"""Shared helpers for ESD eval pipeline."""
import os
import re
import string
from typing import List, Optional, Tuple

import numpy as np
import soundfile as sf
import torch
import torchaudio


def load_wav_mono_16k(path: str, target_sr: int = 16000) -> np.ndarray:
    """Load wav as mono float32 numpy at target_sr (default 16k)."""
    wav, sr = sf.read(path, always_2d=False)
    if wav.ndim > 1:
        wav = wav.mean(axis=-1)
    wav = wav.astype(np.float32)
    if sr != target_sr:
        t = torch.from_numpy(wav).unsqueeze(0)
        t = torchaudio.transforms.Resample(orig_freq=sr, new_freq=target_sr)(t)
        wav = t.squeeze(0).numpy().astype(np.float32)
    return wav


def slice_audio(audio: np.ndarray, max_sec: float = 8.0, sr: int = 16000) -> List[np.ndarray]:
    """Same chunking as InnerEmoEval/utils.py:slice_audio (kept compatible)."""
    n = len(audio)
    L = int(max_sec * sr)
    if n <= L:
        return [audio]
    parts = n // L + 1
    return [a for a in np.array_split(audio, parts, axis=0)]


def detect_lang_for_seg_text(txt: str) -> str:
    """Return 'zh' if any CJK char, else 'en'."""
    for ch in txt:
        if "\u4e00" <= ch <= "\u9fff":
            return "zh"
    return "en"


def detect_lang_for_speaker(spk_id: str) -> str:
    """ESD 0001-0010 -> zh, 0011-0020 -> en."""
    try:
        i = int(spk_id)
    except Exception:
        return "en"
    return "zh" if i <= 10 else "en"


def cosine_sim_np(a: np.ndarray, b: np.ndarray, eps: float = 1e-8) -> float:
    a = a.flatten().astype(np.float32)
    b = b.flatten().astype(np.float32)
    na = float(np.linalg.norm(a))
    nb = float(np.linalg.norm(b))
    if na < eps or nb < eps:
        return 0.0
    return float(np.dot(a, b) / (na * nb))


def cosine_sim_matrix(query: np.ndarray, db: np.ndarray, eps: float = 1e-8) -> np.ndarray:
    """query: (D,) or (1,D); db: (N,D). returns (N,)."""
    q = query.reshape(1, -1).astype(np.float32)
    qn = q / (np.linalg.norm(q, axis=1, keepdims=True) + eps)
    dn = db / (np.linalg.norm(db, axis=1, keepdims=True) + eps)
    return (dn @ qn.T).reshape(-1)


def minmax_norm(x: np.ndarray, eps: float = 1e-8) -> np.ndarray:
    mn, mx = float(np.min(x)), float(np.max(x))
    if mx - mn < eps:
        return np.zeros_like(x)
    return (x - mn) / (mx - mn)


# -------- text normalization for alignment / matching --------
_ALLOWED_PUNC_TRANS = str.maketrans({
    "，": ",", "。": ".", "；": ";", "：": ":", "？": "?", "！": "!",
    "（": "(", "）": ")", "【": "[", "】": "]",
    "“": '"', "”": '"', "‘": "'", "’": "'",
})


def normalize_text_basic(text: str) -> str:
    text = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f\xa0]", "", text)
    text = text.translate(_ALLOWED_PUNC_TRANS)
    return text


def list_files(root: str, exts: Tuple[str, ...] = (".wav",)) -> List[str]:
    out = []
    for dp, _, fns in os.walk(root):
        for fn in fns:
            if fn.lower().endswith(exts):
                out.append(os.path.join(dp, fn))
    out.sort()
    return out