import json
import time
import numpy as np


def fourier_concentration(n, band_fraction=0.5, spatial_fraction=0.75):
    x = np.arange(n)
    F = np.exp(-2j * np.pi * np.outer(x, x) / n) / np.sqrt(n)
    freqs = np.fft.fftfreq(n)
    k = max(1, int(round(band_fraction * n / 2)))
    band = np.abs(freqs) <= (k / n)
    Q = F.conj().T @ np.diag(band.astype(float)) @ F
    m = max(1, int(round(spatial_fraction * n)))
    mask = np.zeros(n); mask[:m] = 1.0
    return (np.diag(mask) @ Q @ np.diag(mask)).real


def kron_factors(factors):
    out = factors[0]
    for f in factors[1:]: out = np.kron(out, f)
    return out


def schatten(a, p):
    s = np.linalg.svd(a, compute_uv=False)
    return float(np.sum(s ** p) ** (1.0 / p))


def sequential_apply(x, factors, shape):
    y = x.reshape(shape)
    for axis, s in enumerate(factors):
        y = np.moveaxis(y, axis, 0)
        old = y.shape
        y = (s @ y.reshape(old[0], -1)).reshape(old)
        y = np.moveaxis(y, 0, axis)
    return y.reshape(-1)


def timed(fn, repeats=7):
    vals = []
    for _ in range(repeats):
        t = time.perf_counter(); fn(); vals.append(time.perf_counter() - t)
    return float(np.median(vals))


def main():
    rng = np.random.default_rng(1441)
    contraction_rows = []
    for d in [1, 2, 3, 4, 5]:
        factors = [fourier_concentration(9 + j, .55, .7) for j in range(d)]
        eigs = [np.linalg.eigvalsh(s) for s in factors]
        # Tensor eigenvalues are products: extrema follow directly since all are nonnegative.
        contraction_rows.append({
            "d": d, "factor_min_eigenvalue": float(min(v.min() for v in eigs)),
            "factor_max_eigenvalue": float(max(v.max() for v in eigs)),
            "global_min_eigenvalue": float(np.prod([v.min() for v in eigs])),
            "global_max_eigenvalue": float(np.prod([v.max() for v in eigs])),
            "predicted_max_bound": 1.0,
        })

    tensor_rows = []
    for d in [2, 3, 4]:
        sizes = [6] * d
        factors = [fourier_concentration(n, .5, .67) for n in sizes]
        # A small explicit Kronecker matrix verifies application; no large SVD is needed.
        K = kron_factors(factors); x = rng.normal(size=6 ** d)
        exact = K @ x; sep = sequential_apply(x, factors, sizes)
        row = {"d": d, "apply_relative_error": float(np.linalg.norm(exact-sep)/np.linalg.norm(exact))}
        for p in [.5, 1., 2.]:
            lhs = schatten(K, p)
            rhs = float(np.prod([schatten(s, p) for s in factors]))
            row[f"schatten_p{p}_relative_error"] = abs(lhs-rhs) / max(rhs, 1e-15)
        tensor_rows.append(row)

    scaling_rows = []
    for d in [2, 3, 4]:
        n = 8 if d <= 3 else 5
        factors = [fourier_concentration(n, .5, .75) for _ in range(d)]
        shape = [n] * d; N = n ** d
        K = kron_factors(factors); x = rng.normal(size=N)
        dense_t = timed(lambda: K @ x)
        axial_t = timed(lambda: sequential_apply(x, factors, shape))
        dense_ops = N * N; axial_ops = N * sum(shape)
        scaling_rows.append({"d": d, "tokens": N, "dense_parameters": N*N,
            "factor_parameters": d*n*n, "dense_matvec_ops": dense_ops,
            "axis_matvec_ops": axial_ops, "predicted_cost_ratio_dense_over_axis": dense_ops/axial_ops,
            "median_dense_seconds": dense_t, "median_axis_seconds": axial_t,
            "observed_time_ratio_dense_over_axis": dense_t/max(axial_t, 1e-15)})

    n1 = n2 = 12; factors = [fourier_concentration(n1, .5, .75), fourier_concentration(n2, .5, .75)]
    K = kron_factors(factors); x = rng.normal(size=n1*n2)
    # Singular values of a tensor product are pairwise products, so this is exact and cheap.
    practical = {"tokens": n1*n2, "dense_parameters": int(K.size),
        "tensor_parameters": int(sum(s.size for s in factors)),
        "dense_apply_seconds": timed(lambda: K @ x),
        "tensor_apply_seconds": timed(lambda: sequential_apply(x, factors, [n1,n2])),
        "dense_largest_singular_value": float(np.prod([np.linalg.svd(s, compute_uv=False)[0] for s in factors])),
        "tensor_largest_singular_value": float(np.prod([np.linalg.svd(s, compute_uv=False)[0] for s in factors]))}
    print(json.dumps({"seed": 1441, "contraction_sweep": contraction_rows,
        "tensor_identity_sweep": tensor_rows, "scaling_sweep": scaling_rows,
        "practical_comparison": practical}, indent=2))

if __name__ == "__main__": main()