"""Chromagram 特征提取。

流程：
1. 音频解码：ffmpeg pipe 输出 22050Hz / Mono / f32le，直接读入内存，无临时文件
2. librosa.feature.chroma_cens() 提取 12×T Chromagram（CENS，对速度/音色鲁棒）
3. 主音对齐：将能量最大的音级滚至第 0 行，实现转调不变性
4. scipy.signal.resample(chroma, 128, axis=1) 时间归一化到 12×128
5. .flatten() 展开为 1536 维向量

子序列 DTW 路径（extract_chroma_fixed_fps_from_samples）跳过第 4、5 步，
改用线性插值到固定帧率，保留音频时长信息供子序列匹配使用。
"""

import logging
import os
import subprocess

import librosa
import numpy as np
from scipy.signal import resample

logger = logging.getLogger(__name__)

# 目标采样率和时间帧数
TARGET_SR = 22050
TARGET_FRAMES = 128
VECTOR_DIM = 12 * TARGET_FRAMES  # 1536


def _load_audio_via_pipe(audio_path: str) -> np.ndarray:
    """使用 ffmpeg pipe 将音频解码为 22050Hz mono float32，不落临时文件到磁盘。"""
    cmd = [
        "ffmpeg", "-y",
        "-i", audio_path,
        "-ar", str(TARGET_SR),
        "-ac", "1",
        "-f", "f32le",
        "pipe:1",
    ]
    result = subprocess.run(cmd, capture_output=True)
    if result.returncode != 0:
        raise RuntimeError(f"ffmpeg 解码失败: {result.stderr.decode(errors='replace')}")
    return np.frombuffer(result.stdout, dtype=np.float32)


def extract_chroma_feature_from_samples(
    samples: np.ndarray,
    sr: int,
    hop_length: int = 512,
    win_len_smooth: int = 41,
) -> np.ndarray:
    """从已加载的音频样本提取 1536 维 Chromagram 特征向量。

    若 sr 不等于 TARGET_SR，先用 librosa.resample 在内存中降采样，
    避免重新走 ffmpeg 流程。

    Args:
        samples: 单声道音频样本（任意采样率）。
        sr: samples 对应的采样率。
        hop_length: CQT hop 大小，增大可成比例降低计算量，不影响最终 128 帧精度。
        win_len_smooth: CENS 平滑窗口帧数，应随 hop_length 等比缩小以保持相同的时间覆盖。

    Returns:
        shape 为 (1536,) 的 numpy 数组。
    """
    y = samples if sr == TARGET_SR else librosa.resample(samples, orig_sr=sr, target_sr=TARGET_SR)

    # 提取 CENS Chromagram (12×T)
    chroma = librosa.feature.chroma_cens(y=y, sr=TARGET_SR, hop_length=hop_length, win_len_smooth=win_len_smooth)

    # 主音对齐
    tonic = int(np.argmax(chroma.sum(axis=1)))
    if tonic != 0:
        chroma = np.roll(chroma, -tonic, axis=0)

    # 时间归一化到 12×128
    if chroma.shape[1] != TARGET_FRAMES:
        chroma = resample(chroma, TARGET_FRAMES, axis=1)

    feature = chroma.flatten().astype(np.float32)
    assert feature.shape == (VECTOR_DIM,), (
        f"特征维度错误: 期望 {VECTOR_DIM}, 实际 {feature.shape}"
    )
    return feature


def extract_chroma_matrix_from_samples(
    samples: np.ndarray,
    sr: int,
    hop_length: int = 512,
    win_len_smooth: int = 41,
) -> np.ndarray:
    """从已加载的音频样本提取 12×128 Chromagram 矩阵（供 DTW 精排使用）。"""
    return extract_chroma_feature_from_samples(samples, sr, hop_length=hop_length, win_len_smooth=win_len_smooth).reshape(12, TARGET_FRAMES)


def extract_chroma_feature(audio_path: str, hop_length: int = 512, win_len_smooth: int = 41) -> np.ndarray:
    """从音频文件提取 1536 维 Chromagram 特征向量。

    Args:
        audio_path: 音频文件路径。
        hop_length: CQT hop 大小。
        win_len_smooth: CENS 平滑窗口帧数。

    Returns:
        shape 为 (1536,) 的 numpy 数组。

    Raises:
        FileNotFoundError: 音频文件不存在。
        RuntimeError: ffmpeg 解码失败。
    """
    if not os.path.isfile(audio_path):
        raise FileNotFoundError(f"音频文件不存在: {audio_path}")

    y = _load_audio_via_pipe(audio_path)
    return extract_chroma_feature_from_samples(y, TARGET_SR, hop_length=hop_length, win_len_smooth=win_len_smooth)


def extract_chroma_matrix(audio_path: str, hop_length: int = 512, win_len_smooth: int = 41) -> np.ndarray:
    """从音频文件提取 12×128 Chromagram 矩阵（未展平，供 DTW 精排使用）。

    Returns:
        shape 为 (12, 128) 的 numpy 数组，已做主音对齐。
    """
    return extract_chroma_feature(audio_path, hop_length=hop_length, win_len_smooth=win_len_smooth).reshape(12, TARGET_FRAMES)


# 子序列 DTW 路径的默认目标帧率（帧/秒）。
# fps=2 → 每 0.5s 一帧；3 分钟歌曲 = 360 帧；副歌片段（40%）= 144 帧。
FIXED_FPS = 2


def load_audio_mono_22050hz(audio_path: str) -> np.ndarray:
    """从音频文件加载 22050Hz 单声道 float32 样本（对外公开接口）。

    直接调用内部的 ffmpeg pipe 解码，无临时文件。

    Raises:
        FileNotFoundError: 音频文件不存在。
        RuntimeError: ffmpeg 解码失败。
    """
    if not os.path.isfile(audio_path):
        raise FileNotFoundError(f"音频文件不存在: {audio_path}")
    return _load_audio_via_pipe(audio_path)


def extract_chroma_fixed_fps_from_samples(
    samples: np.ndarray,
    sr: int,
    target_fps: int = FIXED_FPS,
    hop_length: int = 512,
    win_len_smooth: int = 41,
) -> np.ndarray:
    """提取固定帧率 Chromagram 矩阵，帧数由音频时长决定（供子序列 DTW 使用）。

    与 extract_chroma_feature_from_samples 的区别：
      - 不做固定 128 帧归一化，保留音频绝对时长信息
      - 时间轴插值改用线性插值，避免 FFT 重采样（scipy.signal.resample）的 Gibbs 振铃
      - 返回 12×T 矩阵，T = round(duration_sec * target_fps)
      - 查询片段（如 chorus_only）与全曲均保持同等时间密度（帧/秒），
        子序列 DTW 可在全曲序列中找到最佳对齐窗口

    Args:
        samples: 单声道音频样本（任意采样率）。
        sr: samples 对应的采样率。
        target_fps: 目标帧率（帧/秒），决定最终时间分辨率。
        hop_length: CQT hop 大小。
        win_len_smooth: CENS 平滑窗口帧数。

    Returns:
        shape 为 (12, T) 的 numpy float32 数组，未做主音对齐（由调用方的 12 路位移处理）。
    """
    y = samples if sr == TARGET_SR else librosa.resample(samples, orig_sr=sr, target_sr=TARGET_SR)
    chroma = librosa.feature.chroma_cens(
        y=y, sr=TARGET_SR, hop_length=hop_length, win_len_smooth=win_len_smooth
    )

    # 不做主音对齐：固定帧率路径依赖 _best_shifted_subsequence_dtw 的 12 路遍历
    # 实现转调不变性，argmax 能量估计不稳定，去掉后向量空间更干净

    # 目标帧数由音频时长和目标帧率共同决定
    duration_sec = len(y) / TARGET_SR
    target_t = max(1, round(duration_sec * target_fps))

    # 线性插值到目标帧数
    if chroma.shape[1] != target_t:
        x_old = np.linspace(0, 1, chroma.shape[1])
        x_new = np.linspace(0, 1, target_t)
        chroma = np.array(
            [np.interp(x_new, x_old, chroma[p]) for p in range(12)],
            dtype=np.float32,
        )

    return chroma.astype(np.float32)