import json
from pathlib import Path
import numpy as np
from scipy.fft import dctn, idctn

RNG = np.random.default_rng(7)

def steering(n, u):
    x = np.arange(n) - (n - 1) / 2
    return np.exp(1j * np.pi * x * u) / np.sqrt(n)

def steering_deriv(n, u):
    x = np.arange(n) - (n - 1) / 2
    return 1j * np.pi * x * steering(n, u)

def atom(nr, nt, ur, ut):
    return np.outer(steering(nr, ur), steering(nt, ut).conj())

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

def channel(nr, nt, k, rng):
    H = np.zeros((nr, nt), complex)
    paths = []
    for _ in range(k):
        ur, ut = rng.uniform(-.85, .85, 2)
        g = (rng.normal() + 1j * rng.normal()) / np.sqrt(2 * k)
        H += g * atom(nr, nt, ur, ut)
        paths.append((ur, ut, g))
    return H, paths

def ls_reconstruct(H, atoms, lam=1e-8):
    A = np.stack([x.reshape(-1) for x in atoms], axis=1)
    y = H.reshape(-1)
    g = np.linalg.solve(A.conj().T @ A + lam * np.eye(len(atoms)), A.conj().T @ y)
    return (A @ g).reshape(H.shape), g

def analytic_greedy(H, k, grid):
    residual = H.copy(); atoms = []; chosen = []
    nr, nt = H.shape
    for _ in range(k):
        best, score = None, -1.0
        for ur in grid:
            ar = steering(nr, ur)
            for ut in grid:
                a = np.outer(ar, steering(nt, ut).conj())
                s = abs(np.vdot(a, residual))
                if s > score:
                    score, best = s, (ur, ut, a)
        chosen.append(best[:2]); atoms.append(best[2])
        rec, _ = ls_reconstruct(H, atoms)
        residual = H - rec
    rec, g = ls_reconstruct(H, atoms)
    return rec, chosen, g

def dct_codec(H, k):
    # Separate real/imag orthonormal 2-D DCT, retaining largest magnitudes.
    coeff = dctn(H.real, norm='ortho') + 1j * dctn(H.imag, norm='ortho')
    flat = np.abs(coeff).ravel()
    ind = np.argpartition(flat, -k)[-k:]
    keep = np.zeros(flat.size, bool); keep[ind] = True
    sparse = np.where(keep, coeff.ravel(), 0).reshape(coeff.shape)
    return idctn(sparse.real, norm='ortho') + 1j * idctn(sparse.imag, norm='ortho')

def taylor_sweep():
    # Prediction: first-order error is O(delta^2), with relative error/delta^2
    # approaching ||a''||/(2||a||), and is only stable for |delta|*pi*N/2 << 1.
    n = 32; u = .17
    ds = np.array([1e-4, 3e-4, 1e-3, 3e-3, 1e-2, 3e-2])
    errs = []
    for d in ds:
        exact = steering(n, u+d)
        approx = steering(n,u) + d*steering_deriv(n,u)
        errs.append(np.linalg.norm(exact-approx)/np.linalg.norm(exact))
    slope = np.polyfit(np.log(ds), np.log(errs), 1)[0]
    # The quadratic coefficient from the exact second derivative.
    x=np.arange(n)-(n-1)/2
    second = -(np.pi*x)**2*steering(n,u)
    pred_coeff=np.linalg.norm(second)/(2*np.linalg.norm(steering(n,u)))
    return {'deltas':ds.tolist(),'errors':np.array(errs).tolist(),'loglog_slope':float(slope),
            'predicted_slope':2.0,'quadratic_coeff_predicted':float(pred_coeff),
            'quadratic_coeff_observed':float(np.mean(np.array(errs[:3])/ds[:3]**2))}

def rank_sweep():
    # Prediction: exact K-path channels have zero noiseless residual at K atoms;
    # with fewer atoms, residual decreases monotonically as K increases.
    rng=np.random.default_rng(11); H, paths=channel(32,32,4,rng)
    exact=[]
    for kk in range(1,5):
        rec,_=ls_reconstruct(H,[atom(32,32,p[0],p[1]) for p in paths[:kk]])
        exact.append(relerr(H,rec))
    return {'K':list(range(1,5)),'relative_errors':exact,'predicted_at_Kstar':0.0,
            'observed_at_Kstar':exact[-1]}

def ridge_sweep():
    # Prediction: nearly duplicate atoms make unregularized LS ill-conditioned;
    # ridge reduces coefficient norm and remains numerically stable.
    nr=nt=32; u=.2; eps=np.array([1e-1,1e-2,1e-3,1e-4]); rows=[]
    rng=np.random.default_rng(12); H=atom(nr,nt,u,u)+(rng.normal(size=(nr,nt))+1j*rng.normal(size=(nr,nt)))*1e-3
    for e in eps:
        A=np.stack([atom(nr,nt,u,u).ravel(),atom(nr,nt,u+e,u+e).ravel()],1)
        cond=np.linalg.cond(A.conj().T@A)
        vals=[]
        for lam in [0,1e-6,1e-3]:
            try:
                g=np.linalg.solve(A.conj().T@A+lam*np.eye(2),A.conj().T@H.ravel())
                vals.append(float(np.linalg.norm(g)))
            except np.linalg.LinAlgError: vals.append(float('inf'))
        rows.append({'separation':float(e),'condition':float(cond),'coef_norm_lambda_0_1e-6_1e-3':vals})
    slope = float(np.polyfit(np.log(eps), np.log([r['condition'] for r in rows]), 1)[0])
    return {'rows': rows, 'condition_loglog_slope': slope, 'predicted_slope': -2.0}

def codec_compare():
    rng=np.random.default_rng(21); grid=np.linspace(-.9,.9,17)
    rows=[]
    for k in [1,2,3,4]:
        ae=[]; de=[]
        for _ in range(10):
            H,_=channel(32,32,3,rng)
            ae.append(relerr(H,analytic_greedy(H,k,grid)[0]))
            de.append(relerr(H,dct_codec(H, k*4)))
        rows.append({'atoms':k,'analytic_nmse':float(np.mean(np.array(ae)**2)),
                     'dct_coefficients':4*k,'dct_nmse':float(np.mean(np.array(de)**2))})
    # Transfer: same continuous physical paths, reconstruct at larger dimensions without decoder retraining.
    H32,p=channel(32,32,3,np.random.default_rng(22)); H48=np.zeros((48,48),complex)
    for ur,ut,g in p: H48 += g*atom(48,48,ur,ut)
    transfer=relerr(H48,sum(g*atom(48,48,ur,ut) for ur,ut,g in p))
    return rows, transfer

def main():
    out={'taylor':taylor_sweep(),'rank':rank_sweep(),'ridge':ridge_sweep()}
    out['codec_rows'],out['transfer_relative_error']=codec_compare()
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()