src/data_generator.py

"""
AgentEA-Studio: Real KG data generator.
Generates realistic synthetic KG data for demo purposes,
matching the statistics of DBP15K, ICEWS-WIKI, etc.
"""

import random
import json
import os
from collections import defaultdict

# Entity name pools for generating realistic KGs
PERSON_NAMES = [
    "Barack Obama", "Joe Biden", "Donald Trump", "Angela Merkel", "Emmanuel Macron",
    "Xi Jinping", "Vladimir Putin", "Boris Johnson", "Justin Trudeau", "Narendra Modi",
    "Mahmoud Abbas", "Benjamin Netanyahu", "Emmanuel Macron", "Ursula von der Leyen",
    "Antonio Guterres", "Steffi Lemke", "Olaf Scholz", "Yoshihide Suga",
    "Moon Jae-in", "Scott Morrison", "Jacinda Ardern", "Pedro Sanchez",
    "Mario Draghi", "Andrzej Duda", "Recep Erdogan", "Jair Bolsonaro",
    "Cyril Ramaphosa", "Alberto Fernandez", "Mette Frederiksen", "Sanna Marin",
]

COUNTRY_NAMES = [
    "United States", "China", "Russia", "Germany", "France", "United Kingdom",
    "Japan", "India", "Brazil", "Canada", "Australia", "Italy", "Spain",
    "South Korea", "Mexico", "Indonesia", "Turkey", "Netherlands", "Switzerland",
    "Sweden", "Poland", "Belgium", "Norway", "Austria", "Denmark", "Finland",
    "Portugal", "Greece", "Ireland", "New Zealand", "Singapore", "Malaysia",
    "Thailand", "Vietnam", "Philippines", "Egypt", "South Africa", "Nigeria",
    "Kenya", "Argentina", "Chile", "Colombia", "Peru", "Saudi Arabia", "UAE",
    "Qatar", "Israel", "Iran", "Pakistan", "Bangladesh",
]

ORG_NAMES = [
    "United Nations", "NATO", "European Union", "World Bank", "IMF",
    "World Health Organization", "UNESCO", "OPEC", "ASEAN", "African Union",
    "Red Cross", "Greenpeace", "Amnesty International", "World Trade Organization",
    "International Criminal Court", "OECD", "G20", "G7", "BRICS",
    "World Economic Forum", "International Olympic Committee", "FIFA",
    "Apple Inc.", "Google", "Microsoft", "Amazon", "Meta", "Tesla",
    "Samsung", "Toyota", "Volkswagen", "Shell", "ExxonMobil", "Pfizer",
    "Johnson & Johnson", "Roche", "Novartis", "Siemens", "IBM", "Intel",
    "NVIDIA", "AMD", "Qualcomm", "Huawei", "Alibaba", "Tencent", "Baidu",
]

CITY_NAMES = [
    "New York", "London", "Paris", "Tokyo", "Beijing", "Shanghai", "Moscow",
    "Berlin", "Washington D.C.", "Brussels", "Geneva", "Vienna", "Rome",
    "Madrid", "Lisbon", "Seoul", "Sydney", "Toronto", "Mumbai", "Delhi",
    "Dubai", "Singapore", "Hong Kong", "San Francisco", "Chicago", "Boston",
    "Los Angeles", "Houston", "Miami", "Seattle", "Austin", "Denver",
    "Cairo", "Lagos", "Nairobi", "Cape Town", "Buenos Aires", "Sao Paulo",
    "Mexico City", "Lima", "Bogota", "Santiago", "Bangkok", "Jakarta",
    "Manila", "Hanoi", "Kuala Lumpur", "Taipei", "Osaka", "Kyoto", "Zurich",
    "Stockholm", "Oslo", "Copenhagen", "Helsinki", "Warsaw", "Prague",
    "Budapest", "Athens", "Dublin", "Amsterdam", "Barcelona", "Munich",
    "Frankfurt", "Hamburg", "Milan", "Naples", "Venice", "Florence",
]

EVENT_NAMES = [
    "World War II", "Cold War", "Arab Spring", "Brexit", "COVID-19 Pandemic",
    "2008 Financial Crisis", "Paris Agreement", "Kyoto Protocol", "Iraq War",
    "Syrian Civil War", "Russo-Ukrainian War", "Eurozone Crisis",
    "Belt and Road Initiative", "Trans-Pacific Partnership", "NAFTA",
    "Space Race", "Moon Landing", "Chernobyl Disaster", "Fukushima Disaster",
    "Hurricane Katrina", "Indian Ocean Tsunami", "Haiti Earthquake",
]

# Relation types
RELATIONS = [
    "president_of", "prime_minister_of", "located_in", "capital_of",
    "member_of", "part_of", "founded_by", "headquartered_in",
    "has_leader", "borders_with", "trades_with", "allied_with",
    "sanctions_against", "negotiates_with", "attended_by",
    "born_in", "died_in", "educated_at", "works_for", "owns",
    "subsidiary_of", "competitor_of", "supplier_to", "investor_in",
    "preceded_by", "followed_by", "caused_by", "resulted_in",
    "announced_at", "signed_at", " ratified_by", "participated_in",
    "voted_for", "vetoed_by", "proposed_by", "implemented_by",
    "reported_by", "published_in", "broadcast_by", "regulated_by",
    "funded_by", "audited_by", "certified_by", "accredited_by",
]

# Entity types
ENTITY_TYPES = ["Person", "Country", "City", "Organization", "Event", "Other"]

# Word pools for generating diverse entity names in cross-lingual KGs
FR_WORDS = ["président", "république", "organisation", "traité", "accord",
            "gouvernement", "ministre", "ambassadeur", "conseil", "commission"]
ZH_TRANSLATIONS = {
    "United States": "美国", "China": "中国", "Russia": "俄罗斯",
    "Germany": "德国", "France": "法国", "United Kingdom": "英国",
    "Japan": "日本", "India": "印度", "Brazil": "巴西", "Canada": "加拿大",
}


class KGGenerator:
    """Generate realistic synthetic KG data for entity alignment demo."""

    def __init__(self, seed=42):
        random.seed(seed)
        self.entity_counter = 0
        self.rel_counter = 0

    def _make_entity(self, name, etype, kg_prefix=""):
        eid = f"{kg_prefix}E{self.entity_counter}"
        self.entity_counter += 1
        return {"id": eid, "name": name, "type": etype}

    def _make_relation(self, name):
        rid = f"R{self.rel_counter}"
        self.rel_counter += 1
        return {"id": rid, "name": name}

    def generate_db15k_style(self, n_entities=5000, n_shared=5000):
        """Generate a DBP15K-style cross-lingual KG pair (EN-FR style).
        Two KGs with same entities but different languages/naming.
        """
        entities_kg1 = []
        entities_kg2 = []
        relations = self._build_relation_set()
        triples_kg1 = []
        triples_kg2 = []
        alignments = []

        # Create shared entities
        name_pools = [PERSON_NAMES, COUNTRY_NAMES, ORG_NAMES, CITY_NAMES]
        type_map = ["Person", "Country", "Organization", "City"]

        for i in range(n_shared):
            pool_idx = i % len(name_pools)
            name = random.choice(name_pools[pool_idx])
            etype = type_map[pool_idx]

            e1 = self._make_entity(name, etype, "EN_")
            e2 = self._make_entity(self._frenchify(name), etype, "FR_")
            entities_kg1.append(e1)
            entities_kg2.append(e2)
            alignments.append((e1["id"], e2["id"]))

        # Add KG-specific entities
        for i in range(n_entities - n_shared):
            etype = random.choice(type_map)
            pool = name_pools[type_map.index(etype)]
            name = random.choice(pool)
            e1 = self._make_entity(name, etype, "EN_")
            entities_kg1.append(e1)

        for i in range(n_entities - n_shared):
            etype = random.choice(type_map)
            pool = name_pools[type_map.index(etype)]
            name = random.choice(pool)
            e2 = self._make_entity(self._frenchify(name), etype, "FR_")
            entities_kg2.append(e2)

        # Generate triples
        triples_kg1 = self._generate_triples(entities_kg1, relations, avg_degree=6)
        triples_kg2 = self._generate_triples(entities_kg2, relations, avg_degree=6)

        return {
            "kg1": {"entities": entities_kg1, "relations": relations, "triples": triples_kg1},
            "kg2": {"entities": entities_kg2, "relations": relations, "triples": triples_kg2},
            "alignments": alignments,
        }

    def generate_hhkg_style(self, n_entities_kg1=3000, n_entities_kg2=5000, n_shared=1500):
        """Generate an HHKG-style KG pair (ICEWS-WIKI style).
        Two KGs with very different scales, densities, and structures.
        KG1 is dense (event-focused), KG2 is sparse (encyclopedic).
        """
        entities_kg1 = []
        entities_kg2 = []
        relations = self._build_relation_set()
        triples_kg1 = []
        triples_kg2 = []
        alignments = []

        name_pools = [PERSON_NAMES, COUNTRY_NAMES, ORG_NAMES, CITY_NAMES, EVENT_NAMES]
        type_map = ["Person", "Country", "Organization", "City", "Event"]

        # Shared entities
        for i in range(n_shared):
            pool_idx = i % min(len(name_pools), 4)  # Keep shared entities to common types
            name = random.choice(name_pools[pool_idx])
            etype = type_map[pool_idx]

            e1 = self._make_entity(name, etype, "KG1_")
            e2 = self._make_entity(f"{name} (wiki)", etype, "KG2_")
            entities_kg1.append(e1)
            entities_kg2.append(e2)
            alignments.append((e1["id"], e2["id"]))

        # KG1-specific entities (dense, event-focused)
        for i in range(n_entities_kg1 - n_shared):
            etype = random.choice(["Event", "Person", "Country"])
            pool = name_pools[type_map.index(etype)]
            name = random.choice(pool)
            e1 = self._make_entity(name, etype, "KG1_")
            entities_kg1.append(e1)

        # KG2-specific entities (sparse, encyclopedic)
        for i in range(n_entities_kg2 - n_shared):
            etype = random.choice(type_map)
            pool = name_pools[type_map.index(etype)]
            name = random.choice(pool)
            e2 = self._make_entity(f"{name} (wiki)", etype, "KG2_")
            entities_kg2.append(e2)

        # KG1: very dense (many triples per entity)
        triples_kg1 = self._generate_triples(entities_kg1, relations, avg_degree=25)
        # KG2: sparse (few triples per entity)
        triples_kg2 = self._generate_triples(entities_kg2, relations, avg_degree=4)

        return {
            "kg1": {"entities": entities_kg1, "relations": relations, "triples": triples_kg1},
            "kg2": {"entities": entities_kg2, "relations": relations, "triples": triples_kg2},
            "alignments": alignments,
        }

    def _build_relation_set(self):
        """Build a relation vocabulary."""
        relations = []
        for rel_name in RELATIONS:
            relations.append(self._make_relation(rel_name))
        return relations

    def _generate_triples(self, entities, relations, avg_degree=6):
        """Generate realistic triples connecting entities."""
        triples = []
        n_entities = len(entities)
        total_triples = int(n_entities * avg_degree)

        # Build relation index
        rel_by_name = {r["name"]: r for r in relations}

        # Assign types
        entities_by_type = defaultdict(list)
        for e in entities:
            entities_by_type[e["type"]].append(e)

        for _ in range(total_triples):
            head = random.choice(entities)
            # Choose relation appropriate for the head entity type
            etype = head["type"]

            if etype == "Person":
                possible_rels = ["president_of", "prime_minister_of", "born_in",
                                 "educated_at", "works_for", "has_leader",
                                 "attended_by", "participated_in"]
                tail_pool = entities_by_type["Country"] + entities_by_type["City"] + \
                           entities_by_type["Organization"] + entities_by_type["Event"]
            elif etype == "Country":
                possible_rels = ["located_in", "capital_of", "member_of", "borders_with",
                                 "trades_with", "allied_with", "sanctions_against",
                                 "negotiates_with", "part_of", "has_leader"]
                tail_pool = entities_by_type["Country"] + entities_by_type["City"] + \
                           entities_by_type["Organization"] + entities_by_type["Person"]
            elif etype == "City":
                possible_rels = ["located_in", "capital_of", "part_of",
                                 "headquartered_in", "born_in", "died_in"]
                tail_pool = entities_by_type["Country"] + entities_by_type["Organization"] + \
                           entities_by_type["Person"]
            elif etype == "Organization":
                possible_rels = ["headquartered_in", "member_of", "part_of", "founded_by",
                                 "subsidiary_of", "competitor_of", "funded_by",
                                 "regulated_by", "participated_in"]
                tail_pool = entities_by_type["City"] + entities_by_type["Country"] + \
                           entities_by_type["Organization"] + entities_by_type["Person"]
            elif etype == "Event":
                possible_rels = ["attended_by", "participated_in", "caused_by",
                                 "resulted_in", "announced_at", "signed_at",
                                 "reported_by", "preceded_by", "followed_by"]
                tail_pool = entities_by_type["Person"] + entities_by_type["Country"] + \
                           entities_by_type["Organization"] + entities_by_type["City"] + \
                           entities_by_type["Event"]
            else:
                possible_rels = list(RELATIONS)
                tail_pool = entities

            if not tail_pool:
                continue

            rel_name = random.choice([r for r in possible_rels if r in rel_by_name])
            tail = random.choice(tail_pool)

            if tail["id"] != head["id"]:
                triples.append({
                    "head": head["id"],
                    "relation": rel_by_name[rel_name]["id"],
                    "tail": tail["id"],
                })

        return triples

    def _frenchify(self, name):
        """Create a French-like version of an English entity name."""
        if name in ZH_TRANSLATIONS:
            return name
        parts = name.split()
        if len(parts) > 1:
            return f"{parts[-1]}, {' '.join(parts[:-1])}"
        return name

    def save_dataset(self, dataset, output_dir):
        """Save a generated dataset to disk."""
        os.makedirs(output_dir, exist_ok=True)

        for kg_name in ["kg1", "kg2"]:
            kg = dataset[kg_name]
            kg_dir = os.path.join(output_dir, kg_name)
            os.makedirs(kg_dir, exist_ok=True)

            # Save entities
            with open(os.path.join(kg_dir, "entities.json"), "w") as f:
                json.dump(kg["entities"], f, indent=2)

            # Save relations
            with open(os.path.join(kg_dir, "relations.json"), "w") as f:
                json.dump(kg["relations"], f, indent=2)

            # Save triples
            with open(os.path.join(kg_dir, "triples.json"), "w") as f:
                json.dump(kg["triples"], f, indent=2)

            # Save triples as TSV for easy loading
            with open(os.path.join(kg_dir, "triples.tsv"), "w") as f:
                for t in kg["triples"]:
                    f.write(f"{t['head']}\t{t['relation']}\t{t['tail']}\n")

        # Save alignments
        with open(os.path.join(output_dir, "alignments.json"), "w") as f:
            json.dump(dataset["alignments"], f, indent=2)


def generate_all_datasets(output_dir):
    """Generate all demo datasets."""
    gen = KGGenerator(seed=42)

    print("Generating DBP15K-style dataset...")
    db15k = gen.generate_db15k_style(n_entities=3000, n_shared=3000)
    gen.save_dataset(db15k, os.path.join(output_dir, "db15k_demo"))

    print("Generating HHKG-style dataset (ICEWS-WIKI style)...")
    gen2 = KGGenerator(seed=123)
    hhkg = gen2.generate_hhkg_style(n_entities_kg1=2000, n_entities_kg2=3500, n_shared=1000)
    gen2.save_dataset(hhkg, os.path.join(output_dir, "hhkg_demo"))

    print("Done!")

    # Print statistics
    for name, ds in [("DBP15K-Demo", db15k), ("HHKG-Demo", hhkg)]:
        kg1 = ds["kg1"]
        kg2 = ds["kg2"]
        print(f"\n{name}:")
        print(f"  KG1: {len(kg1['entities'])} entities, {len(kg1['triples'])} triples, "
              f"density={len(kg1['triples'])/len(kg1['entities']):.1f}")
        print(f"  KG2: {len(kg2['entities'])} entities, {len(kg2['triples'])} triples, "
              f"density={len(kg2['triples'])/len(kg2['entities']):.1f}")
        print(f"  Alignments: {len(ds['alignments'])}")


if __name__ == "__main__":
    output_dir = os.path.join(os.path.dirname(__file__), "..", "data")
    generate_all_datasets(output_dir)