#!/usr/bin/env python3
"""
智能推理引擎 - 多跳推理
支持复杂查询的推理链
"""

import json
import sqlite3
from typing import List, Dict, Optional
from dataclasses import dataclass

@dataclass
class ReasoningStep:
    """推理步骤"""
    step_num: int
    query_type: str
    query_params: Dict
    result: Dict
    reasoning: str

class ReasoningEngine:
    """多跳推理引擎"""
    
    def __init__(self, db_path="/home/xzq/gene-family-mas-v2/kg.db"):
        self.db_path = db_path
        self.reasoning_chain = []
        
    def get_db(self):
        return sqlite3.connect(self.db_path)
    
    def multi_hop_query(self, question: str) -> Dict:
        """
        多跳推理查询
        
        支持的查询类型:
        1. 基因-分组-基因 (同组的其他基因)
        2. 基因家族-物种-基因 (某物种某家族的所有基因)
        3. 分组-统计-比较 (分组间的比较)
        """
        question_lower = question.lower()
        
        # 解析查询意图
        if "同组" in question or "同一分组" in question:
            return self._reasoning_same_group(question)
        elif "统计" in question and "比较" in question:
            return self._reasoning_group_comparison(question)
        elif "最大" in question or "最小" in question:
            return self._reasoning_extreme_group(question)
        else:
            return self._reasoning_general(question)
    
    def _reasoning_same_group(self, question: str) -> Dict:
        """推理：查找与某基因同组的其他基因"""
        # Step 1: 提取基因ID
        import re
        gene_match = re.search(r'(CYP\w+)', question)
        if not gene_match:
            return {"error": "无法识别基因ID"}
        
        gene_id = gene_match.group(1)
        
        conn = self.get_db()
        cursor = conn.cursor()
        
        # Step 2: 查找该基因所在分组
        cursor.execute("""
            SELECT pg.name FROM phylogeny_groups pg
            JOIN gene_groups gg ON pg.id = gg.group_id
            WHERE gg.gene_id = ?
        """, (gene_id,))
        row = cursor.fetchone()
        if not row:
            return {"error": f"基因 {gene_id} 未找到分组信息"}
        
        group_name = row[0]
        
        # Step 3: 查找同组的其他基因
        cursor.execute("""
            SELECT g.id, g.name FROM genes g
            JOIN gene_groups gg ON g.id = gg.gene_id
            JOIN phylogeny_groups pg ON gg.group_id = pg.id
            WHERE pg.name = ? AND g.id != ?
        """, (group_name, gene_id))
        
        same_group_genes = [{"id": r[0], "name": r[1]} for r in cursor.fetchall()]
        conn.close()
        
        # 构建推理链
        reasoning_chain = [
            ReasoningStep(1, "entity_resolution", {"gene_id": gene_id}, 
                         {"found": True}, f"识别目标基因: {gene_id}"),
            ReasoningStep(2, "group_lookup", {"gene_id": gene_id}, 
                         {"group": group_name}, f"查询到基因所在分组: {group_name}"),
            ReasoningStep(3, "group_expansion", {"group": group_name}, 
                         {"genes": same_group_genes}, f"分组扩展查询，找到{len(same_group_genes)}个同组基因")
        ]
        
        return {
            "question": question,
            "answer": f"基因 {gene_id} 属于 {group_name}，同组还有 {len(same_group_genes)} 个基因",
            "reasoning_chain": [
                {"step": r.step_num, "type": r.query_type, "reasoning": r.reasoning} 
                for r in reasoning_chain
            ],
            "data": {
                "target_gene": gene_id,
                "group": group_name,
                "same_group_genes": same_group_genes
            }
        }
    
    def _reasoning_group_comparison(self, question: str) -> Dict:
        """推理：分组统计比较"""
        conn = self.get_db()
        cursor = conn.cursor()
        
        # 获取所有分组的统计
        cursor.execute("""
            SELECT pg.name, COUNT(gg.gene_id) as cnt
            FROM phylogeny_groups pg
            LEFT JOIN gene_groups gg ON pg.id = gg.group_id
            WHERE pg.gene_family = 'P450'
            GROUP BY pg.name
            ORDER BY cnt DESC
        """)
        
        groups = [{"name": r[0], "count": r[1]} for r in cursor.fetchall()]
        conn.close()
        
        # 计算统计信息
        max_group = groups[0] if groups else None
        min_group = groups[-1] if groups else None
        avg_count = sum(g["count"] for g in groups) / len(groups) if groups else 0
        
        reasoning_chain = [
            ReasoningStep(1, "aggregation", {"table": "phylogeny_groups"}, 
                         {"total_groups": len(groups)}, "聚合所有分组信息"),
            ReasoningStep(2, "statistics", {"groups": groups}, 
                         {"max": max_group, "min": min_group, "avg": avg_count}, "计算分组统计指标")
        ]
        
        return {
            "question": question,
            "answer": f"P450基因家族共{len(groups)}个分组，最大分组是{max_group['name']}({max_group['count']}个基因)，最小分组是{min_group['name']}({min_group['count']}个基因)，平均每组{avg_count:.1f}个基因",
            "reasoning_chain": [
                {"step": r.step_num, "type": r.query_type, "reasoning": r.reasoning} 
                for r in reasoning_chain
            ],
            "data": {
                "total_groups": len(groups),
                "max_group": max_group,
                "min_group": min_group,
                "average": avg_count,
                "all_groups": groups
            }
        }
    
    def _reasoning_extreme_group(self, question: str) -> Dict:
        """推理：查找最大/最小分组"""
        conn = self.get_db()
        cursor = conn.cursor()
        
        if "最大" in question:
            cursor.execute("""
                SELECT pg.name, COUNT(gg.gene_id) as cnt
                FROM phylogeny_groups pg
                LEFT JOIN gene_groups gg ON pg.id = gg.group_id
                WHERE pg.gene_family = 'P450'
                GROUP BY pg.name
                ORDER BY cnt DESC
                LIMIT 1
            """)
            row = cursor.fetchone()
            conn.close()
            
            return {
                "question": question,
                "answer": f"最大的分组是 {row[0]}，包含 {row[1]} 个基因",
                "data": {"group": row[0], "count": row[1]}
            }
        else:
            return {"error": "暂不支持此类型查询"}
    
    def _reasoning_general(self, question: str) -> Dict:
        """通用推理"""
        return {
            "question": question,
            "answer": "这是一个通用查询，建议使用具体查询如'查找与CYP71D184同组的基因'",
            "suggestions": [
                "查找与CYP71D184同组的基因",
                "P450分组统计比较",
                "最大的分组是哪个"
            ]
        }

if __name__ == "__main__":
    engine = ReasoningEngine()
    
    # 测试多跳推理
    test_questions = [
        "查找与CYP71D184同组的基因",
        "P450分组统计比较",
        "最大的分组是哪个"
    ]
    
    for q in test_questions:
        print(f"\n{'='*60}")
        print(f"问题: {q}")
        result = engine.multi_hop_query(q)
        print(f"答案: {result.get('answer')}")
        if 'reasoning_chain' in result:
            print("\n推理链:")
            for step in result['reasoning_chain']:
                print(f"  Step {step['step']}: {step['reasoning']}")