# -*- coding: utf-8 -*-
"""
知识图谱可视化工具
使用networkx和matplotlib绘制知识图谱
"""

import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt
from matplotlib import font_manager
import matplotlib.patches as mpatches
from pathlib import Path
import numpy as np


# 设置中文字体
def setup_chinese_font():
    """设置中文字体"""
    # 尝试多种中文字体
    chinese_fonts = [
        'Microsoft YaHei',
        'SimHei',
        'SimSun',
        'KaiTi',
        'FangSong',
        'STXihei',
        'STSong',
        'STKaiti',
        'STFangsong'
    ]

    for font in chinese_fonts:
        try:
            plt.rcParams['font.sans-serif'] = [font]
            plt.rcParams['axes.unicode_minus'] = False
            break
        except:
            continue

    print(f"使用字体: {plt.rcParams['font.sans-serif'][0]}")


# 节点类型颜色映射
NODE_TYPE_COLORS = {
    'ICH_Project': '#FF6B6B',      # 红色 - 项目
    'Category': '#4ECDC4',         # 青色 - 类别
    'Batch': '#95E1D3',            # 绿色 - 批次
    'Inheritor': '#FFD93D',        # 黄色 - 传承人
    'Institution': '#6C5CE7'       # 紫色 - 机构
}

# 节点类型大小映射
NODE_TYPE_SIZES = {
    'ICH_Project': 300,
    'Category': 500,
    'Batch': 350,
    'Inheritor': 200,
    'Institution': 250
}


def load_graph_data(nodes_csv, rels_csv):
    """加载图谱数据"""
    print("正在加载数据...")

    # 读取节点和关系
    nodes_df = pd.read_csv(nodes_csv, encoding='utf-8-sig')
    rels_df = pd.read_csv(rels_csv, encoding='utf-8-sig')

    print(f"  - 节点: {len(nodes_df)}")
    print(f"  - 关系: {len(rels_df)}")

    return nodes_df, rels_df


def build_networkx_graph(nodes_df, rels_df):
    """构建NetworkX图"""
    print("正在构建图...")

    G = nx.DiGraph()  # 有向图

    # 添加节点
    for idx, row in nodes_df.iterrows():
        node_id = row['id']
        label = row['label']
        node_type = row['type']

        # 截断过长的标签
        if len(label) > 10:
            display_label = label[:10] + '...'
        else:
            display_label = label

        G.add_node(
            node_id,
            label=display_label,
            full_label=label,
            node_type=node_type,
            color=NODE_TYPE_COLORS.get(node_type, '#CCCCCC'),
            size=NODE_TYPE_SIZES.get(node_type, 200)
        )

    # 添加边
    for idx, row in rels_df.iterrows():
        source = row['source']
        target = row['target']
        rel_type = row['type']

        if source in G.nodes() and target in G.nodes():
            G.add_edge(source, target, rel_type=rel_type)

    print(f"  - 节点数: {G.number_of_nodes()}")
    print(f"  - 边数: {G.number_of_edges()}")

    return G


def filter_graph_by_type(G, include_types=None):
    """按节点类型过滤图"""
    if include_types is None:
        return G

    nodes_to_keep = [n for n, d in G.nodes(data=True)
                     if d.get('node_type') in include_types]

    return G.subgraph(nodes_to_keep).copy()


def draw_graph(G, output_path, title="知识图谱", layout='spring'):
    """绘制知识图谱"""
    print(f"正在绘制图谱: {title}")

    plt.figure(figsize=(20, 16))

    # 选择布局算法
    if layout == 'spring':
        pos = nx.spring_layout(G, k=2, iterations=50, seed=42)
    elif layout == 'circular':
        pos = nx.circular_layout(G)
    elif layout == 'kamada_kawai':
        pos = nx.kamada_kawai_layout(G)
    elif layout == 'random':
        pos = nx.random_layout(G)
    else:
        pos = nx.spring_layout(G, k=2, iterations=50, seed=42)

    # 按节点类型分组
    node_types = {}
    for node, data in G.nodes(data=True):
        node_type = data.get('node_type', 'Unknown')
        if node_type not in node_types:
            node_types[node_type] = []
        node_types[node_type].append(node)

    # 绘制边
    nx.draw_networkx_edges(
        G, pos,
        alpha=0.3,
        width=0.5,
        edge_color='gray',
        arrows=True,
        arrowsize=10,
        arrowstyle='->,head_width=0.2,head_length=0.3'
    )

    # 按类型绘制节点
    for node_type, nodes in node_types.items():
        color = NODE_TYPE_COLORS.get(node_type, '#CCCCCC')
        size = NODE_TYPE_SIZES.get(node_type, 200)

        nx.draw_networkx_nodes(
            G, pos,
            nodelist=nodes,
            node_color=color,
            node_size=size,
            alpha=0.8,
            edgecolors='white',
            linewidths=1
        )

    # 绘制标签（只对重要节点）
    important_nodes = []
    important_labels = {}

    for node, data in G.nodes(data=True):
        node_type = data.get('node_type')
        # 只显示类别、批次和部分重要节点的标签
        if node_type in ['Category', 'Batch'] or (
            node_type == 'ICH_Project' and data.get('size', 0) > 400
        ):
            important_nodes.append(node)
            important_labels[node] = data.get('label', node)

    if len(important_nodes) <= 100:  # 节点不多时显示所有标签
        nx.draw_networkx_labels(
            G, pos,
            labels=important_labels,
            font_size=8,
            font_weight='bold',
            font_family='sans-serif'
        )
    else:
        # 节点太多时只显示类别标签
        category_labels = {n: d['label'] for n, d in G.nodes(data=True)
                          if d.get('node_type') == 'Category'}
        nx.draw_networkx_labels(
            G, pos,
            labels=category_labels,
            font_size=10,
            font_weight='bold'
        )

    # 图例
    legend_patches = []
    for node_type, color in NODE_TYPE_COLORS.items():
        if node_type in node_types:
            patch = mpatches.Patch(color=color, label=node_type)
            legend_patches.append(patch)

    plt.legend(
        handles=legend_patches,
        loc='upper right',
        fontsize=12,
        framealpha=0.9
    )

    plt.title(title, fontsize=16, fontweight='bold', pad=20)
    plt.axis('off')
    plt.tight_layout()

    # 保存图片
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    print(f"  - 保存到: {output_path}")
    plt.close()


def draw_subgraphs(G, output_dir):
    """绘制子图（按节点类型分组）"""
    print("\n正在绘制子图...")

    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # 1. 只显示项目和类别
    print("  1. 项目-类别关系图...")
    G1 = filter_graph_by_type(G, ['ICH_Project', 'Category'])
    draw_graph(
        G1,
        output_dir / 'kg_project_category.png',
        title='非遗项目与类别关系',
        layout='spring'
    )

    # 2. 只显示项目和传承人
    print("  2. 项目-传承人关系图...")
    G2 = filter_graph_by_type(G, ['ICH_Project', 'Inheritor'])
    if G2.number_of_nodes() > 0:
        draw_graph(
            G2,
            output_dir / 'kg_project_inheritor.png',
            title='非遗项目与传承人关系',
            layout='spring'
        )

    # 3. 只显示项目、类别和批次
    print("  3. 项目-类别-批次关系图...")
    G3 = filter_graph_by_type(G, ['ICH_Project', 'Category', 'Batch'])
    draw_graph(
        G3,
        output_dir / 'kg_project_category_batch.png',
        title='非遗项目、类别与批次关系',
        layout='kamada_kawai'
    )

    # 4. 完整图谱（抽样显示）
    print("  4. 完整知识图谱...")
    if G.number_of_nodes() > 500:
        # 节点太多时，只显示连接度高的节点
        degrees = dict(G.degree())
        high_degree_nodes = [n for n, d in degrees.items() if d >= 3]
        G_sample = G.subgraph(high_degree_nodes).copy()
        draw_graph(
            G_sample,
            output_dir / 'kg_full_sampled.png',
            title=f'完整知识图谱（抽样，显示{G_sample.number_of_nodes()}个节点）',
            layout='spring'
        )
    else:
        draw_graph(
            G,
            output_dir / 'kg_full.png',
            title='完整知识图谱',
            layout='spring'
        )


def print_statistics(G):
    """打印图统计信息"""
    print("\n" + "="*60)
    print("图谱统计信息")
    print("="*60)

    print(f"\n节点总数: {G.number_of_nodes()}")
    print(f"边总数: {G.number_of_edges()}")

    # 按类型统计节点
    print("\n节点类型分布:")
    node_types = {}
    for node, data in G.nodes(data=True):
        node_type = data.get('node_type', 'Unknown')
        node_types[node_type] = node_types.get(node_type, 0) + 1

    for node_type, count in sorted(node_types.items()):
        percentage = (count / G.number_of_nodes() * 100)
        print(f"  - {node_type}: {count} ({percentage:.1f}%)")

    # 按类型统计边
    print("\n关系类型分布:")
    rel_types = {}
    for u, v, data in G.edges(data=True):
        rel_type = data.get('rel_type', 'Unknown')
        rel_types[rel_type] = rel_types.get(rel_type, 0) + 1

    for rel_type, count in sorted(rel_types.items()):
        percentage = (count / G.number_of_edges() * 100)
        print(f"  - {rel_type}: {count} ({percentage:.1f}%)")

    # 连接度统计
    degrees = [d for n, d in G.degree()]
    print(f"\n连接度统计:")
    print(f"  - 平均连接度: {np.mean(degrees):.2f}")
    print(f"  - 最大连接度: {max(degrees)}")
    print(f"  - 最小连接度: {min(degrees)}")

    # 找出连接度最高的节点
    top_nodes = sorted(G.degree(), key=lambda x: x[1], reverse=True)[:10]
    print(f"\n连接度最高的10个节点:")
    for node, degree in top_nodes:
        node_data = G.nodes[node]
        label = node_data.get('full_label', node)
        node_type = node_data.get('node_type', '')
        print(f"  - [{node_type}] {label}: {degree}个连接")

    print("="*60)


def visualize_kg(nodes_csv, rels_csv, output_dir):
    """可视化知识图谱"""
    print("="*60)
    print("知识图谱可视化工具")
    print("="*60)

    # 设置中文字体
    setup_chinese_font()

    # 加载数据
    nodes_df, rels_df = load_graph_data(nodes_csv, rels_csv)

    # 构建图
    G = build_networkx_graph(nodes_df, rels_df)

    # 打印统计信息
    print_statistics(G)

    # 绘制子图
    draw_subgraphs(G, output_dir)

    print("\n" + "="*60)
    print("可视化完成！")
    print("="*60)


if __name__ == '__main__':
    # 输入文件
    nodes_csv = r'E:\Project\2026_KG_ICH\dofile\kg_project\output\nodes.csv'
    rels_csv = r'E:\Project\2026_KG_ICH\dofile\kg_project\output\rels.csv'

    # 输出目录
    output_dir = r'E:\Project\2026_KG_ICH\dofile\kg_project\output\visualizations'

    # 执行可视化
    visualize_kg(nodes_csv, rels_csv, output_dir)