import json
import numpy as np

np.set_printoptions(precision=7, suppress=True)
RNG = np.random.default_rng(7)


def random_hermitian(rng, d):
    z = rng.normal(size=(d, d)) + 1j * rng.normal(size=(d, d))
    h = (z + z.conj().T) / 2
    h -= np.trace(h).real * np.eye(d) / d
    return h / np.linalg.norm(h, ord=2)


def hermitian_function_and_derivatives(theta, ops, phi=np.tanh):
    """Matrix functional calculus and its exact Frechet derivatives."""
    H = sum(t * h for t, h in zip(theta, ops))
    lam, U = np.linalg.eigh(H)
    A = (U * phi(lam)) @ U.conj().T
    G = np.empty((len(lam), len(lam)))
    for a in range(len(lam)):
        for b in range(len(lam)):
            if abs(lam[a] - lam[b]) < 1e-9:
                G[a, b] = 1.0 - np.tanh(lam[a]) ** 2
            else:
                G[a, b] = (phi(lam[a]) - phi(lam[b])) / (lam[a] - lam[b])
    dAs = []
    for h in ops:
        hb = U.conj().T @ h @ U
        dAs.append(U @ (G * hb) @ U.conj().T)
    return A, dAs


def predictions_and_grad(theta, ops, rhos, y):
    A, dAs = hermitian_function_and_derivatives(theta, ops)
    yh = np.real(np.einsum('bij,ji->b', rhos, A))
    err = yh - y
    loss = float(np.mean(err * err))
    grad = np.array([2 * np.mean(err * np.real(np.einsum('bij,ji->b', rhos, dA))) for dA in dAs])
    return loss, yh, grad


def adam_fit(ops, rhos, y, steps=700, lr=0.06, seed=1):
    rng = np.random.default_rng(seed)
    theta = rng.normal(0, .12, len(ops))
    m = np.zeros_like(theta); v = np.zeros_like(theta)
    for t in range(1, steps + 1):
        _, _, g = predictions_and_grad(theta, ops, rhos, y)
        m = .9*m + .1*g; v = .999*v + .001*g*g
        theta -= lr * (m/(1-.9**t)) / (np.sqrt(v/(1-.999**t)) + 1e-8)
    return theta, predictions_and_grad(theta, ops, rhos, y)[0]


def classical_fit(features, y, steps=700, lr=.06, seed=1):
    rng = np.random.default_rng(seed)
    w = rng.normal(0, .12, features.shape[1])
    m = np.zeros_like(w); v = np.zeros_like(w)
    for t in range(1, steps + 1):
        e = features @ w-y; g = 2*features.T@e/len(y)
        m=.9*m+.1*g; v=.999*v+.001*g*g
        w -= lr*(m/(1-.9**t))/(np.sqrt(v/(1-.999**t))+1e-8)
    return w, float(np.mean((features@w-y)**2))


def main():
    d, ntrain, ntest = 4, 180, 180
    # Generic fixed Hermitian interactions avoid the special two-Pauli identity.
    ops = [random_hermitian(RNG, d) for _ in range(3)]
    # Commuting control: same parameter count, diagonal interactions.
    cops = [np.diag(RNG.uniform(-1, 1, d)).astype(complex) for _ in range(3)]
    target_theta = np.array([.85, -.65, .55])

    def states(n):
        q = RNG.normal(size=(n,d)) + 1j*RNG.normal(size=(n,d))
        q /= np.linalg.norm(q, axis=1, keepdims=True)
        return q[:,:,None] * q[:,None,:].conj()
    rtrain, rtest = states(ntrain), states(ntest)
    At, _ = hermitian_function_and_derivatives(target_theta, ops)
    ytrain=np.real(np.einsum('bij,ji->b',rtrain,At))
    ytest=np.real(np.einsum('bij,ji->b',rtest,At))

    # Core claim: divided-difference Frechet gradient agrees with finite differences.
    theta=np.array([.31,-.22,.17]); _,_,g=predictions_and_grad(theta, ops, rtrain[:40], ytrain[:40])
    eps=1e-5; fd=[]
    for j in range(3):
        tp=theta.copy(); tm=theta.copy(); tp[j]+=eps; tm[j]-=eps
        lp=predictions_and_grad(tp,ops,rtrain[:40],ytrain[:40])[0]
        lm=predictions_and_grad(tm,ops,rtrain[:40],ytrain[:40])[0]
        fd.append((lp-lm)/(2*eps))
    fd=np.array(fd)
    grad_rel_error=float(np.linalg.norm(g-fd)/(np.linalg.norm(fd)+1e-12))
    comms=[np.linalg.norm(ops[i]@ops[j]-ops[j]@ops[i]) for i in range(3) for j in range(i)]
    ccomms=[np.linalg.norm(cops[i]@cops[j]-cops[j]@cops[i]) for i in range(3) for j in range(i)]

    results={}
    for name, model_ops in [('spectral_noncommuting',ops),('spectral_commuting',cops)]:
        theta_fit, train_loss=adam_fit(model_ops,rtrain,ytrain,seed=3)
        test_loss=predictions_and_grad(theta_fit,model_ops,rtest,ytest)[0]
        results[name]={'train_mse':train_loss,'test_mse':test_loss,'theta':theta_fit.tolist()}
    def feats(r, model_ops):
        return np.column_stack([np.ones(len(r))]+[np.real(np.einsum('bij,ji->b',r,h)) for h in model_ops])
    wc, lc=classical_fit(feats(rtrain,ops),ytrain,seed=3)
    results['classical_linear_energy']={'train_mse':lc,'test_mse':float(np.mean((feats(rtest,ops)@wc-ytest)**2)),'weights':wc.tolist()}
    out={'gradient_check':{'analytic':g.tolist(),'finite_difference':fd.tolist(),'relative_error':grad_rel_error},
         'commutator_frobenius':{'noncommuting':comms,'commuting':ccomms},'results':results,
         'setup':{'seed':7,'train':ntrain,'test':ntest,'dimension':d,'parameters':3,'activation':'tanh','target':'tanh(sum(theta_j H_j))'}}
    print(json.dumps(out, indent=2))

if __name__=='__main__': main()