"""
FAISS向量匹配器
基于向量相似度从候选集中检索最匹配的菜品名
"""
import json
import os
from typing import List, Tuple

import faiss
import numpy as np


class DishNameMatcher:
    """
    菜品名向量匹配器
    加载预构建的FAISS索引，执行相似度检索
    """
    
    def __init__(self, index_dir: str):
        """
        参数：
            index_dir: 索引目录，应包含：
                - dish_names.json: 菜品名列表
                - dish_embeddings.npy: 归一化后的向量
                - faiss_index.bin: FAISS索引文件
        """
        self.index_dir = index_dir
        
        # 加载菜品名列表
        names_path = os.path.join(index_dir, "dish_names.json")
        with open(names_path, "r", encoding="utf-8") as f:
            self.dish_names: List[str] = json.load(f)
        
        # 加载向量（可选，用于调试）
        embeddings_path = os.path.join(index_dir, "dish_embeddings.npy")
        self.embeddings = np.load(embeddings_path)
        
        # 加载FAISS索引
        index_path = os.path.join(index_dir, "faiss_index.bin")
        self.index = faiss.read_index(index_path)
        
        print(
            f"[DishNameMatcher] Loaded index_dir={index_dir} "
            f"num_dishes={len(self.dish_names)} "
            f"embeddings_shape={self.embeddings.shape} "
            f"index_ntotal={self.index.ntotal}"
        )
    
    def match(
        self,
        query_embedding: np.ndarray,
        top_k: int = 3
    ) -> List[Tuple[str, float]]:
        """
        检索最相似的菜品名
        
        参数：
            query_embedding: 查询向量 (1D或2D数组)
            top_k: 返回的候选数量
        
        返回：
            [(dish_name, similarity_score), ...] 按相似度降序排列
        """
        # 确保是2D数组 (1, embedding_dim)
        if query_embedding.ndim == 1:
            query_embedding = query_embedding.reshape(1, -1)
        
        # 归一化（因为索引是IndexFlatIP，需要归一化才能表示余弦相似度）
        norm = np.linalg.norm(query_embedding, axis=1, keepdims=True)
        if norm[0, 0] > 0:
            query_embedding = query_embedding / norm
        
        # 执行检索
        scores, indices = self.index.search(
            query_embedding.astype(np.float32),
            min(top_k, len(self.dish_names))
        )
        
        # 构建结果
        results = []
        for score, idx in zip(scores[0], indices[0]):
            if idx == -1:  # FAISS在结果不足时会返回-1
                break
            dish_name = self.dish_names[idx]
            results.append((dish_name, float(score)))
        
        print(
            f"[DishNameMatcher] Query shape={query_embedding.shape} "
            f"top_k={top_k} results={len(results)}"
        )
        for rank, (name, score) in enumerate(results, 1):
            print(f"  [{rank}] {name} (score={score:.4f})")
        
        return results
    
    def match_text(
        self,
        text: str,
        embedder,
        top_k: int = 3
    ) -> List[Tuple[str, float]]:
        """
        便捷方法：直接从文本检索（自动调用embedder）
        
        参数：
            text: 查询文本（菜品名）
            embedder: OllamaEmbedder实例
            top_k: 返回的候选数量
        """
        query_emb = embedder.encode_single(text)
        return self.match(query_emb, top_k)