import numpy as np
from models.model_loader import model_loader

class AIService:
    # 创建AI处理类
    def __init__(self):
        self.text_model = model_loader.load_text_model()

    # 常用AI方法
    def text_to_embedding(self, text):
        # 1.文本转向量
        if isinstance(text, str):
            text = [text]
        
        embeddings = self.text_model.encode(text)
        # 返回numpy数组，不要转成list
        return embeddings
    
    def calculate_similarity(self, text1, text2):
        # 2.计算文本相似度
        emb1 = self.text_to_embedding(text1)
        emb2 = self.text_to_embedding(text2)
        
        # 确保是1D数组
        if len(emb1.shape) > 1:
            emb1 = emb1[0]
        if len(emb2.shape) > 1:
            emb2 = emb2[0]
            
        similarity = np.dot(emb1, emb2) / (
            np.linalg.norm(emb1) * np.linalg.norm(emb2)
        )
        return float(similarity)
    
    def batch_text_processing(self, texts):
        # 3.批量文本处理
        embeddings = self.text_to_embedding(texts)
        
        # 简单的文本分析（可扩展）
        results = []
        for i, text in enumerate(texts):
            results.append({
                'text': text,
                'embedding_size': embeddings.shape[1],  # 使用shape获取维度
                'text_length': len(text),
                'embedding': embeddings[i].tolist()  # 只在返回给前端时转成list
            })
        
        return results

# 全局AI服务实例
ai_service = AIService()