src/visualization.py

"""
AgentEA-Studio: Visualization helpers for Gradio.
"""

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
import io
import base64
from src.config import ALIGNMENT_RESULTS, ABLATION_RESULTS, AGENTS, PIPELINE_STEPS


def fig_to_base64(fig, dpi=150):
    """Convert matplotlib figure to base64 encoded PNG."""
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=dpi, bbox_inches="tight", facecolor="white")
    buf.seek(0)
    img_base64 = base64.b64encode(buf.read()).decode("utf-8")
    plt.close(fig)
    return img_base64


def plot_alignment_comparison(dataset_name):
    """Plot bar chart comparing alignment methods on a dataset."""
    results = ALIGNMENT_RESULTS.get(dataset_name)
    if not results:
        return None

    methods = list(results.keys())
    hits1 = [results[m]["Hits@1"] * 100 for m in methods]
    hits10 = [results[m]["Hits@10"] * 100 for m in methods]
    mrr = [results[m]["MRR"] * 100 for m in methods]

    x = np.arange(len(methods))
    width = 0.25

    fig, ax = plt.subplots(figsize=(12, 5))
    bars1 = ax.bar(x - width, hits1, width, label="Hits@1", color="#2196F3", edgecolor="white")
    bars2 = ax.bar(x, hits10, width, label="Hits@10", color="#4CAF50", edgecolor="white")
    bars3 = ax.bar(x + width, mrr, width, label="MRR", color="#FF9800", edgecolor="white")

    ax.set_ylabel("Score (%)", fontsize=12)
    ax.set_title(f"Alignment Performance on {dataset_name}", fontsize=14, fontweight="bold")
    ax.set_xticks(x)
    ax.set_xticklabels(methods, rotation=30, ha="right", fontsize=9)
    ax.legend(loc="lower right", fontsize=10)
    ax.set_ylim(0, 105)
    ax.grid(axis="y", alpha=0.3)

    # Highlight our methods
    for i, m in enumerate(methods):
        if "AdaCoAgentEA" in m:
            for bar in [bars1[i], bars2[i], bars3[i]]:
                bar.set_edgecolor("#D32F2F")
                bar.set_linewidth(2)

    plt.tight_layout()
    return fig_to_base64(fig)


def plot_ablation_study(dataset_name):
    """Plot ablation study results."""
    results = ABLATION_RESULTS.get(dataset_name)
    if not results:
        return None

    configs = list(results.keys())
    hits1 = [results[c]["Hits@1"] * 100 for c in configs]
    mrr = [results[c]["MRR"] * 100 for c in configs]

    fig, ax = plt.subplots(figsize=(10, 5))
    x = np.arange(len(configs))
    width = 0.35

    bars1 = ax.bar(x - width/2, hits1, width, label="Hits@1", color="#1565C0", edgecolor="white")
    bars2 = ax.bar(x + width/2, mrr, width, label="MRR", color="#E65100", edgecolor="white")

    ax.set_ylabel("Score (%)", fontsize=12)
    ax.set_title(f"Ablation Study on {dataset_name}", fontsize=14, fontweight="bold")
    ax.set_xticks(x)
    ax.set_xticklabels(configs, rotation=25, ha="right", fontsize=9)
    ax.legend(fontsize=10)
    ax.set_ylim(0, 105)
    ax.grid(axis="y", alpha=0.3)

    plt.tight_layout()
    return fig_to_base64(fig)


def plot_heterogeneity_radar():
    """Plot radar chart showing heterogeneity dimensions across datasets."""
    from src.config import DATASETS

    categories = ["Fact Disparity", "Density Disparity", "Entity Overlap", "Struct. Similarity"]
    datasets_to_plot = ["ICEWS-WIKI", "ICEWS-YAGO", "DBP15K-EN-FR", "DBP-WIKI"]

    N = len(categories)
    angles = [n / float(N) * 2 * np.pi for n in range(N)]
    angles += angles[:1]

    fig, ax = plt.subplots(figsize=(8, 8), subplot_kw=dict(polar=True))
    colors = ["#D32F2F", "#F44336", "#2196F3", "#4CAF50"]

    for i, ds_name in enumerate(datasets_to_plot):
        ds = DATASETS[ds_name]
        values = [
            min(ds["fact_diff_pct"] / 40, 100),
            min(ds["density_diff_pct"] / 35, 100),
            ds["overlapping"],
            ds["structural_similarity"],
        ]
        values += values[:1]
        ax.fill(angles, values, alpha=0.15, color=colors[i])
        ax.plot(angles, values, "o-", linewidth=2, label=ds_name, color=colors[i])

    ax.set_xticks(angles[:-1])
    ax.set_xticklabels(categories, fontsize=10)
    ax.set_ylim(0, 100)
    ax.set_title("KG Heterogeneity Radar", fontsize=14, fontweight="bold", pad=20)
    ax.legend(loc="upper right", bbox_to_anchor=(1.3, 1.1), fontsize=9)

    plt.tight_layout()
    return fig_to_base64(fig)


def plot_pipeline_flow():
    """Plot the multi-agent pipeline flow diagram."""
    fig, ax = plt.subplots(figsize=(16, 8))
    ax.set_xlim(0, 16)
    ax.set_ylim(0, 8)
    ax.axis("off")

    # Colors
    color_input = "#E3F2FD"
    color_area1 = "#C8E6C9"
    color_area2 = "#FFF9C4"
    color_area3 = "#FFCCBC"
    color_gating = "#E1BEE7"
    color_output = "#BBDEFB"

    # Draw boxes
    boxes = [
        (0.5, 5.5, 2.5, 1.5, "Input\nHHKGs", color_input),
        (3.5, 5.5, 4.0, 1.5, "Area 1\nStructural-Semantic\nDecoupling", color_area1),
        (3.5, 3.0, 4.0, 0.8, "Gating G1", color_gating),
        (8.0, 5.5, 4.0, 1.5, "Area 2\nDeep Localization", color_area2),
        (8.0, 3.0, 4.0, 0.8, "Gating G2", color_gating),
        (12.5, 5.5, 3.0, 1.5, "Area 3\nKnowledge\nAggregation", color_area3),
        (12.5, 3.0, 3.0, 0.8, "Gating G3", color_gating),
        (0.5, 1.0, 2.5, 1.0, "Output\nAligned Pairs", color_output),
    ]

    for x, y, w, h, label, color in boxes:
        rect = mpatches.FancyBboxPatch(
            (x, y), w, h, boxstyle="round,pad=0.1",
            facecolor=color, edgecolor="gray", linewidth=1.5
        )
        ax.add_patch(rect)
        ax.text(x + w/2, y + h/2, label, ha="center", va="center", fontsize=9, fontweight="bold")

    # Draw arrows
    arrows = [
        (3.0, 6.25, 3.5, 6.25),  # Input → Area 1
        (7.5, 6.25, 8.0, 6.25),  # Area 1 → Area 2
        (12.0, 6.25, 12.5, 6.25),  # Area 2 → Area 3
        (14.0, 5.5, 14.0, 3.8),  # Area 3 → G3
        (10.0, 3.0, 10.0, 3.8),  # G2 → Area 2
        (5.5, 3.0, 5.5, 3.8),  # G1 → Area 1
        (7.5, 3.4, 12.5, 3.4),  # G2 → G3
        (3.5, 3.4, 0.5, 3.4),  # G1 → Output
        (10.0, 3.4, 0.5, 3.4),  # G2 → Output
    ]

    for x1, y1, x2, y2 in arrows:
        ax.annotate("", xy=(x2, y2), xytext=(x1, y1),
                    arrowprops=dict(arrowstyle="->", color="gray", lw=1.5))

    ax.set_title("AdaCoAgentEA Multi-Agent Pipeline Architecture", fontsize=16, fontweight="bold", pad=10)

    # Legend
    legend_elements = [
        mpatches.Patch(facecolor=color_input, edgecolor="gray", label="Input/Output"),
        mpatches.Patch(facecolor=color_area1, edgecolor="gray", label="Area 1: Structural-Semantic"),
        mpatches.Patch(facecolor=color_area2, edgecolor="gray", label="Area 2: Deep Localization"),
        mpatches.Patch(facecolor=color_area3, edgecolor="gray", label="Area 3: Knowledge Aggregation"),
        mpatches.Patch(facecolor=color_gating, edgecolor="gray", label="Gating Mechanism"),
    ]
    ax.legend(handles=legend_elements, loc="lower center", ncol=5, fontsize=8,
              bbox_to_anchor=(0.5, -0.08))

    plt.tight_layout()
    return fig_to_base64(fig)


def plot_agent_communication():
    """Plot the agent communication protocol diagram."""
    fig, ax = plt.subplots(figsize=(14, 7))
    ax.set_xlim(0, 14)
    ax.set_ylim(0, 7)
    ax.axis("off")

    # Shared Message Pool
    pool = mpatches.FancyBboxPatch(
        (4, 2), 6, 4, boxstyle="round,pad=0.15",
        facecolor="#E8EAF6", edgecolor="#3F51B5", linewidth=2
    )
    ax.add_patch(pool)
    ax.text(7, 5.5, "Shared Message Pool", ha="center", va="center", fontsize=13, fontweight="bold", color="#1A237E")

    # Agents
    agents = [
        (1, 5, "S1-S3", "Small Models\nArea 1", "#C8E6C9"),
        (1, 3, "L1", "LLM\nArea 2", "#FFF9C4"),
        (1, 1, "S4", "Small Model\nArea 2", "#C8E6C9"),
        (12, 5, "L2", "LLM\nArea 2", "#FFF9C4"),
        (12, 3, "L3", "Master\nArea 3", "#FFCCBC"),
        (12, 1, "L4+", "Domain Experts\nArea 3", "#FFCCBC"),
    ]

    for x, y, label, desc, color in agents:
        rect = mpatches.FancyBboxPatch(
            (x - 0.8, y - 0.5), 1.6, 1.0, boxstyle="round,pad=0.05",
            facecolor=color, edgecolor="gray", linewidth=1
        )
        ax.add_patch(rect)
        ax.text(x, y + 0.1, label, ha="center", va="center", fontsize=10, fontweight="bold")
        ax.text(x, y - 0.25, desc, ha="center", va="center", fontsize=7, color="gray")

    # Communication arrows
    for x, y, label, desc, color in agents:
        ax.annotate("", xy=(4.5, y), xytext=(x + 0.8, y),
                    arrowprops=dict(arrowstyle="<->", color="#7986CB", lw=1.2, ls="--"))

    ax.text(7, 1.2, "Structured Agent Communication + Publish-Subscribe Mechanism",
            ha="center", fontsize=10, fontstyle="italic", color="gray")

    ax.set_title("Agent Communication Protocol", fontsize=14, fontweight="bold", pad=10)
    plt.tight_layout()
    return fig_to_base64(fig)


def plot_improvement_waterfall(dataset_name):
    """Plot waterfall chart showing improvement contributions."""
    results = ABLATION_RESULTS.get(dataset_name)
    if not results:
        return None

    base = results["w/o LLM Agents"]["Hits@1"] * 100
    ada = results["AdaCoAgentEA+"]["Hits@1"] * 100

    components = [
        ("Base (w/o LLM)", base, base),
        ("+ LLM Agents", results["w/o Communication Protocols"]["Hits@1"] * 100 - base,
         results["w/o Communication Protocols"]["Hits@1"] * 100),
        ("+ Communication Protocols", results["w/o Expert Role-Playing"]["Hits@1"] * 100 -
         results["w/o Communication Protocols"]["Hits@1"] * 100,
         results["w/o Expert Role-Playing"]["Hits@1"] * 100),
        ("+ Expert Role-Playing", results["w/o Meta-Logic Rules"]["Hits@1"] * 100 -
         results["w/o Expert Role-Playing"]["Hits@1"] * 100,
         results["w/o Meta-Logic Rules"]["Hits@1"] * 100),
        ("+ Meta-Logic Rules", ada - results["w/o Meta-Logic Rules"]["Hits@1"] * 100,
         ada),
    ]

    fig, ax = plt.subplots(figsize=(10, 5))

    labels = [c[0] for c in components]
    values = [c[1] for c in components]
    cumulative = [c[2] for c in components]
    colors = ["#90A4AE", "#42A5F5", "#66BB6A", "#FFA726", "#EF5350"]

    # Plot bars
    bottom = base
    for i, (label, val, cum) in enumerate(components):
        if i == 0:
            ax.bar(i, val, color=colors[i], edgecolor="white")
        else:
            ax.bar(i, val, bottom=bottom, color=colors[i], edgecolor="white")
            bottom += val

    ax.set_xticks(range(len(components)))
    ax.set_xticklabels(labels, rotation=20, ha="right", fontsize=9)
    ax.set_ylabel("Hits@1 (%)", fontsize=12)
    ax.set_title(f"Component Contribution Waterfall ({dataset_name})", fontsize=14, fontweight="bold")
    ax.grid(axis="y", alpha=0.3)

    # Add value labels
    for i, (label, val, cum) in enumerate(components):
        ax.text(i, cum + 1, f"{cum:.1f}%", ha="center", fontsize=9, fontweight="bold")

    plt.tight_layout()
    return fig_to_base64(fig)