parse_data.py

"""
解析JSONL文件，提取节点和边数据用于可视化
"""

import json
from collections import defaultdict
from typing import Dict, List, Set


def parse_jsonl_to_graph(jsonl_path: str) -> Dict:
    """
    解析JSONL文件，构建图数据结构
    
    Args:
        jsonl_path: JSONL文件路径
    
    Returns:
        包含nodes和edges的字典
    """
    nodes = {}  # {node_id: node_data}
    edges = []  # [{source, target, type, ...}]
    
    # 统计信息
    paper_count = 0
    author_count = 0
    reviewer_count = 0
    submission_edges = 0
    review_edges = 0
    
    # 用于去重和统计
    author_ids = set()
    reviewer_ids = set()
    paper_ids = set()
    
    print(f"正在解析文件: {jsonl_path}")
    
    with open(jsonl_path, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            if not line.strip():
                continue
            
            try:
                data = json.loads(line)
                
                # 提取论文信息
                paper_id = data.get('submission_id', f'paper_{line_num}')
                paper_title = data.get('title', f'Paper {data.get("submission_number", line_num)}')
                
                if paper_id not in paper_ids:
                    paper_ids.add(paper_id)
                    paper_count += 1
                    
                    # 添加论文节点
                    nodes[paper_id] = {
                        'id': paper_id,
                        'label': paper_title[:50] + ('...' if len(paper_title) > 50 else ''),
                        'type': 'paper',
                        'title': paper_title,
                        'submission_number': data.get('submission_number', line_num),
                        'group': 1  # 论文节点组
                    }
                
                # 处理作者（投稿关系）
                authors = data.get('authors', [])
                for author in authors:
                    author_id = author.get('id', '')
                    if not author_id:
                        continue
                    
                    if author_id not in author_ids:
                        author_ids.add(author_id)
                        author_count += 1
                        
                        # 添加作者节点
                        author_name = author.get('name') or author_id.replace('~', '').replace('_', ' ')
                        nodes[author_id] = {
                            'id': author_id,
                            'label': author_name,
                            'type': 'author',
                            'group': 2  # 作者节点组
                        }
                    
                    # 添加投稿边（作者 -> 论文）
                    edges.append({
                        'from': author_id,
                        'to': paper_id,
                        'type': 'submission',
                        'arrows': 'to',
                        'color': {'color': '#2B7CE9', 'highlight': '#2B7CE9'},
                        'width': 2
                    })
                    submission_edges += 1
                
                # 处理审稿人（审稿关系）
                reviews = data.get('reviews', [])
                for review in reviews:
                    reviewer_profile = review.get('reviewer_profile')
                    reviewer_id = ''
                    
                    if reviewer_profile and isinstance(reviewer_profile, dict):
                        reviewer_id = reviewer_profile.get('id', '')
                    
                    if not reviewer_id:
                        # 如果没有reviewer_profile，尝试从signature提取
                        signature = review.get('signature', '')
                        if signature:
                            reviewer_id = f'reviewer_{signature}'
                    
                    if reviewer_id:
                        if reviewer_id not in reviewer_ids:
                            reviewer_ids.add(reviewer_id)
                            reviewer_count += 1
                            
                            # 添加审稿人节点
                            if reviewer_profile and isinstance(reviewer_profile, dict):
                                reviewer_name = reviewer_profile.get('name') or reviewer_id.replace('~', '').replace('_', ' ')
                            else:
                                reviewer_name = reviewer_id.replace('~', '').replace('_', ' ')
                            
                            if reviewer_id not in nodes:
                                nodes[reviewer_id] = {
                                    'id': reviewer_id,
                                    'label': reviewer_name,
                                    'type': 'reviewer',
                                    'group': 3  # 审稿人节点组
                                }
                        
                        # 添加审稿边（审稿人 -> 论文）
                        rating = review.get('rating', {}).get('value', 'N/A')
                        # 提取评论内容
                        review_content = {
                            'rating': rating,
                            'summary': review.get('content', {}).get('summary', '') or review.get('summary', ''),
                            'strengths': review.get('content', {}).get('strengths', '') or review.get('strengths', ''),
                            'weaknesses': review.get('content', {}).get('weaknesses', '') or review.get('weaknesses', ''),
                            'questions': review.get('content', {}).get('questions', '') or review.get('questions', ''),
                            'confidence': review.get('content', {}).get('confidence', '') or review.get('confidence', ''),
                            'review_id': review.get('id', ''),
                            'reviewer_id': reviewer_id,
                            'paper_id': paper_id
                        }
                        # 合并所有评论文本用于搜索
                        all_text = ' '.join([
                            review_content['summary'],
                            review_content['strengths'],
                            review_content['weaknesses'],
                            review_content['questions']
                        ]).strip()
                        
                        edges.append({
                            'from': reviewer_id,
                            'to': paper_id,
                            'type': 'review',
                            'arrows': 'to',
                            'color': {'color': '#FF6B6B', 'highlight': '#FF6B6B'},
                            'width': 2,
                            'rating': rating,
                            'label': f'Rating: {rating}' if rating != 'N/A' else '',
                            'review_content': review_content,
                            'review_text': all_text  # 用于搜索的完整文本
                        })
                        review_edges += 1
                
            except json.JSONDecodeError as e:
                print(f"第 {line_num} 行JSON解析错误: {e}")
                continue
            except Exception as e:
                print(f"第 {line_num} 行处理错误: {e}")
                continue
    
    print(f"\n解析完成！")
    print(f"论文节点: {paper_count}")
    print(f"作者节点: {author_count}")
    print(f"审稿人节点: {reviewer_count}")
    print(f"投稿边: {submission_edges}")
    print(f"审稿边: {review_edges}")
    print(f"总节点数: {len(nodes)}")
    print(f"总边数: {len(edges)}")
    
    return {
        'nodes': list(nodes.values()),
        'edges': edges,
        'stats': {
            'papers': paper_count,
            'authors': author_count,
            'reviewers': reviewer_count,
            'submission_edges': submission_edges,
            'review_edges': review_edges
        }
    }


def save_graph_data(graph_data: Dict, output_path: str):
    """保存图数据为JSON文件"""
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(graph_data, f, ensure_ascii=False, indent=2)
    print(f"\n图数据已保存到: {output_path}")


if __name__ == '__main__':
    # 解析JSONL文件
    jsonl_path = r'C:\CODE\openiclr\iclr2026_reviews_10000(1).jsonl'
    graph_data = parse_jsonl_to_graph(jsonl_path)
    
    # 保存为JSON
    output_path = 'graph_data.json'
    save_graph_data(graph_data, output_path)
    
    print(f"\n数据准备完成！可以打开 index.html 查看可视化。")