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


def euler_soft(x, thresholds, temperature=0.1):
    if x.ndim == 3: x = x[:, None]
    t = torch.as_tensor(thresholds, device=x.device, dtype=x.dtype)
    s = torch.sigmoid((x[:, None] - t[None, :, None, None, None]) / temperature)
    p1 = s.mean((-1, -2, -3))
    eh = (s[..., :, 1:] * s[..., :, :-1]).mean((-1, -2, -3))
    ev = (s[..., 1:, :] * s[..., :-1, :]).mean((-1, -2, -3))
    f = (s[..., :-1, :-1] * s[..., 1:, :-1] * s[..., :-1, 1:] * s[..., 1:, 1:]).mean((-1, -2, -3))
    return p1 - eh - ev + f


def euler_hard(x, thresholds):
    if x.ndim == 3: x = x[:, None]
    vals = []
    for v in thresholds:
        s = (x > v).float()
        p1 = s.mean((-1, -2, -3))
        eh = (s[..., :, 1:] * s[..., :, :-1]).mean((-1, -2, -3))
        ev = (s[..., 1:, :] * s[..., :-1, :]).mean((-1, -2, -3))
        f = (s[..., :-1, :-1] * s[..., 1:, :-1] * s[..., :-1, 1:] * s[..., 1:, 1:]).mean((-1, -2, -3))
        vals.append(p1 - eh - ev + f)
    return torch.stack(vals, 1)


def shapes(n, h=16, seed=0):
    rng = np.random.default_rng(seed); yy, xx = np.mgrid[:h, :h]; out=[]
    for _ in range(n):
        im = np.zeros((h,h), np.float32)
        if rng.random() < .5:
            cx,cy=rng.uniform(4,h-4,2); r=rng.uniform(2.5,5.5)
            im[((xx-cx)**2+(yy-cy)**2)<r*r]=1
        else:
            x0,y0=rng.integers(2,h-8,2); w,hh=rng.integers(3,8,2)
            im[y0:y0+hh,x0:x0+w]=1
        im += rng.normal(0,.08,(h,h)).astype(np.float32)
        out.append(np.clip(im,0,1))
    return torch.tensor(np.array(out)[:,None])

class Net(nn.Module):
    def __init__(self):
        super().__init__(); self.a=nn.Conv2d(1,16,3,padding=1); self.b=nn.Conv2d(16,16,3,padding=1); self.c=nn.Conv2d(16,1,3,padding=1)
    def forward(self,x): return torch.sigmoid(self.c(F.relu(self.b(F.relu(self.a(x))))))

def main():
    torch.set_num_threads(2)
    random.seed(7); np.random.seed(7); torch.manual_seed(7)
    thresholds=[.2,.4,.6,.8]; results={}
    # Check 1: T -> 0 agrees with hard excursion away from threshold ties.
    x=torch.rand(32,1,12,12)*.8+.1; hard=euler_hard(x,thresholds)
    errs=[]
    for T in [.2,.1,.05,.02,.01]: errs.append(float((euler_soft(x,thresholds,T)-hard).abs().mean()))
    results['temperature_errors']={'T':[.2,.1,.05,.02,.01],'error':errs}
    # Check 2: constant field has chi density 1 if occupied, 0 otherwise.
    const=torch.full((3,1,10,11),.9); constv=float(euler_soft(const,[.5],.01).mean())
    emptyv=float(euler_soft(torch.zeros_like(const),[.5],.01).mean())
    results['constant_field']={'occupied':constv,'empty':emptyv,'predicted':[0.0,0.0]}
    # Check 3: iid Bernoulli occupancy predicts q-2q^2+q^4.
    q=.35; g=torch.Generator().manual_seed(11); z=(torch.rand((512,1,48,48),generator=g)<q).float()
    empirical=float(euler_hard(z,[.5]).mean()); predicted=q-2*q*q+q**4
    results['iid_bernoulli']={'q':q,'empirical':empirical,'predicted':predicted,'abs_error':abs(empirical-predicted),'note':'finite-grid variance and overlapping windows make exact agreement approximate'}
    # Small equal-step denoising test; targets are clean morphology statistics.
    clean=shapes(128,16,3); noisy=(clean+.22*torch.randn(clean.shape)).clamp(0,1)
    real_target=euler_soft(clean,thresholds,.05).mean(0).detach()
    def train(lam):
        torch.manual_seed(19); net=Net(); opt=torch.optim.Adam(net.parameters(),lr=.004)
        for step in range(80):
            ix=torch.randint(0,len(clean),(32,)); pred=net(noisy[ix]); loss=F.mse_loss(pred,clean[ix])
            if lam: loss=loss+lam*((euler_soft(pred,thresholds,.1).mean(0)-real_target)**2).mean()
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            pred=net(noisy); mse=float(F.mse_loss(pred,clean)); ee=float((euler_hard(pred,thresholds)-euler_hard(clean,thresholds)).abs().mean()); soft=float(((euler_soft(pred,thresholds,.1).mean(0)-real_target)**2).mean())
        return {'mse':mse,'hard_euler_mae':ee,'soft_target_mse':soft}
    results['denoising']={'baseline':train(0.0),'euler_lambda_0.2':train(.2)}
    with open('results.json','w') as f: json.dump(results,f,indent=2)
    print(json.dumps(results,indent=2))

if __name__=='__main__': main()