modules/voiceprint.py

"""
PregoPal - 声纹识别模块
========================
家庭成员声纹注册、识别、管理。

当前方案：频谱特征相似度（轻量 baseline）
后续升级：Whisper-medium encoder embedding（MiniCPM-o 内置）
"""

import json
import numpy as np
import datetime
import hashlib
from pathlib import Path
from config import VOICE_DIR, FAMILY_FILE, VOICEPRINT_SIMILARITY_THRESHOLD


class VoiceprintManager:
    """声纹识别与家庭成员管理"""

    def __init__(self):
        self.family_file = FAMILY_FILE
        self.family_data = self._load_family()

    def _load_family(self):
        """加载家庭成员数据"""
        if self.family_file.exists():
            with open(self.family_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        return {"members": [], "voiceprints": {}}

    def _save_family(self):
        """保存家庭成员数据"""
        with open(self.family_file, 'w', encoding='utf-8') as f:
            json.dump(self.family_data, f, ensure_ascii=False, indent=2)

    def register_member(self, name: str, relation: str, audio_path_input):
        """注册家庭成员声纹"""
        if audio_path_input is None:
            return None, "请先录制或上传语音"

        # 生成声纹ID
        member_id = hashlib.md5(f"{name}_{datetime.datetime.now()}".encode()).hexdigest()[:8]

        # 保存音频文件到数据目录
        audio_path = VOICE_DIR / f"{member_id}.wav"
        try:
            import shutil
            shutil.copy2(audio_path_input, audio_path)
        except Exception as e:
            return None, f"保存音频失败: {str(e)}"

        # 提取简单的声纹特征（使用音频的频谱特征作为简化方案）
        features = self._extract_features_from_file(audio_path)

        # 保存成员信息
        member_info = {
            "id": member_id,
            "name": name,
            "relation": relation,
            "registered_at": datetime.datetime.now().isoformat(),
            "audio_path": str(audio_path),
            "features": features
        }

        self.family_data["members"].append(member_info)
        self.family_data["voiceprints"][member_id] = features
        self._save_family()

        return member_info, f"✅ 成功注册 {relation} - {name}！"

    def _extract_features_from_file(self, audio_path):
        """
        从音频文件提取声纹特征
        当前：频谱统计特征（轻量 baseline）
        后续：Whisper-medium encoder embedding（更鲁棒）
        """
        try:
            import soundfile as sf
            data, samplerate = sf.read(audio_path)
            if len(data.shape) > 1:
                data = data.mean(axis=1)

            # 提取简单特征：MFCC-like 统计量
            features = {
                "mean": float(np.mean(data)),
                "std": float(np.std(data)),
                "max": float(np.max(data)),
                "min": float(np.min(data)),
                "zero_crossing_rate": float(np.sum(np.abs(np.diff(np.sign(data)))) / len(data)),
                "energy": float(np.sum(data ** 2) / len(data)),
                "duration": float(len(data) / samplerate)
            }
            return features
        except Exception:
            # 如果无法解析音频，返回占位特征
            return {
                "mean": 0.0, "std": 0.0, "max": 0.0, "min": 0.0,
                "zero_crossing_rate": 0.0, "energy": 0.0, "duration": 0.0
            }

    def identify_speaker(self, audio_path_input):
        """识别说话人身份"""
        if audio_path_input is None:
            return None, "请先录制或上传语音"

        if not self.family_data["members"]:
            return None, "⚠️ 尚未注册任何家庭成员，请先注册"

        features = self._extract_features_from_file(audio_path_input)

        # 计算与每个注册成员的相似度
        best_match = None
        best_score = -1

        for member in self.family_data["members"]:
            stored = member["features"]
            score = self._compute_similarity(features, stored)
            if score > best_score:
                best_score = score
                best_match = member

        # 相似度阈值判断
        threshold = VOICEPRINT_SIMILARITY_THRESHOLD
        if best_score >= threshold:
            return best_match, f"🔊 识别结果: {best_match['relation']} - {best_match['name']} (置信度: {best_score:.2f})"
        else:
            return None, f"🔊 未能识别说话人 (最高匹配: {best_score:.2f}，需≥{threshold})"

    def _compute_similarity(self, f1, f2):
        """计算两个声纹特征的余弦相似度"""
        keys = ["mean", "std", "max", "min", "zero_crossing_rate", "energy"]
        v1 = np.array([f1.get(k, 0) for k in keys])
        v2 = np.array([f2.get(k, 0) for k in keys])

        # 归一化
        norm1 = np.linalg.norm(v1)
        norm2 = np.linalg.norm(v2)
        if norm1 == 0 or norm2 == 0:
            return 0

        return float(np.dot(v1, v2) / (norm1 * norm2))

    def get_members_list(self):
        """获取家庭成员列表"""
        if not self.family_data["members"]:
            return "📋 暂无注册成员"
        lines = ["📋 已注册家庭成员："]
        for m in self.family_data["members"]:
            lines.append(f"  • {m['relation']} - {m['name']} (注册于 {m['registered_at'][:10]})")
        return "\n".join(lines)

    def delete_member(self, member_id: str):
        """删除家庭成员"""
        self.family_data["members"] = [m for m in self.family_data["members"] if m["id"] != member_id]
        self.family_data["voiceprints"].pop(member_id, None)
        self._save_family()
        return f"已删除成员 {member_id}"