#!/usr/bin/env python3
"""
激进版后处理 - 只保留高置信度关系
"""

import json


def aggressive_filter(input_path: str, output_path: str):
    with open(input_path) as f:
        data = json.load(f)
    
    print("=" * 70)
    print("Aggressive Post-Processing")
    print("=" * 70)
    
    original_rels = len(data.get('relations', []))
    print(f"\nOriginal: {original_rels} relations")
    
    # 只保留这些核心关系类型
    core_types = {'member_of', 'classified_into', 'expressed_in', 'contains', 'involved_in'}
    
    # 只保留这些核心实体参与的关系
    core_entities = {'SPL', 'SBP', 'FtSPL1', 'FtSPL2', 'FtSPL3', 'FtSPL4', 
                     'Group_I', 'Group_II', 'Group_VI', 'flower', 'fruit',
                     'flower_development', 'fruit_development'}
    
    filtered = []
    for r in data.get('relations', []):
        rel_type = r.get('type') or r.get('relation', '')
        source = r.get('source', '')
        target = r.get('target', '')
        
        # 规则1: 必须是核心类型
        if rel_type not in core_types:
            continue
        
        # 规则2: 源实体必须是核心实体
        if source not in core_entities:
            continue
        
        # 规则3: 删除明显错误
        # 错误: SBP --member_of--> SPL
        if rel_type == 'member_of' and source in ['SBP', 'ANK']:
            continue
        
        # 错误: SPL --belongs_to--> organism
        if rel_type == 'belongs_to' and source == 'SPL':
            continue
        
        # 错误: gene --involved_in--> gene
        if rel_type == 'involved_in' and target.startswith('FtSPL'):
            continue
        
        # 标准化字段
        r['type'] = rel_type
        if 'relation' in r:
            del r['relation']
        
        filtered.append(r)
    
    data['relations'] = filtered
    
    print(f"Filtered: {len(filtered)} relations (removed {original_rels - len(filtered)})")
    
    print("\nKept relations:")
    for r in filtered:
        print(f"  {r['source']} --{r['type']}--> {r['target']}")
    
    with open(output_path, 'w') as f:
        json.dump(data, f, indent=2)
    
    print(f"\n✓ Saved to: {output_path}")
    
    # 对比
    print("\n" + "=" * 70)
    import subprocess
    subprocess.run([
        'python3', 'evaluation/compare_extraction.py',
        '--gold', 'evaluation/gold_standard/pmid31286919_layer2_full.json',
        '--pred', output_path
    ])


if __name__ == "__main__":
    aggressive_filter(
        "evaluation/auto_results/pmid31286919_optimized_v3.json",
        "evaluation/auto_results/pmid31286919_aggressive.json"
    )