#!/usr/bin/env python3
"""JAM 2 instrument: measure rotation/injection geometry by depth.

For node-count vector c_d modulo q=3^k, reverse-Collatz growth gives

    c_(d+1) = P c_d + i_d,

where P doubles every residue (a permutation modulo q) and i_d counts
the odd children injected by branch nodes.  After subtracting the uniform
vectors of the appropriate masses, the same identity holds for deviations.

This instrument measures the angle between P*dev(c_d) and dev(i_d), their
cancellation ratio, and the finite scaling of absolute and relative L2 error.
It does not promote finite decorrelation to an asymptotic theorem.
"""

from __future__ import annotations

import argparse
from dataclasses import dataclass

import numpy as np


DEFAULT_MODULI = (9, 27, 81, 243, 729)
UINT64_MAX = np.iinfo(np.uint64).max


@dataclass(frozen=True)
class Reading:
    depth: int
    cosine: float
    cancellation: float
    next_relative_l2: float
    current_relative_l2: float
    nodes: int


def counts(values: np.ndarray, modulus: int) -> np.ndarray:
    residues = np.asarray(values % modulus, dtype=np.int64)
    return np.bincount(residues, minlength=modulus).astype(np.int64)


def deviation(vector: np.ndarray) -> np.ndarray:
    return vector.astype(np.float64) - vector.sum() / vector.size


def cosine(left: np.ndarray, right: np.ndarray) -> float:
    denominator = np.linalg.norm(left) * np.linalg.norm(right)
    return float(np.dot(left, right) / denominator) if denominator else float("nan")


def next_parts(level: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    if int(level.max()) > UINT64_MAX // 2:
        raise OverflowError("uint64 horizon reached; lower --depth")
    branch = level[(level % 6 == 4) & (level != 4)]
    doubled = level * np.uint64(2)
    injected = ((branch - 1) // 3).astype(np.uint64)
    return doubled, injected


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--depth", type=int, default=63)
    parser.add_argument("--tail", type=int, default=20)
    parser.add_argument("--fit-from", type=int, default=30)
    parser.add_argument("--moduli", type=int, nargs="+", default=DEFAULT_MODULI)
    args = parser.parse_args()
    if not 1 <= args.depth <= 63:
        parser.error("--depth must be between 1 and 63 for exact uint64 arithmetic")
    if args.tail < 1:
        parser.error("--tail must be positive")
    if any(q < 3 or q % 3 for q in args.moduli):
        parser.error("every modulus must be a positive multiple of 3")

    history: dict[int, list[Reading]] = {q: [] for q in args.moduli}
    level = np.array([1], dtype=np.uint64)
    for depth in range(args.depth):
        doubled, injected = next_parts(level)
        nxt = np.concatenate((doubled, injected))
        for modulus in args.moduli:
            base_counts = counts(level, modulus)
            doubled_counts = counts(doubled, modulus)
            injected_counts = counts(injected, modulus)
            next_counts = counts(nxt, modulus)
            assert np.array_equal(doubled_counts + injected_counts, next_counts)

            rotated = deviation(doubled_counts)
            arrival = deviation(injected_counts)
            next_error = deviation(next_counts)
            assert np.allclose(rotated + arrival, next_error)
            denominator = np.linalg.norm(rotated) + np.linalg.norm(arrival)
            history[modulus].append(
                Reading(
                    depth=depth,
                    cosine=cosine(rotated, arrival),
                    cancellation=float(np.linalg.norm(next_error) / denominator)
                    if denominator
                    else float("nan"),
                    next_relative_l2=float(np.linalg.norm(next_error) / len(nxt)),
                    current_relative_l2=float(
                        np.linalg.norm(deviation(base_counts)) / len(level)
                    ),
                    nodes=len(level),
                )
            )
        level = nxt

    print("modulus,tail_cosine_min,mean,max,cancellation_mean,relative_L2_start,end,fit_exponent")
    for modulus in args.moduli:
        readings = history[modulus]
        tail = readings[-min(args.tail, len(readings)) :]
        cosines = [r.cosine for r in tail]
        cancellations = [r.cancellation for r in tail]

        fit = [r for r in readings if r.depth >= args.fit_from]
        log_nodes = np.log([r.nodes for r in fit])
        log_absolute_l2 = np.log([r.current_relative_l2 * r.nodes for r in fit])
        exponent, _ = np.polyfit(log_nodes, log_absolute_l2, 1)
        print(
            f"{modulus},{min(cosines):.9f},{sum(cosines)/len(cosines):.9f},"
            f"{max(cosines):.9f},{sum(cancellations)/len(cancellations):.9f},"
            f"{tail[0].current_relative_l2:.9f},"
            f"{tail[-1].next_relative_l2:.9f},{exponent:.9f}"
        )

    print(
        "VERDICT: through this finite horizon, rotated injections are nearly "
        "orthogonal at the finer rungs and absolute L2 deviation grows roughly "
        "like sqrt(N), so relative error falls roughly like 1/sqrt(N). "
        "COMPUTED, not proved."
    )


if __name__ == "__main__":
    main()