"""Generate all documentation charts from real experiment data.

Outputs to docs/images/ for use in README and GitHub pages.
"""

import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np

OUT = Path("docs/images")
OUT.mkdir(parents=True, exist_ok=True)

# Shared style
plt.rcParams.update({
    "figure.facecolor": "#0d1117",
    "axes.facecolor": "#0d1117",
    "axes.edgecolor": "#30363d",
    "axes.labelcolor": "#c9d1d9",
    "text.color": "#c9d1d9",
    "xtick.color": "#8b949e",
    "ytick.color": "#8b949e",
    "grid.color": "#21262d",
    "grid.alpha": 0.6,
    "font.family": "monospace",
    "font.size": 11,
})


def load_meta(path):
    with open(path) as f:
        return json.load(f)


# ─────────────────────────────────────────────────────────────
# 1. OPS SCATTER PLOT (hero image)
# ─────────────────────────────────────────────────────────────
print("1/6  Ops scatter plot...")

# Use the no-mutation run — shows dramatic crystallization transition
data = np.load("data/bff_runs/final_10000000.npz")
ops_log = data["ops_log"]

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

# Bin into windows and show density as a 2D histogram
# This avoids the scatter plot density problem entirely
n_xbins = 500
n_ybins = 200
window_size = len(ops_log) // n_xbins

# Build mean ops per window + fraction of high-ops interactions
x_centers = []
frac_active = []      # fraction with ops > 100
frac_maxed = []       # fraction hitting 10k cap
mean_ops = []

for i in range(n_xbins):
    start = i * window_size
    end = start + window_size
    window = ops_log[start:end].astype(np.float64)
    x_centers.append(start + window_size // 2)
    mean_ops.append(window.mean())
    frac_active.append((window > 100).sum() / len(window))
    frac_maxed.append((window >= 10_000).sum() / len(window))

x_centers = np.array(x_centers)
mean_ops = np.array(mean_ops)
frac_active = np.array(frac_active)
frac_maxed = np.array(frac_maxed)

# Two-panel: mean ops + fraction of active/maxed interactions
ax.bar(x_centers, mean_ops, width=window_size * 0.9,
       color="#58a6ff", alpha=0.7, linewidth=0)
ax.set_xlabel("Interaction")
ax.set_ylabel("Mean Operations per Window")
ax.set_title("Phase Transition: Self-Replicators Emerge from Random Noise",
             fontsize=13, fontweight="bold", pad=12)
ax.set_xlim(0, len(ops_log))

# Overlay: fraction of max-steps interactions as orange line on twin axis
ax2 = ax.twinx()
ax2.plot(x_centers, frac_maxed * 100, color="#d29922", linewidth=1.0,
         alpha=0.8, label="% hitting max steps")
ax2.set_ylabel("% Hitting Max Steps", color="#d29922")
ax2.tick_params(axis="y", labelcolor="#d29922")
ax2.set_ylim(0, max(frac_maxed * 100) * 1.3)
ax2.legend(loc="upper right", fontsize=9,
           facecolor="#161b22", edgecolor="#30363d", labelcolor="#d29922")

# Format x-axis as millions
ax.xaxis.set_major_formatter(
    matplotlib.ticker.FuncFormatter(lambda v, _: f"{v/1e6:.0f}M")
)

ax.grid(True, axis="y", linewidth=0.3)
fig.tight_layout()
fig.savefig(OUT / "ops_scatter.png", dpi=200, bbox_inches="tight")
plt.close(fig)
print("   -> ops_scatter.png")


# ─────────────────────────────────────────────────────────────
# 2. COMPRESSION CHART
# ─────────────────────────────────────────────────────────────
print("2/6  Compression chart...")

meta = load_meta("data/bff_runs/final_10000000_meta.json")
interactions = meta["snapshot_interactions"]
compression = meta["snapshot_compression"]

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

ax.plot(interactions, compression, linewidth=1.0, color="#58a6ff")
ax.axhline(y=1.0, color="#8b949e", linestyle="--", alpha=0.5, linewidth=0.8)
ax.text(interactions[-1] * 0.82, 1.02, "random baseline",
        color="#8b949e", fontsize=9, alpha=0.7)

ax.fill_between(interactions, compression, alpha=0.08, color="#58a6ff")
ax.set_xlabel("Interaction")
ax.set_ylabel("Compression Ratio")
ax.set_title("Structure Emerging from Randomness",
             fontsize=13, fontweight="bold", pad=12)
ax.set_ylim(bottom=0, top=1.1)

ax.xaxis.set_major_formatter(
    matplotlib.ticker.FuncFormatter(lambda v, _: f"{v/1e6:.0f}M")
)

ax.grid(True, linewidth=0.3)
fig.tight_layout()
fig.savefig(OUT / "compression.png", dpi=200, bbox_inches="tight")
plt.close(fig)
print("   -> compression.png")


# ─────────────────────────────────────────────────────────────
# 3. PHASE DIAGRAM (4 mutation rates)
# ─────────────────────────────────────────────────────────────
print("3/6  Phase diagram...")

runs = [
    ("No mutation (crystal)",     "data/bff_runs/final_10000000_meta.json",          "#f85149"),
    ("Mutation 1e-5 (soft crystal)", "data/bff_mutation_1e-5/final_10000000_meta.json", "#d29922"),
    ("Mutation 1e-3 (ecology)",   "data/bff_mutation_1e-3/final_10000000_meta.json",  "#3fb950"),
    ("Mutation 1e-1 (dissolution)", "data/bff_mutation_1e-1/final_10000000_meta.json", "#8b949e"),
]

fig, ax = plt.subplots(figsize=(12, 5.5))

for label, path, color in runs:
    m = load_meta(path)
    ax.plot(m["snapshot_interactions"], m["snapshot_compression"],
            linewidth=1.2, color=color, label=label, alpha=0.9)

ax.axhline(y=1.0, color="#8b949e", linestyle="--", alpha=0.3, linewidth=0.8)
ax.set_xlabel("Interaction")
ax.set_ylabel("Compression Ratio")
ax.set_title("Phase Diagram: Mutation Rate Controls Everything",
             fontsize=13, fontweight="bold", pad=12)
ax.set_ylim(bottom=0, top=1.1)
ax.legend(loc="upper right", fontsize=10,
          facecolor="#161b22", edgecolor="#30363d", labelcolor="#c9d1d9")

ax.xaxis.set_major_formatter(
    matplotlib.ticker.FuncFormatter(lambda v, _: f"{v/1e6:.0f}M")
)

ax.grid(True, linewidth=0.3)
fig.tight_layout()
fig.savefig(OUT / "phase_diagram.png", dpi=200, bbox_inches="tight")
plt.close(fig)
print("   -> phase_diagram.png")


# ─────────────────────────────────────────────────────────────
# 4. REPLICATOR TIMELINE
# ─────────────────────────────────────────────────────────────
print("4/6  Replicator timeline...")

# Use 1e-5 run — has highest viral/cellular diversity
meta_rep = load_meta("data/bff_mutation_1e-5/final_10000000_meta.json")
rep_interactions = meta_rep["snapshot_interactions"]
rep_counts = meta_rep["snapshot_replicator_counts"]

inanimate = np.array([s.get("inanimate", 0) for s in rep_counts])
viral = np.array([s.get("viral", 0) for s in rep_counts])
cellular = np.array([s.get("cellular", 0) for s in rep_counts])

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 7), height_ratios=[2, 1],
                                sharex=True)

# Top panel: total replicators (inanimate dominates)
ax1.fill_between(rep_interactions, inanimate, alpha=0.3, color="#8b949e")
ax1.plot(rep_interactions, inanimate, linewidth=0.7, color="#8b949e",
         label="Inanimate")
ax1.set_ylabel("Count")
ax1.set_title("Replicator Emergence (mutation=1e-5)",
              fontsize=13, fontweight="bold", pad=12)
ax1.legend(loc="upper left", fontsize=10,
           facecolor="#161b22", edgecolor="#30363d", labelcolor="#c9d1d9")
ax1.grid(True, linewidth=0.3)

# Bottom panel: viral + cellular (the interesting ones, zoomed in)
ax2.plot(rep_interactions, viral, linewidth=1.2, color="#d29922",
         label="Viral", alpha=0.9)
ax2.plot(rep_interactions, cellular, linewidth=1.2, color="#3fb950",
         label="Cellular", alpha=0.9)
ax2.fill_between(rep_interactions, viral, alpha=0.15, color="#d29922")
ax2.fill_between(rep_interactions, cellular, alpha=0.15, color="#3fb950")
ax2.set_xlabel("Interaction")
ax2.set_ylabel("Count")
ax2.legend(loc="upper left", fontsize=10,
           facecolor="#161b22", edgecolor="#30363d", labelcolor="#c9d1d9")
ax2.grid(True, linewidth=0.3)

ax2.xaxis.set_major_formatter(
    matplotlib.ticker.FuncFormatter(lambda v, _: f"{v/1e6:.0f}M")
)

fig.tight_layout()
fig.savefig(OUT / "replicator_timeline.png", dpi=200, bbox_inches="tight")
plt.close(fig)
print("   -> replicator_timeline.png")

# Also load the ecology run meta for depth chart below
meta_eco = load_meta("data/bff_mutation_1e-3/final_10000000_meta.json")


# ─────────────────────────────────────────────────────────────
# 5. DEPTH TIMELINE
# ─────────────────────────────────────────────────────────────
print("5/6  Depth timeline...")

# Compare depth spread (max - mean) across mutation rates to show
# how different regimes build complexity differently
depth_runs = [
    ("1e-5 (soft crystal)", "data/bff_mutation_1e-5/final_10000000_meta.json", "#d29922"),
    ("1e-3 (ecology)",      "data/bff_mutation_1e-3/final_10000000_meta.json", "#3fb950"),
    ("1e-1 (dissolution)",  "data/bff_mutation_1e-1/final_10000000_meta.json", "#8b949e"),
]

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 7), sharex=True)

for label, path, color in depth_runs:
    m = load_meta(path)
    interactions = m["snapshot_interactions"]
    max_d = m.get("snapshot_max_depth", [])
    mean_d = m.get("snapshot_mean_depth", [])
    if max_d and mean_d:
        ax1.plot(interactions, max_d, linewidth=1.0, color=color,
                 label=label, alpha=0.9)
        # Depth spread = max - mean (measures how uneven complexity is)
        spread = [mx - mn for mx, mn in zip(max_d, mean_d)]
        ax2.plot(interactions, spread, linewidth=1.0, color=color,
                 label=label, alpha=0.9)

ax1.set_ylabel("Max Depth")
ax1.set_title("Symbiogenesis Depth Across Regimes",
              fontsize=13, fontweight="bold", pad=12)
ax1.legend(loc="lower right", fontsize=10,
           facecolor="#161b22", edgecolor="#30363d", labelcolor="#c9d1d9")
ax1.grid(True, linewidth=0.3)

ax2.set_xlabel("Interaction")
ax2.set_ylabel("Depth Spread (max - mean)")
ax2.legend(loc="upper left", fontsize=10,
           facecolor="#161b22", edgecolor="#30363d", labelcolor="#c9d1d9")
ax2.grid(True, linewidth=0.3)

ax2.xaxis.set_major_formatter(
    matplotlib.ticker.FuncFormatter(lambda v, _: f"{v/1e6:.0f}M")
)

fig.tight_layout()
fig.savefig(OUT / "depth_timeline.png", dpi=200, bbox_inches="tight")
plt.close(fig)
print("   -> depth_timeline.png")


# ─────────────────────────────────────────────────────────────
# 6. CORE LOOP DIAGRAM
# ─────────────────────────────────────────────────────────────
print("6/6  Core loop diagram...")

fig, ax = plt.subplots(figsize=(12, 4))
ax.set_xlim(0, 12)
ax.set_ylim(0, 4)
ax.axis("off")

# Box style
box_kw = dict(boxstyle="round,pad=0.4", facecolor="#161b22",
              edgecolor="#58a6ff", linewidth=1.5)
arrow_kw = dict(arrowstyle="->,head_width=0.3,head_length=0.15",
                color="#58a6ff", linewidth=1.5)

# Boxes
boxes = [
    (1.0, 2.0, "Pick 2\nrandom tapes"),
    (3.5, 2.0, "Concatenate\n(128 bytes)"),
    (6.0, 2.0, "Execute\n7-instruction BF"),
    (8.5, 2.0, "Split back\n(2 x 64 bytes)"),
    (11.0, 2.0, "Write\nhome"),
]

for x, y, text in boxes:
    ax.text(x, y, text, ha="center", va="center", fontsize=10,
            color="#c9d1d9", bbox=box_kw, fontfamily="monospace")

# Arrows between boxes
for i in range(len(boxes) - 1):
    x1 = boxes[i][0] + 0.9
    x2 = boxes[i+1][0] - 0.9
    ax.annotate("", xy=(x2, 2.0), xytext=(x1, 2.0),
                arrowprops=arrow_kw)

# Mutation branch
ax.text(8.5, 0.6, "Mutate?\n(bit-flip)", ha="center", va="center",
        fontsize=9, color="#d29922",
        bbox=dict(boxstyle="round,pad=0.3", facecolor="#161b22",
                  edgecolor="#d29922", linewidth=1.2, linestyle="--"),
        fontfamily="monospace")
ax.annotate("", xy=(8.5, 1.1), xytext=(8.5, 1.55),
            arrowprops=dict(arrowstyle="->,head_width=0.2,head_length=0.12",
                            color="#d29922", linewidth=1.2, linestyle="--"))

# Loop-back arrow (from Write back to Pick)
ax.annotate("",
            xy=(1.0, 3.2), xytext=(11.0, 3.2),
            arrowprops=dict(arrowstyle="->,head_width=0.3,head_length=0.15",
                            color="#3fb950", linewidth=1.5,
                            connectionstyle="arc3,rad=0.0"))
ax.text(6.0, 3.5, "repeat", ha="center", va="center",
        fontsize=10, color="#3fb950", fontfamily="monospace",
        fontstyle="italic")

fig.tight_layout()
fig.savefig(OUT / "core_loop.png", dpi=200, bbox_inches="tight")
plt.close(fig)
print("   -> core_loop.png")


print(f"\nAll charts saved to {OUT}/")
print("Files:", [f.name for f in sorted(OUT.glob("*.png"))])
