"""Fast condition number analysis - minimal computation."""
import numpy as np
from scipy.linalg import expm
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import time

sx = np.array([[0,1],[1,0]], dtype=complex)
sy = np.array([[0,-1j],[1j,0]], dtype=complex)
sz = np.array([[1,0],[0,-1]], dtype=complex)
I2 = np.eye(2, dtype=complex)

def tensor_op(op, site, N):
    r = np.eye(1, dtype=complex)
    for i in range(N):
        r = np.kron(r, op if i == site else I2)
    return r

def two_site_op(o1, o2, s1, s2, N):
    r = np.eye(1, dtype=complex)
    for i in range(N):
        if i == s1: r = np.kron(r, o1)
        elif i == s2: r = np.kron(r, o2)
        else: r = np.kron(r, I2)
    return r

N = 5  # smaller system for speed
dim = 2**N
print(f"System: N={N}, dim={dim}")

# Build H
H = np.zeros((dim,dim), dtype=complex)
H += 1.0*(two_site_op(sx,sx,0,1,N) + two_site_op(sy,sy,0,1,N) + two_site_op(sz,sz,0,1,N))
for i in range(2, N-1):
    H += 0.5*two_site_op(sx,sx,i,i+1,N)
H += 0.45*two_site_op(sz,sz,1,2,N)
for i in range(N):
    H += 0.12*tensor_op(sz,i,N) + 0.08*tensor_op(sx,i,N)

O0 = tensor_op(sz, 0, N)

# Compute moments directly: m_k = Tr[O0 * L^k(O0)] / Tr[O0^2]
print("Computing moments...")
norm = np.trace(O0 @ O0).real / dim
moments = [1.0]
L_k = O0.copy()
for k in range(1, 26):
    L_k = 1j*(H @ L_k - L_k @ H)
    m = (np.trace(O0.conj().T @ L_k).real / dim) / norm
    moments.append(m)
moments = np.array(moments)
print(f"Moments computed: {len(moments)}")
print(f"Even moments: {moments[0::2][:8]}")

# EXPERIMENT 1: Hankel condition numbers
print("\n--- Condition numbers ---")
cond_nums = []
n_range = range(2, 12)
for n in n_range:
    Hk = np.zeros((n,n))
    for i in range(n):
        for j in range(n):
            idx = 2*(i+j)
            if idx < len(moments):
                Hk[i,j] = moments[idx]
    svs = np.linalg.svd(Hk, compute_uv=False)
    c = svs[0]/max(svs[-1], 1e-30)
    cond_nums.append(c)
    print(f"  n={n:2d}: kappa = {c:.4e}")

# EXPERIMENT 2: Forward perturbation sensitivity
print("\n--- Forward: perturb H -> change in C_op ---")
tau_vals = np.linspace(0, 0.8, 50)

def autocorr(H, O0, taus):
    d = H.shape[0]
    nm = np.trace(O0@O0).real/d
    C = []
    for t in taus:
        U = expm(-1j*H*t)
        Ot = U.conj().T @ O0 @ U
        C.append((np.trace(Ot@O0).real/d)/nm)
    return np.array(C)

C0 = autocorr(H, O0, tau_vals)

eps_list = [0.001, 0.01, 0.05, 0.1, 0.3, 0.5]
fwd_errs = []
for eps in eps_list:
    np.random.seed(42)
    dH = eps*np.random.randn(dim,dim).astype(complex)
    dH = (dH+dH.conj().T)/2
    Cp = autocorr(H+dH, O0, tau_vals)
    e = np.sqrt(np.mean((C0-Cp)**2))
    fwd_errs.append(e)
    print(f"  eps={eps:.3f}: RMSE(C) = {e:.6f}, ratio = {e/eps:.3f}")

# EXPERIMENT 3: Inverse - perturb moments, recover b_n
print("\n--- Inverse: perturb moments -> recover b_n ---")

# Get clean b_n via Lanczos
def lanczos(H, O0, K=15):
    d = H.shape[0]
    def ip(A,B): return np.trace(A.conj().T@B).real/d
    n0 = np.sqrt(ip(O0,O0))
    Op = np.zeros_like(O0); Oc = O0/n0
    bs = []
    for _ in range(K):
        An = 1j*(H@Oc - Oc@H)
        if bs: An -= bs[-1]*Op
        bn = np.sqrt(ip(An,An))
        if bn < 1e-12: break
        bs.append(bn)
        Op = Oc; Oc = An/bn
    return np.array(bs)

b_clean = lanczos(H, O0, 15)
print(f"Clean b_n: {b_clean[:6]}")

inv_errs = []
for eps in eps_list:
    errs = []
    for trial in range(10):
        mp = moments.copy()
        for k in range(0, len(mp), 2):
            mp[k] *= (1 + eps*np.random.randn())
        
        n_h = min(6, len(b_clean))
        Hk = np.zeros((n_h, n_h))
        for i in range(n_h):
            for j in range(n_h):
                idx = 2*(i+j)
                if idx < len(mp): Hk[i,j] = mp[idx]
        
        try:
            L = np.linalg.cholesky(Hk)
            b_rec = [abs(L[k+1,k]/L[k,k]) for k in range(min(n_h-1, len(b_clean))) if abs(L[k,k]) > 1e-15]
            if b_rec:
                nc = min(len(b_rec), len(b_clean))
                errs.append(np.sqrt(np.mean((b_clean[:nc]-np.array(b_rec[:nc]))**2)))
        except:
            errs.append(float('inf'))
    
    valid = [e for e in errs if np.isfinite(e)]
    me = np.mean(valid) if valid else float('nan')
    inv_errs.append(me)
    nfail = sum(1 for e in errs if not np.isfinite(e))
    print(f"  eps={eps:.3f}: RMSE(b) = {me:.6f}, ratio = {me/eps:.1f}, fails = {nfail}/10")

# PLOTTING
fig, axes = plt.subplots(1, 3, figsize=(15, 4.5))

# (a) Condition numbers
ax = axes[0]
ns = list(n_range)
ax.semilogy(ns, cond_nums, 'o-', color='#c0392b', lw=2, ms=7)
lc = np.log(cond_nums)
co = np.polyfit(ns, lc, 1)
gr = np.exp(co[0])
ax.semilogy(ns, np.exp(np.polyval(co, ns)), '--', color='gray', label=f'Fit: $\\kappa_n \\sim {gr:.2f}^n$')
ax.set_xlabel('Hankel matrix size n')
ax.set_ylabel('Condition number')
ax.set_title('(a) Exponential ill-conditioning')
ax.legend()

# (b) Forward vs Inverse
ax = axes[1]
ax.loglog(eps_list, fwd_errs, 's-', color='#2980b9', lw=2, ms=7, label='Forward: $\\delta H \\to \\delta C$')
valid_inv = [(e,f) for e,f in zip(eps_list, inv_errs) if np.isfinite(f)]
if valid_inv:
    ei, fi = zip(*valid_inv)
    ax.loglog(ei, fi, 'D-', color='#c0392b', lw=2, ms=7, label='Inverse: $\\delta m_k \\to \\delta b_n$')
ax.set_xlabel('Perturbation size')
ax.set_ylabel('Error')
ax.set_title('(b) Forward vs. inverse sensitivity')
ax.legend()

# (c) Autocorrelation fingerprint
ax = axes[2]
ax.plot(tau_vals, C0, 'k-', lw=2.5, label='Clean')
for eps, col, ls in [(0.01,'#27ae60','--'),(0.1,'#e67e22','-.'),(0.5,'#c0392b',':')]:
    np.random.seed(int(eps*1000))
    dH = eps*np.random.randn(dim,dim).astype(complex)
    dH = (dH+dH.conj().T)/2
    Cp = autocorr(H+dH, O0, tau_vals)
    ax.plot(tau_vals, Cp, color=col, ls=ls, lw=1.5, label=f'$\\|\\delta H\\|={eps}$')
ax.set_xlabel('$\\tau$')
ax.set_ylabel('$C_{\\mathrm{op}}(\\tau)$')
ax.set_title('(c) Fingerprint sensitivity')
ax.legend(fontsize=9)

plt.tight_layout()
plt.savefig('/home/ubuntu/owf_condition_analysis.png', dpi=200, bbox_inches='tight')
print("\nFigure saved.")

print(f"\n{'='*60}")
print(f"SUMMARY")
print(f"{'='*60}")
print(f"Condition number growth: ~{gr:.2f}^n")
print(f"Forward: smooth, linear scaling (ratio ~ constant)")
print(f"Inverse: exponentially amplified errors")
print(f"This asymmetry = ONE-WAY FUNCTION signature")