"""Generate synthetic research paper dataset for the hybrid demo notebook.

Creates:
- 100 papers with titles, abstracts, topics
- Citation graph (papers cite other papers)
- 384-dim abstract embeddings (random but topic-clustered)

Output: papers.json with all data ready for loading.
"""

import json
import random
import math

random.seed(42)

TOPICS = [
    "graph neural networks",
    "natural language processing",
    "computer vision",
    "reinforcement learning",
    "generative models",
]

TOPIC_TITLES = {
    "graph neural networks": [
        "Graph Attention Networks for Node Classification",
        "Message Passing Neural Networks: A Survey",
        "Spectral Graph Convolutions Revisited",
        "Heterogeneous Graph Transformers",
        "Graph Neural Networks for Molecular Property Prediction",
        "Inductive Representation Learning on Large Graphs",
        "Dynamic Graph Neural Networks",
        "Subgraph Sampling for Graph Neural Networks",
        "Graph Contrastive Learning Framework",
        "Temporal Graph Networks for Dynamic Interactions",
        "Scalable Graph Neural Networks via Sampling",
        "Graph Isomorphism Networks",
        "Relational Graph Convolutional Networks",
        "Multi-Scale Graph Neural Networks",
        "Graph Neural Networks for Recommendation Systems",
        "Knowledge Graph Embeddings via GNNs",
        "Expressive Power of Graph Neural Networks",
        "Graph Neural Architecture Search",
        "Equivariant Graph Neural Networks",
        "Over-Smoothing in Deep Graph Networks",
    ],
    "natural language processing": [
        "Attention Is All You Need: Revisiting Transformers",
        "BERT Pre-training for Language Understanding",
        "Large Language Models as Few-Shot Learners",
        "Retrieval Augmented Generation for Knowledge Tasks",
        "Chain of Thought Prompting in Large Models",
        "Efficient Transformers: A Comprehensive Survey",
        "Multilingual Transfer Learning",
        "Named Entity Recognition with Transformers",
        "Semantic Textual Similarity Benchmarks",
        "Question Answering over Knowledge Graphs",
        "Text Summarization with Pre-trained Models",
        "Dialogue Systems with Large Language Models",
        "Machine Translation Quality Estimation",
        "Document Understanding with Layout Models",
        "Sentiment Analysis in Low-Resource Languages",
        "Code Generation with Language Models",
        "Instruction Tuning for Language Models",
        "Constitutional AI Training Methods",
        "Hallucination Detection in Language Models",
        "Long Context Window Transformers",
    ],
    "computer vision": [
        "Vision Transformers for Image Recognition",
        "Self-Supervised Learning for Visual Representations",
        "Object Detection with Anchor-Free Methods",
        "Semantic Segmentation with Diffusion Models",
        "Video Understanding with Temporal Attention",
        "3D Object Detection from Point Clouds",
        "Neural Radiance Fields for Novel View Synthesis",
        "Image Generation with Diffusion Processes",
        "Multi-Modal Vision-Language Models",
        "Depth Estimation from Single Images",
        "Medical Image Segmentation Networks",
        "Real-Time Instance Segmentation",
        "Visual Question Answering Systems",
        "Image Super-Resolution with GANs",
        "Contrastive Learning for Visual Features",
        "Pose Estimation with Graph Networks",
        "Scene Graph Generation Methods",
        "Visual Grounding in Natural Language",
        "Domain Adaptation for Visual Recognition",
        "Efficient Neural Architecture for Mobile Vision",
    ],
    "reinforcement learning": [
        "Model-Based Reinforcement Learning Survey",
        "Multi-Agent Reinforcement Learning",
        "Offline Reinforcement Learning Methods",
        "Reward Shaping for Sparse Environments",
        "Hierarchical Reinforcement Learning",
        "Safe Reinforcement Learning with Constraints",
        "Meta-Reinforcement Learning Algorithms",
        "Reinforcement Learning from Human Feedback",
        "Distributional Reinforcement Learning",
        "Curriculum Learning for RL Agents",
        "Exploration Strategies in RL",
        "Inverse Reinforcement Learning",
        "Policy Gradient Methods Revisited",
        "World Models for Planning",
        "Multi-Task Reinforcement Learning",
        "Robust Reinforcement Learning",
        "Reinforcement Learning for Robotics",
        "Game-Playing AI with Deep RL",
        "RL for Combinatorial Optimization",
        "Sample-Efficient RL Algorithms",
    ],
    "generative models": [
        "Denoising Diffusion Probabilistic Models",
        "Score-Based Generative Modeling",
        "Variational Autoencoders with Improved Priors",
        "Generative Adversarial Networks: Recent Advances",
        "Flow-Based Generative Models",
        "Text-to-Image Generation with Diffusion",
        "Conditional Generation with Classifier Guidance",
        "Audio Generation with Neural Networks",
        "Video Generation with Temporal Diffusion",
        "3D Shape Generation Methods",
        "Molecular Generation with Deep Learning",
        "Music Generation with Transformers",
        "Layout Generation for Documents",
        "Graph Generation Methods",
        "Protein Structure Generation",
        "Latent Diffusion Models",
        "Consistency Models for Fast Generation",
        "Controllable Text Generation",
        "Image Editing with Generative Models",
        "Few-Shot Generation Methods",
    ],
}

DIM = 384


def make_topic_vector(topic_idx: int) -> list[float]:
    """Generate a random vector clustered around a topic centroid."""
    vec = [random.gauss(0, 0.3) for _ in range(DIM)]
    # Boost dimensions associated with this topic
    start = topic_idx * (DIM // len(TOPICS))
    end = start + (DIM // len(TOPICS))
    for i in range(start, end):
        vec[i] += random.gauss(1.5, 0.3)
    # Normalize
    norm = math.sqrt(sum(v * v for v in vec))
    return [v / norm for v in vec]


def main():
    papers = []
    paper_id = 1

    for topic_idx, topic in enumerate(TOPICS):
        titles = TOPIC_TITLES[topic]
        for title in titles:
            paper = {
                "id": paper_id,
                "title": title,
                "topic": topic,
                "year": random.randint(2020, 2025),
                "citations": random.randint(0, 500),
                "embedding": make_topic_vector(topic_idx),
            }
            papers.append(paper)
            paper_id += 1

    # Generate citation graph: papers tend to cite papers in same/related topics
    edges = []
    for paper in papers:
        # Each paper cites 2-5 other papers
        num_citations = random.randint(2, 5)
        candidates = [p for p in papers if p["id"] != paper["id"]]
        # Weight toward same topic
        weights = []
        for c in candidates:
            if c["topic"] == paper["topic"]:
                weights.append(3.0)
            else:
                weights.append(1.0)

        cited = random.choices(candidates, weights=weights, k=num_citations)
        for c in cited:
            edges.append({"from_id": paper["id"], "to_id": c["id"]})

    # Deduplicate edges
    seen = set()
    unique_edges = []
    for e in edges:
        key = (e["from_id"], e["to_id"])
        if key not in seen:
            seen.add(key)
            unique_edges.append(e)

    data = {"papers": papers, "citations": unique_edges}

    with open("papers.json", "w") as f:
        json.dump(data, f)

    print(f"Generated {len(papers)} papers with {len(unique_edges)} citations")


if __name__ == "__main__":
    main()
