"""Toy numerical verification for phase-blind checkpoint scheduling.

This deliberately tests the mathematical mechanism, not distributed PyTorch
I/O.  Run with: python3 phase_blind.py
"""
import json
from pathlib import Path
import numpy as np


def anonymous_pair_map(delta, T=1.0, d=0.1):
    """Return map for identical periodic jobs on an anonymous resource."""
    if not (0 < d < T):
        raise ValueError("require 0 < d < T")
    return np.asarray(delta, dtype=float).copy()


def priority_pair_map(delta, alpha=0.25):
    """Toy identity/age-priority comparator with deterministic contraction."""
    return np.asarray(delta, dtype=float) * (1.0 - alpha)


def fleet_map(gaps, f):
    """Linearized ordered-gap map from the anonymous-throughput formula."""
    gaps = np.asarray(gaps, dtype=float)
    n = gaps.size + 1
    out = np.empty_like(gaps)
    for j in range(1, n):
        out[j - 1] = ((n - j) * f(j)) / (j * f(n - j)) * gaps[j - 1]
    return out


def verify_pair_map():
    rows = []
    for T, d in [(0.5, 0.01), (1.0, 0.1), (2.0, 1.9)]:
        deltas = np.linspace(1e-4, T / 2 - 1e-4, 100)
        nxt = anonymous_pair_map(deltas, T, d)
        slope, intercept = np.polyfit(deltas, nxt, 1)
        rows.append({"T": T, "d": d, "slope": float(slope),
                     "intercept": float(intercept),
                     "max_identity_error": float(np.max(abs(nxt - deltas)))})
    deltas = np.linspace(1e-4, .49, 100)
    priority_slope = np.polyfit(deltas, priority_pair_map(deltas), 1)[0]
    return {"sweep": rows, "priority_slope": float(priority_slope)}


def verify_fleet():
    rows = []
    functions = [("constant", lambda k: 1.0),
                 ("inverse_n", lambda k: 1.0 / k),
                 ("sqrt_n", lambda k: 1.0 / np.sqrt(k))]
    for n in range(2, 11):
        for name, fn in functions:
            base = np.linspace(.01, .01 * n, n - 1)
            lambdas = fleet_map(base, fn) / base
            rows.append({"N": n, "f": name, "det": float(np.prod(lambdas)),
                         "reciprocal_error": float(max(
                             abs(lambdas[j] * lambdas[-j - 1] - 1)
                             for j in range(n - 1))),
                         "max_lambda": float(max(lambdas))})
    return rows


def hitting_time(m, sigma, trials=3000, max_steps=200000, seed=0):
    """First passage of a gap random walk; each endpoint has noise sigma."""
    rng = np.random.default_rng(seed)
    x = np.full(trials, float(m))
    alive = np.ones(trials, dtype=bool)
    hit = np.full(trials, max_steps, dtype=np.int64)
    for step in range(1, max_steps + 1):
        idx = np.flatnonzero(alive)
        if not len(idx):
            break
        x[idx] += rng.normal(0.0, np.sqrt(2.0) * sigma, size=len(idx))
        crossed = x[idx] <= 0
        if np.any(crossed):
            hit[idx[crossed]] = step
            alive[idx[crossed]] = False
    return hit


def verify_noise_scaling():
    m = 0.20
    sigmas = np.array([.004, .006, .009, .013, .019])
    medians = np.array([np.median(hitting_time(m, s, seed=100 + i))
                        for i, s in enumerate(sigmas)])
    slope = np.polyfit(np.log(sigmas), np.log(medians), 1)[0]
    ratios = medians * sigmas ** 2 / m ** 2
    # A direct fourfold prediction test, interpolated from the nearest pair.
    ratio_006_012 = medians[1] / np.median(hitting_time(m, .012, seed=222))
    return {"m": m, "sigmas": sigmas.tolist(), "median_steps": medians.tolist(),
            "loglog_exponent": float(slope), "C_median": ratios.tolist(),
            "C_median_mean": float(np.mean(ratios)),
            "median_ratio_sigma_.006_to_.012": float(ratio_006_012)}


def summarize(result):
    pair = result["pair"]["sweep"]
    fleet = result["fleet"]
    noise = result["noise"]
    max_det_err = max(abs(x["det"] - 1) for x in fleet)
    max_recip_err = max(x["reciprocal_error"] for x in fleet)
    max_pair_err = max(x["max_identity_error"] for x in pair)
    # Predictions are mechanism checks, not claims about production storage.
    result["prediction_checks"] = {
        "P1_pair_return_slope_predicted_1_observed": [x["slope"] for x in pair],
        "P1_max_absolute_identity_error": max_pair_err,
        "P2_determinant_predicted_1_max_abs_error": max_det_err,
        "P2_reciprocal_eigenvalue_max_error": max_recip_err,
        "P3_noise_exponent_predicted_-2_observed": noise["loglog_exponent"],
        "P3_sigma_doubling_lifetime_ratio_predicted_4_observed":
            noise["median_ratio_sigma_.006_to_.012"],
        "all_mechanism_checks_pass": bool(
            max_pair_err < 1e-12 and max_det_err < 1e-12 and
            max_recip_err < 1e-12 and abs(noise["loglog_exponent"] + 2) < .15 and
            2.5 < noise["median_ratio_sigma_.006_to_.012"] < 6.0)
    }


def main():
    result = {"pair": verify_pair_map(), "fleet": verify_fleet(),
              "noise": verify_noise_scaling()}
    d = .30
    anon, priority = d, d
    for _ in range(20):
        anon = float(anonymous_pair_map(anon))
        priority = float(priority_pair_map(priority))
    result["comparison_20_cycles"] = {"initial_gap": d,
        "anonymous_gap": anon, "priority_gap": priority}
    summarize(result)
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()