src/data_loader.py

"""
AgentEA-Studio: Data loading and statistics module.
"""

import json
import numpy as np
from src.config import DATASETS, ALIGNMENT_RESULTS, SAMPLE_ALIGNMENTS, ABLATION_RESULTS


def get_dataset_list():
    """Return list of available datasets."""
    return list(DATASETS.keys())


def get_dataset_info(dataset_name):
    """Get dataset statistics as a formatted string."""
    ds = DATASETS.get(dataset_name)
    if not ds:
        return "Dataset not found."

    lines = [
        f"## {ds['name']} ({ds['type']})",
        "",
        "| Property | {kg1} | {kg2} |".format(kg1=ds['kg1'], kg2=ds['kg2']),
        "|---|---|---|",
        f"| Entities | {ds['entities_kg1']:,} | {ds['entities_kg2']:,} |",
        f"| Relations | {ds['relations_kg1']} | {ds['relations_kg2']} |",
        f"| Facts | {ds['facts_kg1']:,} | {ds['facts_kg2']:,} |",
        f"| Density | {ds['density_kg1']:.3f} | {ds['density_kg2']:.3f} |",
        "",
        "| Metric | Value |",
        "|---|---|",
        f"| Fact Count Difference | {ds['fact_diff_pct']}% |",
        f"| Density Difference | {ds['density_diff_pct']}% |",
        f"| Anchor Pairs | {ds['anchors']:,} |",
        f"| Entity Overlap | {ds['overlapping']}% |",
        f"| Structural Similarity | {ds['structural_similarity']}% |",
        f"| Temporal | {'Yes' if ds['temporal'] else 'No'} |",
    ]
    return "\n".join(lines)


def get_dataset_comparison():
    """Return a comparison table of all datasets."""
    datasets = [
        ("ICEWS-WIKI", "HHKG"),
        ("ICEWS-YAGO", "HHKG"),
        ("DBP15K-EN-FR", "Cross-lingual"),
        ("DBP-WIKI", "Heterogeneous"),
    ]

    header = "| Dataset | Type | Entities | Facts | Density | Overlap | Struct. Sim. |"
    sep = "|---|---|---|---|---|---|---|"

    rows = []
    for name, dtype in datasets:
        ds = DATASETS[name]
        entities = f"{ds['entities_kg1']:,} / {ds['entities_kg2']:,}"
        facts = f"{ds['facts_kg1']:,} / {ds['facts_kg2']:,}"
        density = f"{ds['density_kg1']:.1f} / {ds['density_kg2']:.1f}"
        overlap = f"{ds['overlapping']}%"
        ssim = f"{ds['structural_similarity']}%"
        rows.append(f"| {name} | {dtype} | {entities} | {facts} | {density} | {overlap} | {ssim} |")

    return "\n".join([header, sep] + rows)


def get_heterogeneity_explanation(dataset_name):
    """Explain why a dataset is highly heterogeneous."""
    ds = DATASETS.get(dataset_name)
    if not ds:
        return ""

    points = []
    if ds['fact_diff_pct'] > 100:
        points.append(f"- **Fact count disparity**: {ds['kg1']} has {ds['fact_diff_pct']}% more facts than {ds['kg2']}")
    if ds['density_diff_pct'] > 100:
        points.append(f"- **Density disparity**: {ds['kg1']} density is {ds['density_diff_pct']}% higher than {ds['kg2']}")
    if ds['overlapping'] < 100:
        points.append(f"- **Low entity overlap**: Only {ds['overlapping']}% of entities appear in both KGs")
    if ds['structural_similarity'] < 50:
        points.append(f"- **Low structural similarity**: Only {ds['structural_similarity']}% — vastly different graph structures")

    if not points:
        return "This dataset has relatively homogeneous KGs."

    return "\n".join(points)


def get_alignment_results(dataset_name):
    """Get pre-computed alignment results as a markdown table."""
    results = ALIGNMENT_RESULTS.get(dataset_name)
    if not results:
        return "No results available."

    header = "| Method | Hits@1 | Hits@5 | Hits@10 | MRR |"
    sep = "|---|---|---|---|---|"

    rows = []
    for method, metrics in results.items():
        h1 = f"{metrics['Hits@1']:.3f}"
        h5 = f"{metrics['Hits@5']:.3f}"
        h10 = f"{metrics['Hits@10']:.3f}"
        mrr = f"{metrics['MRR']:.3f}"
        rows.append(f"| {method} | {h1} | {h5} | {h10} | {mrr} |")

    return "\n".join([header, sep] + rows)


def get_sample_alignments(dataset_name):
    """Get sample alignment pairs as a markdown table."""
    samples = SAMPLE_ALIGNMENTS.get(dataset_name, [])
    if not samples:
        return "No samples available."

    header = "| # | Entity in KG1 | Entity in KG2 | Type | Confidence |"
    sep = "|---|---|---|---|---|"

    rows = []
    for i, s in enumerate(samples, 1):
        rows.append(f"| {i} | {s['entity_kg1']} | {s['entity_kg2']} | {s['type']} | {s['confidence']:.1%} |")

    return "\n".join([header, sep] + rows)


def get_ablation_results(dataset_name):
    """Get ablation study results as a markdown table."""
    results = ABLATION_RESULTS.get(dataset_name)
    if not results:
        return "No ablation results available."

    header = "| Configuration | Hits@1 | Hits@10 | MRR | Δ Hits@1 |"
    sep = "|---|---|---|---|---|"

    base_h1 = results["AdaCoAgentEA+"]["Hits@1"]
    rows = []
    for config, metrics in results.items():
        h1 = f"{metrics['Hits@1']:.3f}"
        h10 = f"{metrics['Hits@10']:.3f}"
        mrr = f"{metrics['MRR']:.3f}"
        delta = f"{metrics['Hits@1'] - base_h1:+.3f}"
        rows.append(f"| {config} | {h1} | {h10} | {mrr} | {delta} |")

    return "\n".join([header, sep] + rows)


def get_improvement_summary(dataset_name):
    """Get a summary of the improvement over best baseline."""
    results = ALIGNMENT_RESULTS.get(dataset_name)
    if not results:
        return ""

    ada_h1 = results["AdaCoAgentEA+"]["Hits@1"]
    best_baseline_h1 = 0
    best_baseline_name = ""

    for method, metrics in results.items():
        if method not in ["AdaCoAgentEA+", "AdaCoAgentEA"]:
            if metrics["Hits@1"] > best_baseline_h1:
                best_baseline_h1 = metrics["Hits@1"]
                best_baseline_name = method

    if best_baseline_h1 == 0:
        return ""

    relative_improvement = (ada_h1 - best_baseline_h1) / best_baseline_h1 * 100

    return (
        f"AdaCoAgentEA+ achieves **Hits@1 = {ada_h1:.1%}**, "
        f"improving over the best baseline **{best_baseline_name}** "
        f"(Hits@1 = {best_baseline_h1:.1%}) by **{relative_improvement:.1f}%** relative improvement."
    )


def get_agent_details(agent_id):
    """Get details for a specific agent."""
    from src.config import AGENTS

    for area_key, area in AGENTS.items():
        for agent in area["agents"]:
            if agent["id"] == agent_id:
                return f"""### {agent['icon']} {agent['name']} ({agent['id']})
**Type**: {agent['type']}
**Area**: {area['name']}

{agent['description']}
"""
    return "Agent not found."


def get_pipeline_step(step_num):
    """Get details for a specific pipeline step."""
    from src.config import PIPELINE_STEPS

    for step in PIPELINE_STEPS:
        if step["step"] == step_num:
            return f"""### Step {step['step']}: {step['name']}
**Area**: {step['area']}

{step['description']}

**Details**: {step['details']}
"""
    return "Step not found."