#!/usr/bin/env python3
"""
Scaffold Jump Simulation
------------------------
Simulates the progressive, evidence-gated fading of human involvement
according to the Scaffold Jump formal model.

Produces three figures:
1. Main trajectory (H(t) + P(t) with phase shading and jump markers)
2. Cumulative human effort
3. Phase timeline
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.lines import Line2D

# -------------------------------------------------
# 1. Simulation parameters (easy to experiment with)
# -------------------------------------------------
np.random.seed(42)

n_windows = 120
H0 = 1.0
P_star = 0.92
lambda_fade = 3.5
alpha = 0.78          # fade factor
beta = 1.7            # recovery factor
window_size_for_jump = 4
safety_breach_prob = 0.025

# Phase boundaries (human involvement fraction)
phase_bounds = {
    "Full Scaffold":       (0.75, 1.01),
    "Selective Scaffold":  (0.40, 0.75),
    "Monitor Scaffold":    (0.12, 0.40),
    "Scaffold Removed":    (0.00, 0.12)
}

# -------------------------------------------------
# 2. Helper
# -------------------------------------------------
def get_phase(h):
    for name, (lo, hi) in phase_bounds.items():
        if lo <= h < hi:
            return name
    return "Scaffold Removed"

# -------------------------------------------------
# 3. Core simulation
# -------------------------------------------------
def simulate_scaffold_jump():
    H = np.zeros(n_windows)
    P = np.zeros(n_windows)
    safety = np.ones(n_windows)
    phase = []
    jumps = []          # list of (window_index, "forward" | "reverse")
    effort = np.zeros(n_windows)

    H[0] = H0
    P[0] = 0.55
    phase.append(get_phase(H[0]))
    effort[0] = H[0]

    consecutive_good = 0

    for t in range(1, n_windows):
        # Performance improves with diminishing returns + noise
        improvement = 0.012 * (1 - P[t-1]) + np.random.normal(0, 0.018)
        P[t] = np.clip(P[t-1] + improvement, 0.40, 0.98)

        # Occasional safety breach
        if np.random.rand() < safety_breach_prob:
            safety[t] = 0
            P[t] *= 0.88
        else:
            safety[t] = 1

        # Continuous model suggestion
        if safety[t] == 1:
            H_cont = H0 * np.exp(-lambda_fade * (P[t] / P_star))
        else:
            H_cont = H[t-1] * beta

        H_cont = np.clip(H_cont, 0.0, 1.0)

        # Discrete jump logic
        if safety[t] == 0:
            H[t] = min(1.0, H[t-1] * beta)
            consecutive_good = 0
            if H[t] > H[t-1] + 0.05:
                jumps.append((t, "reverse"))
        else:
            if P[t] >= 0.78 and H[t-1] > 0.05:
                consecutive_good += 1
            else:
                consecutive_good = 0

            if consecutive_good >= window_size_for_jump and H[t-1] > 0.08:
                H[t] = max(0.0, H[t-1] * alpha)
                consecutive_good = 0
                if H[t] < H[t-1] - 0.04:
                    jumps.append((t, "forward"))
            else:
                # gentle drift toward continuous model
                H[t] = 0.85 * H[t-1] + 0.15 * H_cont

        H[t] = np.clip(H[t], 0.0, 1.0)
        phase.append(get_phase(H[t]))
        effort[t] = effort[t-1] + H[t]

    df = pd.DataFrame({
        "window": np.arange(n_windows),
        "H": H,
        "P": P,
        "safety": safety,
        "phase": phase,
        "cumulative_effort": effort
    })
    return df, jumps

# -------------------------------------------------
# 4. Run simulation
# -------------------------------------------------
df, jumps = simulate_scaffold_jump()

# -------------------------------------------------
# 5. Visualisation
# -------------------------------------------------
plt.style.use("seaborn-v0_8-whitegrid")
fig = plt.figure(figsize=(14, 11))

# Colour map for phases
colors = {
    "Full Scaffold":      "#a8dadc",
    "Selective Scaffold": "#457b9d",
    "Monitor Scaffold":   "#1d3557",
    "Scaffold Removed":   "#0d1b2a"
}

# ---- Plot 1: Main trajectory ----
ax1 = fig.add_subplot(3, 1, 1)

for phase_name, color in colors.items():
    mask = df["phase"] == phase_name
    if mask.any():
        ax1.fill_between(df["window"], 0, 1.08, where=mask,
                         color=color, alpha=0.25, linewidth=0)

ax1.plot(df["window"], df["H"], color="#1d3557", lw=2.4, label="Human involvement $H(t)$")
ax1.plot(df["window"], df["P"], color="#e63946", lw=1.9, ls="--", label="Performance $P(t)$")

# Jump markers
for t, jtype in jumps:
    if jtype == "forward":
        ax1.axvline(t, color="#2a9d8f", ls=":", alpha=0.7)
        ax1.plot(t, df.loc[t, "H"], "v", color="#2a9d8f", ms=9, zorder=5)
    else:
        ax1.axvline(t, color="#e76f51", ls=":", alpha=0.7)
        ax1.plot(t, df.loc[t, "H"], "^", color="#e76f51", ms=9, zorder=5)

ax1.set_ylim(0, 1.08)
ax1.set_ylabel("Involvement / Performance")
ax1.set_title("Scaffold Jump Simulation – Human Involvement & Performance", fontsize=13, pad=10)

# Legend
phase_patches = [mpatches.Patch(color=c, alpha=0.35, label=n) for n, c in colors.items()]
jump_handles = [
    Line2D([0], [0], marker="v", color="w", markerfacecolor="#2a9d8f", markersize=9, label="Forward jump"),
    Line2D([0], [0], marker="^", color="w", markerfacecolor="#e76f51", markersize=9, label="Reverse jump")
]
ax1.legend(handles=phase_patches + jump_handles + 
           [Line2D([0], [0], color="#1d3557", lw=2.4, label="Human involvement $H(t)$"),
            Line2D([0], [0], color="#e63946", lw=1.9, ls="--", label="Performance $P(t)$")],
           loc="upper right", fontsize=8, framealpha=0.95)

# ---- Plot 2: Cumulative effort ----
ax2 = fig.add_subplot(3, 1, 2)
ax2.plot(df["window"], df["cumulative_effort"], color="#1d3557", lw=2.3)
ax2.fill_between(df["window"], 0, df["cumulative_effort"], color="#1d3557", alpha=0.15)
ax2.set_ylabel("Cumulative human effort\n(sum of $H$)")
ax2.set_title("Cumulative Human Effort under Scaffold Jump", fontsize=12)

# ---- Plot 3: Phase timeline ----
ax3 = fig.add_subplot(3, 1, 3)
phase_ids = {name: i for i, name in enumerate(colors.keys())}

for i, (name, color) in enumerate(colors.items()):
    mask = df["phase"] == name
    ax3.fill_between(df["window"], i - 0.4, i + 0.4, where=mask,
                     color=color, alpha=0.85)

ax3.set_yticks(range(len(colors)))
ax3.set_yticklabels(list(colors.keys()))
ax3.set_xlabel("Evaluation window")
ax3.set_title("Phase Timeline", fontsize=12)
ax3.set_ylim(-0.6, 3.6)

plt.tight_layout()
plt.savefig("scaffold_jump_simulation.png", dpi=300, bbox_inches="tight")
plt.savefig("scaffold_jump_simulation.pdf", bbox_inches="tight")
print("Figures saved as: scaffold_jump_simulation.png and .pdf")
plt.show()

# -------------------------------------------------
# 6. Summary statistics
# -------------------------------------------------
print("\n=== Simulation Summary ===")
print(f"Total windows:              {n_windows}")
print(f"Final H(t):                 {df['H'].iloc[-1]:.3f}")
print(f"Final performance P(t):     {df['P'].iloc[-1]:.3f}")
print(f"Forward jumps:              {sum(1 for _, t in jumps if t == 'forward')}")
print(f"Reverse jumps:              {sum(1 for _, t in jumps if t == 'reverse')}")
print(f"Final cumulative effort:    {df['cumulative_effort'].iloc[-1]:.1f}")
print("\nTime spent in each phase:")
print(df["phase"].value_counts().reindex(phase_bounds.keys()))