"""Digital-recurrence monitor and a small logistic-map validation experiment."""
import json
import math
import numpy as np


class RecurrenceMonitor:
    def __init__(self, quantization=0.0, max_states=1000000):
        self.quantization = float(quantization)
        self.max_states = int(max_states)

    def _key(self, z):
        a = np.asarray(z)
        if self.quantization > 0:
            # The infinity-norm tolerance is represented by a deterministic grid key.
            a = np.round(a / self.quantization).astype(np.int64)
        return a.tobytes()

    def first_recurrence(self, trajectory):
        """Return (tau, period), where tau is the first repeated state's time."""
        seen = {}
        for t, z in enumerate(trajectory):
            if len(seen) >= self.max_states:
                break
            key = self._key(z)
            if key in seen:
                s = seen[key]
                return {"tau": int(t), "start": int(s), "period": int(t - s)}
            seen[key] = t
        return {"tau": None, "start": None, "period": None}


def logistic_trajectory(x0, n, dtype):
    x = np.array(x0, dtype=dtype)
    out = np.empty(n + 1, dtype=dtype)
    out[0] = x
    one = dtype(1.0)
    four = dtype(4.0)
    for i in range(n):
        x = dtype(four * x * (one - x))
        out[i + 1] = x
    return out


def finite_difference_slope(x0, horizon, dtype, delta=1e-7):
    """Slope of log forecast separation, matching the stated lambda estimator."""
    # Keep the perturbation representable in the selected precision.
    a = np.array(x0, dtype=dtype)
    b = np.array(x0 + delta, dtype=dtype)
    xa = logistic_trajectory(a, horizon, dtype)
    xb = logistic_trajectory(b, horizon, dtype)
    d = np.abs(xb.astype(np.float64) - xa.astype(np.float64))
    valid = d[1:] > 0
    if valid.sum() < 3:
        return float("nan")
    y = np.log(np.maximum(d[1:][valid], np.finfo(np.float64).tiny))
    t = np.arange(1, horizon + 1, dtype=np.float64)[valid]
    return float(np.polyfit(t, y, 1)[0])


def restart_lambda(x0, horizon, dtype, repeats=32):
    vals = []
    # Independent starts are deterministic, but separated by a fixed irrational-ish offset.
    for r in range(repeats):
        start = (float(x0) + (r + 1) * 0.00123456789) % 0.999
        v = finite_difference_slope(start, horizon, dtype)
        if np.isfinite(v):
            vals.append(v)
    return float(np.mean(vals)) if vals else float("nan")


def run():
    seed = 2386
    np.random.seed(seed)
    x0 = 0.1234567
    eps = 1e-3
    results = []
    # Long trajectories locate the recurrence scale; horizons test the predicted kink.
    for name, dtype, max_n in [("float16", np.float16, 20000),
                                ("float32", np.float32, 12000),
                                ("float64", np.float64, 20000)]:
        tr = logistic_trajectory(x0, max_n, dtype)
        rec = RecurrenceMonitor(max_states=max_n + 2).first_recurrence(tr)
        row = {"dtype": name, "recurrence": rec, "horizons": []}
        for n in [32, 128, 512, 2048, 4096, 8192, 12000]:
            if n > max_n:
                continue
            one = finite_difference_slope(x0, n, dtype)
            # Restart segments are short relative to the long-rollout test.
            k = min(128, n)
            rest = restart_lambda(x0, k, dtype)
            collapse = abs(one - rest) / (abs(rest) + eps) if np.isfinite(one) and np.isfinite(rest) else float("nan")
            row["horizons"].append({"N": n, "lambda_one": one,
                                    "lambda_restart": rest, "collapse": collapse,
                                    "alarm": bool(rec["tau"] is not None and n >= rec["tau"] and collapse > 1.0)})
        results.append(row)
    return {"seed": seed, "results": results}


def json_safe(value):
    if isinstance(value, dict):
        return {k: json_safe(v) for k, v in value.items()}
    if isinstance(value, list):
        return [json_safe(v) for v in value]
    if isinstance(value, float) and not math.isfinite(value):
        return None
    return value


if __name__ == "__main__":
    print(json.dumps(json_safe(run()), indent=2, allow_nan=False))