"""
fractal.py -- Fractal (Cantor-spectrum) RoPE frequency schedule.

Ported from the wheelerv2 / Quazimoto-LM codebase. That file also carried a
Mandelbrot phase-seeding table for Quazimoto's Kuramoto oscillator bank; Byrne
has no oscillator bank, so only the RoPE half is carried over here.

Standard RoPE places its inverse frequencies on a geometric (log-uniform) grid.
`fractal_rope_inv_freq` keeps the same endpoints and band count -- so short- and
long-range positional coverage is preserved -- but repositions the exponents onto
the middle-thirds Cantor set, blended with the geometric schedule by `gamma`.
gamma=0 is exact standard RoPE; gamma=1 is a fully fractal spectrum.
"""

import torch


@torch.no_grad()
def _cantor_map(u, bits=24):
    """Monotone map [0,1] -> the middle-thirds Cantor set.

    Read the binary digits b_k of u and use them as base-3 digits 2*b_k:
        u = 0.b1 b2 b3...(2)   ->   sum_k (2 b_k) / 3^k
    This is the inverse of the Cantor 'devil's staircase'. It is monotone
    non-decreasing (so it preserves frequency ordering) and its image is the
    self-similar Cantor set, giving a genuinely fractal placement of points.

    NOTE (fixed vs. the wheelerv2 original): the work is done in float64 and the
    upper clamp is 1 - 1e-12. In float32 the original's `1.0 - 1e-9` rounds back
    to exactly 1.0, so the clamp never bit; u=1 then made the first digit step
    compute floor(2.0) = 2 -- a "bit" of 2 -- giving cantor(1) = 4/3 instead of 1.
    That overshot the top of the exponent range and broke the endpoint-preserving
    property this schedule relies on."""
    u = u.double().clamp(0.0, 1.0 - 1e-12)
    out = torch.zeros_like(u)
    frac = u.clone()
    p3 = 1.0
    for _ in range(bits):
        frac = frac * 2.0
        bit = torch.floor(frac)
        frac = frac - bit
        p3 = p3 / 3.0
        out = out + 2.0 * bit * p3
    return out


@torch.no_grad()
def fractal_rope_inv_freq(dim, theta=10000.0, gamma=1.0, bits=24):
    """RoPE inverse-frequency schedule placed on a fractal (Cantor) spectrum.

        e_f = (1-g)*e_geom_f  +  g*e_max*cantor(e_geom_f / e_max)
        w_f = theta^(-e_f)

    The map is monotone, so w stays sorted (matching RotaryEmbedding's cache
    layout). Returns a [dim//2] tensor."""
    e_geom = torch.arange(0, dim, 2).double() / dim         # [F], 0 .. (d-2)/d
    e_max = float(e_geom[-1]) if e_geom.numel() > 1 else 1.0
    if e_max <= 0:
        return (theta ** (-e_geom)).float()
    e_frac = e_max * _cantor_map(e_geom / e_max, bits=bits)
    e = (1.0 - gamma) * e_geom + gamma * e_frac
    return (theta ** (-e)).float()