# -*- coding: utf-8 -*-
"""
整合基础政务实体和深层语义实体
方案3：合并项目节点，完全整合
"""

import pandas as pd
import json
from pathlib import Path

def merge_kg_data():
    """整合知识图谱数据"""

    output_dir = Path(__file__).parent / 'output'

    # 读取4个文件
    print("=== 读取原始文件 ===")
    nodes_basic = pd.read_csv(output_dir / 'nodes.csv', encoding='utf-8-sig')
    nodes_deep = pd.read_csv(output_dir / 'nodes_desc.csv', encoding='utf-8-sig')
    rels_basic = pd.read_csv(output_dir / 'rels.csv', encoding='utf-8-sig')
    rels_deep = pd.read_csv(output_dir / 'rels_desc.csv', encoding='utf-8-sig')

    print(f"基础节点: {len(nodes_basic)}")
    print(f"深层节点: {len(nodes_deep)}")
    print(f"基础关系: {len(rels_basic)}")
    print(f"深层关系: {len(rels_deep)}")

    # 1. 合并ICH项目节点
    print("\n=== 合并ICH项目节点 ===")
    basic_projects = nodes_basic[nodes_basic['type'] == 'ICH_Project'].copy()
    deep_projects = nodes_deep[nodes_deep['type'] == 'ICH_Project'].copy()

    print(f"基础项目节点: {len(basic_projects)}")
    print(f"深层项目节点: {len(deep_projects)}")

    # 合并项目节点的属性
    merged_projects = []

    for pid in basic_projects['id']:
        basic_row = basic_projects[basic_projects['id'] == pid].iloc[0]

        # 查找深层项目节点
        deep_row = deep_projects[deep_projects['id'] == pid]

        if len(deep_row) > 0:
            # 合并属性
            deep_row = deep_row.iloc[0]

            # 解析属性JSON
            basic_props = json.loads(basic_row['properties']) if basic_row['properties'] else {}
            deep_props = json.loads(deep_row['properties']) if deep_row['properties'] else {}

            # 合并属性（深层属性优先）
            merged_props = {**basic_props, **deep_props}

            merged_projects.append({
                'id': pid,
                'label': basic_row['label'],  # 保留基础标签
                'type': 'ICH_Project',
                'properties': json.dumps(merged_props, ensure_ascii=False)
            })
        else:
            # 只在基础中存在
            merged_projects.append({
                'id': pid,
                'label': basic_row['label'],
                'type': 'ICH_Project',
                'properties': basic_row['properties']
            })

    print(f"合并后项目节点: {len(merged_projects)}")

    # 2. 合并其他节点（去除ICH_Project）
    print("\n=== 合并其他节点 ===")
    other_basic_nodes = nodes_basic[nodes_basic['type'] != 'ICH_Project']
    other_deep_nodes = nodes_deep[nodes_deep['type'] != 'ICH_Project']

    all_nodes = pd.concat([
        pd.DataFrame(merged_projects),
        other_basic_nodes,
        other_deep_nodes
    ], ignore_index=True)

    print(f"合并后总节点数: {len(all_nodes)}")
    print(f"  - ICH_Project: {len(merged_projects)}")
    print(f"  - 其他节点: {len(all_nodes) - len(merged_projects)}")

    # 3. 合并关系
    print("\n=== 合并关系 ===")
    all_rels = pd.concat([rels_basic, rels_deep], ignore_index=True)

    # 去重关系
    all_rels = all_rels.drop_duplicates(subset=['source', 'target', 'type'], keep='first')

    print(f"合并后关系数: {len(all_rels)}")

    # 4. 统计信息
    print("\n=== 整合后统计 ===")
    print(f"总节点数: {len(all_nodes)}")
    print(f"总关系数: {len(all_rels)}")

    print("\n节点类型分布:")
    for ntype, count in all_nodes.groupby('type').size().items():
        print(f"  {ntype}: {count}")

    print(f"\n关系类型数量: {len(all_rels['type'].unique())}")

    # 5. 保存整合后的文件
    print("\n=== 保存整合文件 ===")
    output_file_nodes = output_dir / 'kg_merged_nodes.csv'
    output_file_rels = output_dir / 'kg_merged_rels.csv'

    all_nodes.to_csv(output_file_nodes, index=False, encoding='utf-8-sig')
    all_rels.to_csv(output_file_rels, index=False, encoding='utf-8-sig')

    print(f"节点已保存: {output_file_nodes}")
    print(f"关系已保存: {output_file_rels}")

    # 6. 验证
    print("\n=== 验证 ===")

    # 检查项目节点完整性
    project_count = len(all_nodes[all_nodes['type'] == 'ICH_Project'])
    print(f"ICH项目节点数: {project_count} (应该是268)")

    # 检查关系完整性
    rel_sources = set(all_rels['source'].unique())
    node_ids = set(all_nodes['id'].unique())

    missing_sources = rel_sources - node_ids
    if missing_sources:
        print(f"警告: {len(missing_sources)} 个关系的源节点不在节点文件中")
    else:
        print("所有关系的源节点都存在于节点文件中")

    # 统计每个项目的关系数
    project_rels = all_rels[all_rels['source'].str.startswith('ICH-')].groupby('source').size()
    print(f"\n有关系的项目数: {len(project_rels)}/268")

    # 展示几个示例项目的统计
    print("\n示例项目关系统计:")
    sample_projects = ['ICH-1', 'ICH-118', 'ICH-229']
    for pid in sample_projects:
        basic_rel_count = len(rels_basic[rels_basic['source'] == pid])
        deep_rel_count = len(rels_deep[rels_deep['source'] == pid])
        total_rel_count = len(all_rels[all_rels['source'] == pid])
        print(f"  {pid}: 基础{basic_rel_count} + 深层{deep_rel_count} = 总计{total_rel_count} 条关系")

    print("\n完成！")

    return all_nodes, all_rels


if __name__ == '__main__':
    merge_kg_data()