#!/usr/bin/env python3
"""
Gold Standard 人工标注工具
用于人工标注论文中的实体和关系，作为评估自动提取的基准
"""

import json
import re
from dataclasses import dataclass, asdict
from typing import List, Optional, Dict
from datetime import datetime


@dataclass
class Entity:
    """实体标注"""
    id: str
    name: str
    type: str
    species: Optional[str] = None
    primary_id: Optional[str] = None
    attributes: Dict = None
    description: Optional[str] = None
    text_span: Optional[Dict] = None  # {"start": int, "end": int, "text": str}
    
    def __post_init__(self):
        if self.attributes is None:
            self.attributes = {}


@dataclass
class Relation:
    """关系标注"""
    id: str
    source: str  # 源实体ID
    target: str  # 目标实体ID
    type: str
    confidence: str  # "high" | "medium" | "low"
    qualifiers: Optional[Dict] = None
    description: Optional[str] = None
    evidence: Optional[Dict] = None  # 证据信息
    
    def __post_init__(self):
        if self.qualifiers is None:
            self.qualifiers = {}


@dataclass
class GoldStandard:
    """单篇论文的Gold Standard"""
    doc_id: str
    pmid: Optional[str]
    doi: Optional[str]
    title: str
    text: str  # 标注的文本内容
    species: str
    family: str
    entities: List[Entity]
    relations: List[Relation]
    key_findings: List[str]
    annotated_by: str
    annotated_at: str
    notes: Optional[str] = None


class GoldStandardAnnotator:
    """人工标注工具类"""
    
    # 25种核心实体类型
    ENTITY_TYPES = [
        "gene", "protein", "gene_family", "protein_domain", "motif",
        "organism", "cultivar", "tissue", "cell_type", "developmental_stage",
        "treatment", "hormone", "environmental_factor", "phenotype", "trait",
        "biological_process", "signaling_pathway", "metabolite", "mirna", 
        "transcription_factor", "promoter",
        "dataset", "publication", "method", "assay"
    ]
    
    # 35种核心关系类型
    RELATION_TYPES = [
        # 调控
        "activates", "inhibits", "regulates",
        # 表达
        "expressed_in", "upregulated_in", "downregulated_in", 
        "induced_by", "repressed_by",
        # 功能
        "involved_in", "controls", "affects", "promotes", "mediates",
        # 进化
        "homologous_to", "orthologous_to", "paralogous_to", "syntenic_with",
        # 层级
        "member_of", "contains", "part_of", "belongs_to", "classified_into",
        # 互作
        "binds_to", "interacts_with", "targets",
        # 数据
        "annotated_in", "stored_in", "originated_from"
    ]
    
    def __init__(self, doc_id: str, text: str, title: str = ""):
        self.doc_id = doc_id
        self.text = text
        self.title = title
        self.pmid = None
        self.doi = None
        self.species = None
        self.family = None
        self.entities: List[Entity] = []
        self.relations: List[Relation] = []
        self.key_findings: List[str] = []
        self.notes = ""
        
    def set_metadata(self, pmid: str = None, doi: str = None, 
                     species: str = None, family: str = None):
        """设置论文元数据"""
        self.pmid = pmid
        self.doi = doi
        self.species = species
        self.family = family
        
    def add_entity(self, id: str, name: str, type: str,
                   species: str = None, primary_id: str = None,
                   text_start: int = None, text_end: int = None,
                   description: str = None, **attributes) -> Entity:
        """
        添加实体标注
        
        Args:
            id: 实体唯一ID（如FtSPL1）
            name: 实体显示名称
            type: 实体类型（必须在25种中）
            species: 所属物种
            primary_id: 稳定外部ID
            text_start: 在原文中的起始位置
            text_end: 在原文中的结束位置
            description: 描述
            **attributes: 其他属性
        """
        if type not in self.ENTITY_TYPES:
            print(f"警告: 实体类型 '{type}' 不在标准25种中")
            print(f"标准类型: {self.ENTITY_TYPES}")
        
        text_span = None
        if text_start is not None and text_end is not None:
            text_span = {
                "start": text_start,
                "end": text_end,
                "text": self.text[text_start:text_end]
            }
        
        entity = Entity(
            id=id,
            name=name,
            type=type,
            species=species,
            primary_id=primary_id,
            attributes=attributes,
            description=description,
            text_span=text_span
        )
        self.entities.append(entity)
        return entity
    
    def add_relation(self, source: str, target: str, type: str,
                     confidence: str = "high",
                     text_snippet: str = None,
                     section: str = None,
                     figure: str = None,
                     table: str = None,
                     evidence_level: str = "experiment",
                     method: str = None,
                     description: str = None,
                     **qualifiers) -> Relation:
        """
        添加关系标注
        
        Args:
            source: 源实体ID
            target: 目标实体ID
            type: 关系类型（必须在35种中）
            confidence: 置信度 (high/medium/low)
            text_snippet: 原文证据片段
            section: 所在章节
            figure: 图号
            table: 表号
            evidence_level: 证据等级 (experiment/omics/prediction/curation)
            method: 实验/分析方法
            description: 描述
            **qualifiers: 其他限定条件
        """
        if type not in self.RELATION_TYPES:
            print(f"警告: 关系类型 '{type}' 不在标准35种中")
            print(f"标准类型: {self.RELATION_TYPES}")
        
        # 检查source和target是否存在
        source_exists = any(e.id == source for e in self.entities)
        target_exists = any(e.id == target for e in self.entities)
        
        if not source_exists:
            print(f"警告: 源实体 '{source}' 不存在于entities中")
        if not target_exists:
            print(f"警告: 目标实体 '{target}' 不存在于entities中")
        
        evidence = {
            "source_type": "publication",
            "source_id": f"pmid:{self.pmid}" if self.pmid else self.doi,
            "snippet": text_snippet,
            "evidence_level": evidence_level,
            "method": method
        }
        
        if section or figure or table:
            evidence["locator"] = {}
            if section:
                evidence["locator"]["section"] = section
            if figure:
                evidence["locator"]["figure"] = figure
            if table:
                evidence["locator"]["table"] = table
        
        relation = Relation(
            id=f"rel_{len(self.relations):03d}",
            source=source,
            target=target,
            type=type,
            confidence=confidence,
            qualifiers=qualifiers,
            description=description,
            evidence=evidence
        )
        self.relations.append(relation)
        return relation
    
    def add_key_finding(self, finding: str):
        """添加关键发现"""
        self.key_findings.append(finding)
        
    def add_note(self, note: str):
        """添加标注笔记"""
        self.notes += note + "\n"
    
    def validate(self) -> List[str]:
        """验证标注完整性，返回问题列表"""
        issues = []
        
        # 检查必需字段
        if not self.species:
            issues.append("缺少species")
        if not self.family:
            issues.append("缺少family")
        
        # 检查实体ID唯一性
        ids = [e.id for e in self.entities]
        if len(ids) != len(set(ids)):
            issues.append("实体ID有重复")
        
        # 检查关系的source/target存在性
        for rel in self.relations:
            if not any(e.id == rel.source for e in self.entities):
                issues.append(f"关系 {rel.id} 的source '{rel.source}' 不存在")
            if not any(e.id == rel.target for e in self.entities):
                issues.append(f"关系 {rel.id} 的target '{rel.target}' 不存在")
        
        # 检查high/medium confidence必须有evidence
        for rel in self.relations:
            if rel.confidence in ["high", "medium"]:
                if not rel.evidence or not rel.evidence.get("snippet"):
                    issues.append(f"关系 {rel.id} (confidence={rel.confidence}) 缺少evidence.snippet")
        
        return issues
    
    def to_dict(self) -> dict:
        """转换为字典格式"""
        return {
            "doc_id": self.doc_id,
            "pmid": self.pmid,
            "doi": self.doi,
            "title": self.title,
            "text": self.text,
            "species": self.species,
            "family": self.family,
            "entities": [asdict(e) for e in self.entities],
            "relations": [asdict(r) for r in self.relations],
            "key_findings": self.key_findings,
            "annotated_by": "manual",
            "annotated_at": datetime.now().isoformat(),
            "notes": self.notes
        }
    
    def save(self, filepath: str):
        """保存到JSON文件"""
        data = self.to_dict()
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=2, ensure_ascii=False)
        print(f"已保存到: {filepath}")
    
    @classmethod
    def load(cls, filepath: str) -> 'GoldStandardAnnotator':
        """从JSON文件加载"""
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        annotator = cls(data['doc_id'], data['text'], data['title'])
        annotator.set_metadata(data.get('pmid'), data.get('doi'), 
                              data.get('species'), data.get('family'))
        
        for e_data in data.get('entities', []):
            annotator.entities.append(Entity(**e_data))
        
        for r_data in data.get('relations', []):
            annotator.relations.append(Relation(**r_data))
        
        annotator.key_findings = data.get('key_findings', [])
        annotator.notes = data.get('notes', '')
        
        return annotator
    
    def print_summary(self):
        """打印标注摘要"""
        print("="*60)
        print(f"Gold Standard 摘要: {self.doc_id}")
        print("="*60)
        print(f"PMID: {self.pmid}")
        print(f"物种: {self.species}")
        print(f"家族: {self.family}")
        print(f"实体数: {len(self.entities)}")
        print(f"关系数: {len(self.relations)}")
        print(f"关键发现: {len(self.key_findings)}")
        
        print("\n实体列表:")
        for e in self.entities:
            print(f"  - {e.id} ({e.type}): {e.name}")
        
        print("\n关系列表:")
        for r in self.relations:
            print(f"  - {r.source} --[{r.type}]--> {r.target} ({r.confidence})")
        
        issues = self.validate()
        if issues:
            print("\n⚠️ 验证问题:")
            for issue in issues:
                print(f"  - {issue}")
        else:
            print("\n✅ 验证通过")


# ==================== 使用示例 ====================

def example_usage():
    """使用示例"""
    
    # 论文文本（Abstract + Results开头）
    text = """Background: SPL (SQUAMOSA promoter binding protein-like) is a class 
    of plant-specific transcription factors that play important roles in many 
    growth and developmental processes... 24 Fagopyrum tataricum SPL (FtSPL) 
    genes were identified... According to the amino acid sequence of the SBP 
    domain and gene structure, the SPL genes were divided into eight groups..."""
    
    # 创建标注器
    annotator = GoldStandardAnnotator(
        doc_id="pmid31286919",
        text=text,
        title="Genome-wide identification of the SPL gene family in Tartary Buckwheat"
    )
    
    # 设置元数据
    annotator.set_metadata(
        pmid="31286919",
        doi="10.1186/s12870-019-1916-6",
        species="Fagopyrum tataricum",
        family="SPL"
    )
    
    # 添加实体
    annotator.add_entity(
        id="FtSPL",
        name="FtSPL genes",
        type="gene",
        species="Fagopyrum tataricum",
        text_start=150,
        text_end=158,
        description="Fagopyrum tataricum SPL genes, 24 members"
    )
    
    annotator.add_entity(
        id="SPL",
        name="SPL gene family",
        type="gene_family",
        description="SQUAMOSA promoter binding protein-like family"
    )
    
    annotator.add_entity(
        id="SBP",
        name="SBP domain",
        type="protein_domain",
        primary_id="PF03110",
        description="SQUAMOSA PROMOTER BINDING PROTEIN domain"
    )
    
    annotator.add_entity(
        id="flower",
        name="flower",
        type="tissue"
    )
    
    # 添加关系
    annotator.add_relation(
        source="FtSPL",
        target="SPL",
        type="member_of",
        confidence="high",
        text_snippet="24 Fagopyrum tataricum SPL (FtSPL) genes were identified",
        section="abstract",
        evidence_level="prediction",
        method="genome annotation",
        member_count=24
    )
    
    annotator.add_relation(
        source="FtSPL",
        target="SBP",
        type="contains",
        confidence="high",
        text_snippet="According to the amino acid sequence of the SBP domain",
        section="abstract",
        evidence_level="prediction",
        method="domain annotation"
    )
    
    annotator.add_relation(
        source="FtSPL",
        target="flower",
        type="involved_in",
        confidence="medium",
        text_snippet="genes that might regulate flower and fruit development",
        section="abstract",
        evidence_level="experiment",
        method="qRT-PCR"
    )
    
    # 添加关键发现
    annotator.add_key_finding("24 FtSPL genes identified in tartary buckwheat")
    annotator.add_key_finding("FtSPL genes divided into eight groups")
    annotator.add_key_finding("FtSPLs may regulate flower and fruit development")
    
    # 验证并保存
    annotator.print_summary()
    annotator.save("gold_standard/pmid31286919.json")


if __name__ == "__main__":
    example_usage()