import json, math
import numpy as np
from block_tt3d import tt_svd_matrix, dense_from_cores, tt_apply, params


def relerr(a, b): return float(np.linalg.norm(a-b) / np.linalg.norm(a))


def main():
    rng = np.random.default_rng(123)
    modes = [2, 4, 4, 4]  # channel,x,y,z
    N = int(np.prod(modes))
    # Full-rank controlled-spectrum operator: no accidental exact low-rank result.
    q, _ = np.linalg.qr(rng.normal(size=(N, N)))
    p, _ = np.linalg.qr(rng.normal(size=(N, N)))
    singular = np.geomspace(1.0, 1e-4, N)
    A = q @ np.diag(singular) @ p.T
    normA = np.linalg.norm(A)

    tolerance = []
    for frac in [0.50, 0.20, 0.10, 0.05, 0.01]:
        eps = frac * normA
        cores, ranks = tt_svd_matrix(A, modes, modes, eps=eps)
        Ar = dense_from_cores(cores)
        e = relerr(A, Ar)
        tolerance.append({'relative_eps': frac, 'observed_relative_error': e,
                          'bound_ratio': e/frac, 'ranks': ranks,
                          'parameters': params(cores)})

    rank_sweep = []
    for r in [1, 2, 4, 8, 16, 32]:
        cores, ranks = tt_svd_matrix(A, modes, modes, max_rank=r)
        e = relerr(A, dense_from_cores(cores))
        rank_sweep.append({'requested_rank': r, 'ranks': ranks,
                           'parameters': params(cores), 'relative_error': e,
                           'parameter_ratio_to_r2': params(cores)/(r*r)})

    # Block-TT test: four channel blocks, each spatial 4x4x4 operator.
    # Keep the same global matrix and compare independent block storage to one TT.
    T = A.reshape(modes + modes)
    block_parameters = 0
    block_sqerr = 0.0
    block_ranks = []
    for o in range(2):
        for i in range(2):
            B = T[o, :, :, :, i, :, :, :].reshape(64, 64)
            bc, br = tt_svd_matrix(B, modes[1:], modes[1:], eps=0.05*np.linalg.norm(B))
            block_parameters += params(bc)
            block_sqerr += np.linalg.norm(B-dense_from_cores(bc))**2
            block_ranks.append(br)
    mono, mr = tt_svd_matrix(A, modes, modes, eps=0.05*normA)
    mono_rec = dense_from_cores(mono)

    # Apply correctness and a tiny standard dense baseline memory comparison.
    x = rng.normal(size=tuple(modes))
    apply_error = relerr(A @ x.reshape(-1), tt_apply(mono, x).reshape(-1))
    dense_parameters = A.size
    result = {
      'operator': {'shape': [N, N], 'dense_parameters': dense_parameters, 'frobenius_norm': float(normA)},
      'prediction_checks': {
        'svd_bound': 'relative error <= requested relative epsilon',
        'rank_scaling': 'interior TT storage scales approximately as r^2',
        'block_semantics': 'separate blocks can be rounded independently, but may cost more storage'},
      'tolerance_sweep': tolerance,
      'rank_sweep': rank_sweep,
      'block_vs_monolithic': {
        'block_parameters': block_parameters,
        'block_relative_error': math.sqrt(block_sqerr)/normA,
        'block_ranks': block_ranks,
        'monolithic_parameters': params(mono),
        'monolithic_relative_error': relerr(A, mono_rec),
        'dense_parameters': dense_parameters},
      'apply_relative_error': apply_error
    }
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()