import json, math
import numpy as np


def tt_svd_matrix(A, out_shape, in_shape, eps=0.0, max_rank=None):
    d = len(out_shape)
    T = np.asarray(A).reshape(tuple(out_shape) + tuple(in_shape))
    # Interleave (output,input) so each TT core owns one Cartesian mode pair.
    T = np.transpose(T, sum(([k, d+k] for k in range(d)), []))
    T = T.reshape([out_shape[k]*in_shape[k] for k in range(d)])
    cores, ranks = [], [1]
    work = T
    total_sq = np.linalg.norm(work.ravel())**2
    for k in range(d-1):
        work = work.reshape(ranks[-1]*out_shape[k]*in_shape[k], -1)
        U, s, Vt = np.linalg.svd(work, full_matrices=False)
        # Distribute a global tolerance over d-1 cuts, as in the TT error bound.
        budget = (eps / math.sqrt(max(1, d-1)))**2
        tail = np.cumsum(s[::-1]**2)[::-1]
        keep = len(s) if eps == 0 else int(np.searchsorted(tail <= budget, True))
        # searchsorted on a boolean is awkward for nonmonotonic bool; explicit minimum.
        if eps != 0:
            valid = np.where(tail <= budget)[0]
            keep = int(valid[0]+1) if len(valid) else len(s)
        if max_rank is not None: keep = min(keep, max_rank)
        keep = max(1, keep)
        cores.append(U[:, :keep].reshape(ranks[-1], out_shape[k], in_shape[k], keep))
        work = (s[:keep,None] * Vt[:keep])
        ranks.append(keep)
    cores.append(work.reshape(ranks[-1], out_shape[-1], in_shape[-1], 1))
    return cores, ranks


def tt_reconstruct(cores):
    x = cores[0]
    for g in cores[1:]:
        x = np.tensordot(x, g, axes=([-1],[0]))
    # x: (1,o1,i1,...,od,id,1), remove boundaries then regroup
    x = np.squeeze(x, (0,-1))
    d = len(cores)
    x = np.transpose(x, list(range(0, 2*d, 2)) + list(range(1, 2*d, 2)))
    shape = [cores[k].shape[1] for k in range(d)] + [cores[k].shape[2] for k in range(d)]
    return x.reshape(shape)


def tt_apply(cores, x):
    """Apply a TT matrix by contracting input modes and neighboring TT ranks."""
    d = len(cores)
    # Unique einsum symbols: input/output modes plus TT ranks.
    alphabet = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
    need = 3*d + 1
    if need > len(alphabet): raise ValueError('too many modes')
    inp = alphabet[:d]
    out = alphabet[d:2*d]
    ranks = alphabet[2*d:3*d+1]
    terms = [inp]
    for k, g in enumerate(cores):
        terms.append(ranks[k] + out[k] + inp[k] + ranks[k+1])
    equation = ','.join(terms) + '->' + out
    return np.einsum(equation, x, *cores, optimize=True)

def dense_from_cores(cores):
    t = tt_reconstruct(cores)
    d = len(cores); os=t.shape[:d]; ins=t.shape[d:]
    return t.reshape(int(np.prod(os)), int(np.prod(ins)))


def params(cores): return int(sum(g.size for g in cores))


def make_spatial_operator(seed=4, n=4, rank=2):
    rng=np.random.default_rng(seed)
    # CP sum of separable 1D factors; exact spatial TT rank <= rank.
    S=np.zeros((n,n,n,n,n,n))
    for q in range(rank):
        fs=[rng.normal(size=(n,n)) for _ in range(3)]
        term=fs[0][:,:,None,None,None,None]*fs[1][None,None,:,:,None,None]*fs[2][None,None,None,None,:,:]
        S += term
    return S


def main():
    rng=np.random.default_rng(7)
    n=4; C=2; spatial=make_spatial_operator(n=n,rank=2)
    # Channel semantics: diagonal blocks share one spatial map, off-diagonal use a distinct map.
    spatial2=make_spatial_operator(seed=12,n=n,rank=2)
    # tensor mode order (channel,x,y,z), matrix axes output then input
    A=np.zeros((C,n,n,n,C,n,n,n))
    A[0, :, :, :, 0] = spatial
    A[1, :, :, :, 1] = spatial
    A[0, :, :, :, 1] = 0.35*spatial2
    A[1, :, :, :, 0] = -0.25*spatial2
    dense=A.reshape(C*n**3,C*n**3)

    # Prediction 1: TT-SVD tolerance gives global Frobenius error <= eps (up to roundoff).
    tol_rows=[]
    for eps in [1e-1, 3e-2, 1e-2, 3e-3]:
        cs,rs=tt_svd_matrix(dense,[C,n,n,n],[C,n,n,n],eps=eps)
        rec=dense_from_cores(cs)
        rel=np.linalg.norm(dense-rec)/np.linalg.norm(dense)
        tol_rows.append({'eps':eps,'relative_error':float(rel),'ranks':rs,'bound_ratio':float(rel/eps)})

    # Prediction 2: uniform TT parameter count is affine/quadratic in intermediate rank.
    rank_rows=[]
    for r in [1,2,3,4,6]:
        cs,rs=tt_svd_matrix(dense,[C,n,n,n],[C,n,n,n],max_rank=r)
        rank_rows.append({'requested_rank':r,'actual_ranks':rs,'parameters':params(cs),'param_over_r2':params(cs)/(r*r)})

    # Prediction 3: block representation preserves independent channel blocks.
    # Decompose each of four channel blocks as a spatial TT and sum storage.
    block_p=0; block_err=0.; block_details=[]
    for o in range(C):
        for i in range(C):
            b=A[o,:,:,:,i,:,:,:].reshape(n**3,n**3)
            bc,br=tt_svd_matrix(b,[n,n,n],[n,n,n],eps=1e-10)
            block_p += params(bc)
            block_err += np.linalg.norm(b-dense_from_cores(bc))**2
            block_details.append(br)
    mono, mr=tt_svd_matrix(dense,[C,n,n,n],[C,n,n,n],eps=1e-10)
    mono_err=np.linalg.norm(dense-dense_from_cores(mono))/np.linalg.norm(dense)

    # Direct contraction correctness on a small random tensor.
    x=rng.normal(size=(C,n,n,n)); y1=dense@x.reshape(-1); y2=tt_apply(mono,x).reshape(-1)
    apply_rel=np.linalg.norm(y1-y2)/np.linalg.norm(y1)
    result={'tolerance_sweep':tol_rows,'rank_sweep':rank_rows,
            'block_vs_monolithic':{'block_parameters':block_p,'block_relative_error':float(math.sqrt(block_err)/np.linalg.norm(dense)),
              'block_ranks':block_details,'monolithic_parameters':params(mono),'monolithic_ranks':mr,'monolithic_relative_error':float(mono_err)},
            'apply_relative_error':float(apply_rel),
            'predictions':{'tolerance':'relative error should be <= eps','rank':'parameters should grow approximately quadratically in interior rank','blocks':'block TT should represent independent channel blocks without channel rank mixing'}}
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()