# Author - Deepak Soni
"""Generate architecture + performance figures for the Antahkarana-base model card. Run locally (matplotlib)."""
import os
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
A = os.path.join(HERE, "assets"); os.makedirs(A, exist_ok=True)
ORG = {"manas": "#4C9BE8", "buddhi": "#7B61FF", "ahamkara": "#E8804C", "citta": "#2BB673"}
NAVY = "#1b2a4a"

# ---------------------------------------------------------------- 1. architecture
fig, ax = plt.subplots(figsize=(11, 7.2)); ax.set_xlim(0, 11); ax.set_ylim(0, 7.2); ax.axis("off")
ax.text(5.5, 6.85, "Antaḥkaraṇa-base", fontsize=22, fontweight="bold", ha="center", color=NAVY)
ax.text(5.5, 6.45, "a backbone-agnostic continual-learning mind  ·  four organs, one loop", fontsize=11,
        ha="center", color="#555", style="italic")

# the core box
ax.add_patch(FancyBboxPatch((0.4, 3.0), 10.2, 3.0, boxstyle="round,pad=0.05,rounding_size=0.15",
                            fc="#f3f5fb", ec=NAVY, lw=1.8))
ax.text(5.5, 5.72, "ANTAHKARANA CORE  (the loop is identical for every domain)", fontsize=10.5,
        ha="center", fontweight="bold", color=NAVY)
organs = [("manas", "intake · avidyā\n(energy-OOD novelty)"),
          ("buddhi", "decide · pramāṇa\n(abstain) · guṇa · viveka"),
          ("ahamkara", "identity ·\nboundary detection"),
          ("citta", "memory · saṃskāra ·\nreplay (DER++) · nidrā")]
for i, (name, desc) in enumerate(organs):
    x = 0.85 + i * 2.45
    ax.add_patch(FancyBboxPatch((x, 3.35), 2.2, 1.85, boxstyle="round,pad=0.04,rounding_size=0.12",
                                fc=ORG[name], ec="none", alpha=0.92))
    ax.text(x + 1.1, 4.92, name, fontsize=13, fontweight="bold", ha="center", color="white")
    ax.text(x + 1.1, 4.05, desc, fontsize=8.3, ha="center", color="white")

# adapter + backbone
ax.add_patch(FancyBboxPatch((3.0, 1.7), 5.0, 0.85, boxstyle="round,pad=0.04,rounding_size=0.1",
                            fc="#ffd166", ec="#b8860b", lw=1.4))
ax.text(5.5, 2.12, "BackboneAdapter   (5 methods — the ONLY thing a new domain writes)", fontsize=9.5,
        ha="center", fontweight="bold", color="#5a4500")
for i, (lab, col) in enumerate([("Language\nMistral-7B (LLM)", "#4C9BE8"),
                                ("Vision\nViT-B/16", "#7B61FF"),
                                ("Security\nIDS encoder", "#E8804C"),
                                ("…any domain\n(coding/audio/…)", "#888")]):
    x = 0.9 + i * 2.45
    ax.add_patch(FancyBboxPatch((x, 0.35), 2.05, 0.95, boxstyle="round,pad=0.03,rounding_size=0.1",
                                fc="white", ec=col, lw=1.8))
    ax.text(x + 1.02, 0.82, lab, fontsize=8.6, ha="center", color=col, fontweight="bold")
for a in (5.5,):
    ax.add_patch(FancyArrowPatch((a, 3.0), (a, 2.56), arrowstyle="-|>", mutation_scale=14, color=NAVY, lw=1.4))
    ax.add_patch(FancyArrowPatch((a, 1.7), (a, 1.32), arrowstyle="-|>", mutation_scale=14, color="#b8860b", lw=1.4))
plt.tight_layout(); plt.savefig(os.path.join(A, "architecture.png"), dpi=150, bbox_inches="tight"); plt.close()

# ---------------------------------------------------------------- 2. performance (3 panels)
fig, axs = plt.subplots(1, 3, figsize=(15, 4.4))

# (a) cross-modal forgetting: core vs naive
mods = ["Language\n(7B)", "Vision\n(ViT)", "Security\n(IDS)"]
naive = [0.012, 0.133, 0.152]; core = [0.000, 0.022, 0.069]
x = np.arange(3); w = 0.36
axs[0].bar(x - w/2, naive, w, label="naive", color="#c44e52")
axs[0].bar(x + w/2, core, w, label="Antaḥkaraṇa core", color="#2BB673")
axs[0].set_xticks(x); axs[0].set_xticklabels(mods); axs[0].set_ylabel("catastrophic forgetting  (lower = better)")
axs[0].set_title("Less forgetting on every modality", fontweight="bold")
axs[0].legend(frameon=False)
for i in range(3):
    axs[0].text(i - w/2, naive[i] + .004, f"{naive[i]:.3f}", ha="center", fontsize=8)
    axs[0].text(i + w/2, core[i] + .004, f"{core[i]:.3f}", ha="center", fontsize=8)

# (b) zero-day AUROC per held-out attack family (no cherry-pick)
fam = ["dos", "probe", "r2l", "u2r"]; auroc = [0.79, 0.90, 0.78, 0.96]
axs[1].bar(fam, auroc, color="#4C9BE8"); axs[1].axhline(0.5, ls="--", color="#888", lw=1)
axs[1].set_ylim(0, 1.05); axs[1].set_ylabel("zero-day AUROC (avidyā, label-free)")
axs[1].set_title("Zero-day detection — every held-out family", fontweight="bold")
axs[1].text(3.0, 0.54, "chance", color="#888", fontsize=8)
for i, v in enumerate(auroc):
    axs[1].text(i, v + .02, f"{v:.2f}", ha="center", fontsize=9)

# (c) adaptive evasion curve
ev = [0.0, 0.25, 0.5, 0.75, 1.0]; ea = [0.945, 0.937, 0.912, 0.806, 0.246]
axs[2].plot(ev, ea, "-o", color="#E8804C", lw=2.2, ms=7)
axs[2].axhline(0.5, ls="--", color="#888", lw=1)
axs[2].set_ylim(0, 1.05); axs[2].set_xlabel("attack disguised toward normal →")
axs[2].set_ylabel("AUROC"); axs[2].set_title("Honest limit: evadable under full disguise", fontweight="bold")
axs[2].text(0.55, 0.54, "chance", color="#888", fontsize=8)
plt.tight_layout(); plt.savefig(os.path.join(A, "performance.png"), dpi=150, bbox_inches="tight"); plt.close()
print("wrote", os.listdir(A))