import json
import math
import random
from pathlib import Path

import numpy as np


def phi(x):
    x = np.asarray(x)
    return np.stack([
        0.55 * np.sin(2 * np.pi * x) + 0.25 * np.cos(4 * np.pi * x),
        0.55 * np.cos(2 * np.pi * x) - 0.25 * np.sin(6 * np.pi * x),
    ], axis=-1)


def generate(r, ell=2, n=7000, burn=45, seed=0):
    # Sample the stationary attractor without long forward floating-point
    # doubling orbits (which eventually lose all mantissa bits).  For a
    # current uniform x, choose inverse branches x'=(x+k)/ell and accumulate
    # h = sum_{j>=0} r^j phi(x_{-1-j}).
    rng = np.random.default_rng(seed)
    x = rng.random(n)
    h = np.zeros((n, 2))
    weight = 1.0
    for _ in range(burn):
        branch = rng.integers(0, ell, size=n)
        x = (x + branch) / ell
        h += weight * phi(x)
        weight *= r
    return np.column_stack([x, h])


def box_dimension(z, scales=(6, 8, 12, 16, 24, 32)):
    # Normalize coordinates so grid resolution has comparable meaning.
    lo, hi = z.min(0), z.max(0)
    q = (z - lo) / (hi - lo + 1e-12)
    counts = []
    for m in scales:
        cells = np.floor(np.minimum(q, 1 - 1e-12) * m).astype(np.int64)
        counts.append(np.unique(cells, axis=0).shape[0])
    x = np.log(np.asarray(scales, float))
    y = np.log(np.asarray(counts, float))
    # Fine-scale endpoint bins are most informative, while retaining robustness.
    slope = float(np.polyfit(x[-5:], y[-5:], 1)[0])
    return slope, counts


def correlation_dimension(z, seed=123, pairs=30000):
    rng = np.random.default_rng(seed)
    n = len(z)
    ii = rng.integers(0, n, pairs)
    jj = rng.integers(0, n, pairs)
    keep = ii != jj
    d = np.linalg.norm(z[ii[keep]] - z[jj[keep]], axis=1)
    d = d[d > 1e-15]
    if len(d) < 100:
        return float("nan")
    # Pairwise distance quantiles avoid the zero/finite-sample extremes.
    radii = np.geomspace(np.quantile(d, .01), np.quantile(d, .35), 14)
    c = np.array([(d < q).mean() for q in radii])
    good = (c > 3 / len(d)) & (c < .4)
    slope = float(np.polyfit(np.log(radii[good]), np.log(c[good]), 1)[0])
    return slope


def contraction_check(r=.8, ell=2, n=30):
    # Same driver, two fibers: exact difference is r^n times its initial value.
    x = .371234
    a, b = np.array([1.0, -.4]), np.array([-.3, .8])
    vals = []
    for t in range(n + 1):
        vals.append(np.linalg.norm(a - b))
        a = r * a + phi(x)
        b = r * b + phi(x)
        x = (ell * x) % 1
    slope = float(np.polyfit(np.arange(1, n + 1), np.log(vals[1:]), 1)[0])
    return slope, math.log(r), float(vals[-1] / vals[0]), r ** n


def main():
    np.random.seed(7)
    random.seed(7)
    ell = 2
    rs = [0.55, 0.65, 0.70, 1 / math.sqrt(2), 0.75, 0.85, 0.95]
    rows = []
    for k, r in enumerate(rs):
        z = generate(r, ell=ell, seed=100 + k)
        predicted = 1.0 + min(2.0, math.log(ell) / abs(math.log(r)))
        bd, boxes = box_dimension(z)
        cd = correlation_dimension(z, seed=200 + k)
        g = math.log(ell) + 2 * math.log(r)
        rows.append({
            "r": r, "ell_det": ell * r * r, "g": g,
            "predicted_dimension": predicted,
            "box_dimension": bd, "correlation_dimension": cd,
            "boxes": boxes,
        })
    cslope, clog, cfinal, cpred = contraction_check()
    result = {
        "seed": 7, "ell": ell, "dimension": 2,
        "threshold_r": 1 / math.sqrt(ell),
        "contraction": {"observed_log_slope": cslope, "predicted_log_slope": clog,
                        "observed_ratio_n30": cfinal, "predicted_ratio_n30": cpred},
        "sweep": rows,
        "notes": "Volume exponent is exact for A=rI: log(ell)+log|det(A)|=log(ell*r^2).",
    }
    Path("toy_results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()