"""
fractal.py -- Mandelbrot phase seeding for Quazimoto-LM's oscillator bank.

The recommended (and only) fractal integration: instead of a generic learned
phase initializer, give each TOKEN a characteristic dynamical signature drawn
from the Mandelbrot iteration z <- z^2 + c, and use the ANGLES of that orbit to
seed the N oscillator phases. The Mandelbrot map is itself an iterated dynamical
system, so this hands the Kuramoto ring bank a token-specific, deterministic,
parameter-free phase prior congruent with what the block already does.

We build a frozen [vocab_size, n_osc] table once: token id -> a complex point c
(spread over the Mandelbrot region by a 2D Halton low-discrepancy sequence so
coverage is even and deterministic) -> phase_k = angle(z_k) for the first n_osc
orbit points. The model adds this, through a zero-init gate, to to_theta(h) inside
each QuazimotoBlock -- a no-op at init that the optimizer can choose to open.

Smooth-by-construction: we read the orbit ANGLE (always defined, bounded to
(-pi, pi]) rather than escape-time, avoiding the chaotic boundary discontinuities
that raw escape counts would inject.
"""

import torch


def _halton(i, base):
    """Radical-inverse (van der Corput) value of i in the given base, in [0,1)."""
    f, r = 1.0, 0.0
    while i > 0:
        f /= base
        r += f * (i % base)
        i //= base
    return r


@torch.no_grad()
def mandelbrot_phase_table(vocab_size, n_osc, region=(-2.5, 1.0, -1.25, 1.25),
                           clamp_mag=1e3):
    """Return a frozen [vocab_size, n_osc] tensor of orbit-angle phase seeds.

    token id -> c via 2D Halton(base 2,3) over `region`; phase_k = angle(z_k) for
    k = 0..n_osc-1 of the iteration z <- z^2 + c (z0 = 0). This is the FLAT
    (tokenizer-agnostic) map; for the hierarchical byte-merge map build the table
    offline with build_fractal_table.py and load it via load_phase_table()."""
    x0, x1, y0, y1 = region
    ids = torch.arange(1, vocab_size + 1)
    hx = torch.tensor([_halton(int(i), 2) for i in ids])
    hy = torch.tensor([_halton(int(i), 3) for i in ids])
    cr = x0 + hx * (x1 - x0)
    ci = y0 + hy * (y1 - y0)
    return phases_from_c(torch.complex(cr, ci), n_osc, clamp_mag)


@torch.no_grad()
def phases_from_c(c, n_osc, clamp_mag=1e3):
    """Orbit-angle phases for a batch of complex seeds c [V] -> [V, n_osc].
    phase_k = angle(z_k) of z <- z^2 + c (z0=0); magnitude clamped (angle kept)."""
    z = torch.zeros_like(c)
    phases = torch.empty(c.shape[0], n_osc)
    for k in range(n_osc):
        z = z * z + c
        mag = z.abs()
        over = mag > clamp_mag                        # rescale escaped orbits, keep direction
        if over.any():
            z = torch.where(over, z / mag * clamp_mag, z)
        phases[:, k] = torch.angle(z)                 # in (-pi, pi]
    return phases


def load_phase_table(vocab_size, n_osc, path):
    """Load a precomputed phase table if it matches (vocab_size, n_osc); else None.
    Returns (phases, mode) or (None, None)."""
    import os
    if not os.path.exists(path):
        return None, None
    d = torch.load(path, map_location="cpu", weights_only=False)
    ph = d.get("phases")
    if ph is not None and tuple(ph.shape) == (vocab_size, n_osc):
        return ph.float(), d.get("mode", "precomputed")
    return None, None