import json, math, random
import numpy as np
import torch
import torch.nn.functional as F

SEED=386
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DT=torch.float64

def make_problem(n=24):
    # A spatially varying detector mask after a mild blur: right side has sparse coverage.
    yy,xx=np.mgrid[0:n,0:n]; N=n*n
    # periodic detector coverage creates a deliberately nonuniform field of view
    coverage=(0.15 + 0.85*(xx < int(.58*n)).astype(float) + 0.18*(xx>=int(.58*n)))
    # separable Gaussian blur matrix
    coords=np.arange(n)
    G=np.exp(-((coords[:,None]-coords[None,:])**2)/(2*1.35**2)); G/=G.sum(1,keepdims=True)
    B=np.kron(G,G)
    # detector rows are image pixels, with coverage weights
    m=coverage.reshape(-1)
    A=(m[:,None]**0.5)*B
    # two equal objects, one in high and one in low sensitivity areas
    u=np.zeros((n,n))
    for cy,cx in [(12,6),(12,19)]:
        u[((yy-cy)**2+(xx-cx)**2)<=3.0**2]=1.0
    # slight asymmetric shape so centroid/size are informative
    u[9:13,18:21]=1.15
    u0=u.reshape(-1)
    d=A@u0
    noise=.10*np.std(d)*np.random.randn(N)
    return A, u0, d+noise, coverage.reshape(-1)

def tv_torch(u,w,n,eps=1e-3):
    im=u.reshape(n,n)
    dx=im[:,1:]-im[:,:-1]; dy=im[1:,:]-im[:-1,:]
    # node weights on corresponding left/top nodes, a standard anisotropic boundary convention
    wx=w.reshape(n,n)
    tx=(wx[:,:-1]*torch.sqrt(dx*dx+eps*eps)).sum()
    ty=(wx[:-1,:]*torch.sqrt(dy*dy+eps*eps)).sum()
    return (tx+ty)/(n*n)

def solve(A,d,w,n,lam,steps=700,lr=.08):
    At=torch.tensor(A,dtype=DT); dt=torch.tensor(d,dtype=DT)
    u=torch.zeros(n*n,dtype=DT,requires_grad=True)
    opt=torch.optim.Adam([u],lr=lr)
    for k in range(steps):
        opt.zero_grad()
        res=At@u-dt
        loss=.5*(res@res)/len(d)+lam*tv_torch(u,w,n)
        loss.backward(); opt.step()
        with torch.no_grad(): u.clamp_(0,1.4)
    return u.detach().numpy()

def metrics(u,truth,A,d,w,n):
    im=u.reshape(n,n); gt=truth.reshape(n,n)
    mse=np.mean((u-truth)**2); psnr=10*np.log10(1.4**2/max(mse,1e-15))
    # centroid and mass/area separately in high and low sensitivity halves
    out={'psnr':float(psnr),'residual':float(np.linalg.norm(A@u-d)/np.sqrt(len(d))),
         'tv':float(tv_torch(torch.tensor(u,dtype=DT),torch.tensor(w,dtype=DT),n).item())}
    for name,sel in [('high',np.arange(n)<int(.58*n)),('low',np.arange(n)>=int(.58*n))]:
        mask=np.broadcast_to(sel[:,None],(n,n)) if False else np.broadcast_to(sel[None,:],(n,n))
        # mask columns; centroid x and thresholded area, using positive reconstructed mass
        mass=np.maximum(im,0)[mask].sum(); true_mass=gt[mask].sum()
        xs=np.tile(np.arange(n),(n,1))[mask]
        cx=(xs*np.maximum(im,0)[mask]).sum()/max(mass,1e-12)
        truecx=(xs*gt[mask]).sum()/max(true_mass,1e-12)
        area=(im[mask]>.5).sum(); truearea=(gt[mask]>.5).sum()
        out[name+'_mass_ratio']=float(mass/max(true_mass,1e-12)); out[name+'_centroid_abs']=float(abs(cx-truecx)); out[name+'_area_ratio']=float(area/max(truearea,1))
    return out

def main():
    n=24; A,truth,d,coverage=make_problem(n)
    # Exact discrete sensitivity and direct numerical response check.
    s=np.linalg.norm(A,axis=0); delta=1e-10; w=(s+delta)/np.mean(s+delta)
    rng=np.random.default_rng(SEED+1); inds=rng.choice(n*n,12,replace=False)
    observed=np.array([np.linalg.norm(A[:,i]) for i in inds]); relerr=float(np.max(abs(observed-s[inds])/(observed+1e-15)))
    # finite perturbation response K(u+h e_i)-K(u), confirming linear scaling and weights
    h=1e-4; base=A@truth
    fd=np.array([np.linalg.norm((A@(truth+h*np.eye(n*n)[i])-base)/h) for i in inds])
    scale_err=float(np.max(abs(fd-observed)/(observed+1e-15)))
    mathcheck={'max_column_norm_relative_error':relerr,'max_finite_difference_relative_error':scale_err,
               'sensitivity_min_max_ratio':float(s.min()/s.max()),'weight_min_max': [float(w.min()),float(w.max())]}
    # same lambda, plus a small validation-like sweep for scalar TV vs weighted TV
    torch.manual_seed(SEED)
    zero=np.zeros(n*n); results={}
    for name,ww,lam in [('no_tv',zero,0.0),('ordinary_tv',np.ones(n*n),0.030),('weighted_tv',w,0.030)]:
        rec=solve(A,d,torch.tensor(ww,dtype=DT),n,lam)
        results[name]=metrics(rec,truth,A,d,ww,n)
    # lambda sweep illustrates whether improvement is merely coefficient tuning
    sweep=[]
    for lam in [0.01,0.02,0.03,0.05,0.08]:
        for name,ww in [('ordinary',np.ones(n*n)),('weighted',w)]:
            rec=solve(A,d,torch.tensor(ww,dtype=DT),n,lam,steps=500)
            mm=metrics(rec,truth,A,d,ww,n); sweep.append({'method':name,'lambda':lam,**mm})
    best_ord=max([x for x in sweep if x['method']=='ordinary'],key=lambda x:x['psnr'])
    best_w=max([x for x in sweep if x['method']=='weighted'],key=lambda x:x['psnr'])
    report={'seed':SEED,'n':n,'noise_std':float(.10*np.std(A@truth)),'mathcheck':mathcheck,'results':results,'sweep_best_psnr':{'ordinary':best_ord,'weighted':best_w}}
    with open('results.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()