#!/usr/bin/env python3
"""
Scaffold Jump – Parameter Sensitivity Heatmap
Shows how final phase / time in Scaffold Removed varies with λ and safety breach probability
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

np.random.seed(42)

def run_single_simulation(lambda_fade, safety_breach_prob, n_windows=100):
    H = 1.0
    P = 0.55
    alpha = 0.78
    beta = 1.7
    consecutive_good = 0
    time_in_removed = 0

    for t in range(1, n_windows):
        improvement = 0.012 * (1 - P) + np.random.normal(0, 0.018)
        P = np.clip(P + improvement, 0.4, 0.98)

        if np.random.rand() < safety_breach_prob:
            safety = 0
            P *= 0.88
        else:
            safety = 1

        if safety == 0:
            H = min(1.0, H * beta)
            consecutive_good = 0
        else:
            if P >= 0.78 and H > 0.05:
                consecutive_good += 1
            else:
                consecutive_good = 0

            if consecutive_good >= 4 and H > 0.08:
                H = max(0.0, H * alpha)
                consecutive_good = 0
            else:
                H_cont = np.exp(-lambda_fade * (P / 0.92))
                H = 0.85 * H + 0.15 * np.clip(H_cont, 0, 1)

        H = np.clip(H, 0, 1)
        if H < 0.12:
            time_in_removed += 1

    return time_in_removed / n_windows   # proportion of time in Scaffold Removed

# Create grid
lambda_values = np.linspace(1.5, 6.0, 12)
breach_values = np.linspace(0.005, 0.08, 12)

results = np.zeros((len(breach_values), len(lambda_values)))

print("Running parameter sweep...")
for i, breach in enumerate(breach_values):
    for j, lam in enumerate(lambda_values):
        # Average over a few runs for stability
        scores = [run_single_simulation(lam, breach) for _ in range(6)]
        results[i, j] = np.mean(scores)

# Plot
plt.figure(figsize=(10, 7))
sns.heatmap(results,
            xticklabels=np.round(lambda_values, 1),
            yticklabels=np.round(breach_values, 3),
            cmap="YlGnBu",
            annot=False,
            cbar_kws={'label': 'Proportion of time in Scaffold Removed'})

plt.xlabel("Fade speed (λ)", fontsize=12)
plt.ylabel("Safety breach probability", fontsize=12)
plt.title("Scaffold Jump Sensitivity Analysis\nProportion of time spent in Scaffold Removed", fontsize=13, pad=12)
plt.tight_layout()

plt.savefig("scaffold_jump_heatmap.png", dpi=300, bbox_inches="tight")
plt.savefig("scaffold_jump_heatmap.pdf", bbox_inches="tight")
print("Heatmap saved as scaffold_jump_heatmap.png / .pdf")
plt.show()