import json
import time
from pathlib import Path
import numpy as np
from scipy.optimize import linear_sum_assignment


def standardize_features(F):
    F = np.asarray(F, dtype=np.float64)
    F = F - F.mean(axis=0, keepdims=True)
    s = F.std(axis=0, keepdims=True)
    s[s < 1e-12] = 1.0
    return F / s


def cosine_score(X, Y):
    Xn = X / np.maximum(np.linalg.norm(X, axis=1, keepdims=True), 1e-12)
    Yn = Y / np.maximum(np.linalg.norm(Y, axis=1, keepdims=True), 1e-12)
    return Xn @ Yn.T


def tree_features(X, max_depth=5, branches=(2, 3, 4), normalize=True):
    """Finite rooted-tree contractions computed from the sample Gram matrix.

    A coordinate vertex of degree two contracts to A=XX^T/d. Repeated A
    multiplication gives rooted path contractions; powers of A1 represent
    wide rooted trees with repeated branches. Thus all features are exactly
    invariant under right multiplication X -> XQ for orthogonal Q.
    """
    X = np.asarray(X, dtype=np.float64)
    A = X @ X.T / max(X.shape[1], 1)
    one = np.ones((X.shape[0], 1), dtype=np.float64)
    messages, v = [], one
    for _ in range(max_depth):
        v = A @ v
        messages.append(v[:, 0])
    root = (A @ one)[:, 0]
    for b in branches:
        messages.append(root ** b)
    F = np.stack(messages, axis=1)
    return standardize_features(F) if normalize else F


def tree_score(X, Y):
    FX, FY = tree_features(X), tree_features(Y)
    return FX @ FY.T / FX.shape[1]


def hungarian_accuracy(score, truth=None):
    r, c = linear_sum_assignment(-score)
    if truth is None:
        truth = np.arange(score.shape[0])
    return float(np.mean(c == np.asarray(truth)[r]))


def random_orthogonal(d, rng):
    q, r = np.linalg.qr(rng.normal(size=(d, d)))
    q *= np.sign(np.diag(r))[None, :]
    return q


def make_pair(n, d, rho, rng):
    X = rng.normal(size=(n, d))
    Q = random_orthogonal(d, rng)
    # Y has row-wise correlation rho with the rotated matched X.
    Y = rho * (X @ Q) + np.sqrt(max(1 - rho*rho, 0.0)) * rng.normal(size=(n, d))
    return X, Y


def run(seed=7, n=96, d=64, rhos=(0.55, 0.65, 0.75, 0.85), repeats=8):
    rng = np.random.default_rng(seed)
    invariant_errors, rows = [], []
    X = rng.normal(size=(n, d)); Q = random_orthogonal(d, rng)
    f0 = tree_features(X, normalize=False)
    f1 = tree_features(X @ Q, normalize=False)
    invariant_errors.append(float(np.max(np.abs(f0-f1)) / max(np.max(np.abs(f0)), 1e-12)))
    for rho in rhos:
        ca, ta = [], []
        for _ in range(repeats):
            X, Y = make_pair(n, d, rho, rng)
            t = time.perf_counter(); cs = cosine_score(X, Y); tc = time.perf_counter()-t
            t = time.perf_counter(); ts = tree_score(X, Y); tt = time.perf_counter()-t
            ca.append(hungarian_accuracy(cs)); ta.append(hungarian_accuracy(ts))
        rows.append({'rho': rho, 'cosine_mean': float(np.mean(ca)), 'cosine_std': float(np.std(ca)),
                     'tree_mean': float(np.mean(ta)), 'tree_std': float(np.std(ta)),
                     'cosine_sec': tc, 'tree_sec': tt})
    out = {'invariance_relative_max_error': invariant_errors[0], 'rows': rows}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))
    return out


if __name__ == '__main__':
    run()