#!/usr/bin/env python3
"""
应用 locus ID 到基因符号的映射
"""

import json


def load_gene_list():
    """加载基因列表"""
    with open('ftspl_gene_list.json') as f:
        return json.load(f)


def create_mapping(gene_list):
    """创建 locus ID -> gene symbol 映射"""
    mapping = {}
    for gene in gene_list:
        locus = gene['locus_id']
        symbol = gene['symbol']
        # 处理多种格式
        mapping[locus] = symbol
        mapping[locus.replace('.T01', '')] = symbol
        mapping[locus.replace('.01.T01', '')] = symbol
    return mapping


def apply_mapping(data, mapping):
    """应用映射到实体和关系"""
    # 映射实体 ID
    for e in data.get('entities', []):
        old_id = e['id']
        if old_id in mapping:
            e['id'] = mapping[old_id]
            e['name'] = mapping[old_id]
    
    # 映射关系中的 ID
    for r in data.get('relations', []):
        source = r.get('source', '')
        target = r.get('target', '')
        
        if source in mapping:
            r['source'] = mapping[source]
        if target in mapping:
            r['target'] = mapping[target]
    
    return data


def main():
    print("=" * 70)
    print("Applying Locus ID to Gene Symbol Mapping")
    print("=" * 70)
    
    # 加载基因列表
    print("\n1. Loading gene list...")
    gene_list = load_gene_list()
    mapping = create_mapping(gene_list)
    print(f"   Created mapping for {len(mapping)} genes")
    
    # 显示映射示例
    print("\n2. Mapping examples:")
    for i, (locus, symbol) in enumerate(list(mapping.items())[:5]):
        print(f"   {locus} -> {symbol}")
    
    # 应用映射到各个结果文件
    files_to_process = [
        ('evaluation/auto_results/pmid31286919_phylogeny.json', 
         'evaluation/auto_results/pmid31286919_phylogeny_mapped.json'),
    ]
    
    for input_path, output_path in files_to_process:
        print(f"\n3. Processing: {input_path}")
        
        with open(input_path) as f:
            data = json.load(f)
        
        # 应用映射
        data = apply_mapping(data, mapping)
        
        # 保存
        with open(output_path, 'w') as f:
            json.dump(data, f, indent=2)
        
        print(f"   ✓ Saved to: {output_path}")
        
        # 显示结果
        print(f"\n   Mapped relations:")
        for r in data.get('relations', []):
            rel_type = r.get('type') or r.get('relation', '')
            if 'classified' in rel_type.lower():
                print(f"     {r['source']} --{rel_type}--> {r['target']}")


if __name__ == "__main__":
    main()