"""MVP for a thermodynamic-confidence controller for minibatch SGD.
Run: /home/maxwelhelp/main/bin/python3 thermo_controller.py
"""
import math, random, json
import numpy as np
import torch
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

SEED = 3166

def seed_all(seed=SEED):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)

# h from Eq. (9), evaluated stably for small u.
def h(u):
    u = np.asarray(u, dtype=float)
    return 1.0 + u - np.sqrt(1.0 + 2.0*u)

def rate(a, C, D):
    """I(a)=(C/D^2)h(Da/C), with the D=0 equilibrium limit."""
    if D <= 1e-12: return a*a/(2.0*C)
    return float(C/(D*D) * h(D*a/C))

def radius(T, delta, C, D, N=1.0):
    """Smallest a for 2N exp(-T I(a)) <= delta, by bisection."""
    target = math.log(2*N/delta)/T
    if target <= 0: return 0.0
    lo, hi = 0., max(math.sqrt(2*C*target), C*target + D*target, 1e-8)
    while rate(hi,C,D) < target: hi *= 2
    for _ in range(70):
        mid=(lo+hi)/2
        if rate(mid,C,D) >= target: hi=mid
        else: lo=mid
    return hi

def autocorr_time(x):
    x=np.asarray(x,float); x=x-x.mean()
    if len(x)<4 or np.dot(x,x)<1e-14: return 1.0
    ac=np.correlate(x,x,mode='full')[len(x)-1:]/np.dot(x,x)
    # Initial-positive-sequence estimate, conservatively clipped.
    s=1.0
    for z in ac[1:]:
        if z <= 0: break
        s += 2*z
    return float(np.clip(s,1.,len(x)/2))

class CurrentController:
    def __init__(self, window=32, delta=.05, batch=32, max_batch=128):
        self.window=window; self.delta=delta; self.batch=batch; self.max_batch=max_batch
        self.values=[]; self.radii=[]; self.means=[]; self.good=0
    def observe(self, j):
        self.values.append(float(j))
        if len(self.values) >= self.window:
            x=np.array(self.values[-self.window:]); mean=float(x.mean())
            var=float(max(x.var(ddof=1),1e-10)); tau=autocorr_time(x)
            # Empirical sub-gamma proxy: effective variance inflated by correlation.
            C=2*var*tau; D=max(1e-8, 0.25*math.sqrt(var)*max(tau-1,0))
            r=radius(self.window,self.delta,C,D)
            self.radii.append(r); self.means.append(mean)
            # The prescribed relative trigger, with a tiny noise-scale floor for m ~= 0.
            if r > .25*max(abs(mean), .1*math.sqrt(var)) and self.batch < self.max_batch:
                self.batch=min(self.max_batch,self.batch*2); self.good=0
            elif r < .1*max(abs(mean), .1*math.sqrt(var)):
                self.good += 1
                if self.good >= 4: self.good=0 # permit LR increase in a full optimizer; not used here
            else: self.good=0
            return r, mean, tau
        return None

def math_check():
    # Exact bound proxy: radius should fall ~ T^-1/2 in Gaussian regime and rise with tau/variance.
    delta=.05; Ts=np.array([32,64,128,256,512])
    rs=np.array([radius(int(T),delta,C=2.,D=0.) for T in Ts])
    slope=np.polyfit(np.log(Ts),np.log(rs),1)[0]
    r_var=radius(128,delta,2.,0.)
    r_corr=radius(128,delta,8.,0.)
    return {'T':Ts.tolist(),'radii':rs.tolist(),'loglog_slope':float(slope),
            'radius_var1':float(r_var),'radius_var4':float(r_corr),
            'scaling_ok':bool(-.65 < slope < -.35),'variance_ok':bool(r_corr>r_var)}

class Net(torch.nn.Module):
    def __init__(self):
        super().__init__(); self.net=torch.nn.Sequential(torch.nn.Linear(64,64),torch.nn.ReLU(),torch.nn.Linear(64,10))
    def forward(self,x): return self.net(x)

def train(controlled, device, steps=300):
    seed_all(SEED)
    d=load_digits(); X=d.data.astype('float32')/16.; y=d.target.astype('int64')
    Xtr,Xte,ytr,yte=train_test_split(X,y,test_size=.25,random_state=SEED,stratify=y)
    xt=torch.tensor(Xtr,device=device); yt=torch.tensor(ytr,device=device)
    xval=torch.tensor(Xte,device=device); yval=torch.tensor(yte,device=device)
    model=Net().to(device); opt=torch.optim.SGD(model.parameters(),lr=.08)
    ctrl=CurrentController() if controlled else None
    v=torch.randn(sum(p.numel() for p in model.parameters()),device=device); v/=v.norm()
    rng=np.random.default_rng(SEED); losses=[]; batch_history=[]; n=len(xt)
    for step in range(steps):
        bs=ctrl.batch if ctrl else 32
        batch_history.append(bs)
        ii=torch.tensor(rng.integers(0,n,size=bs),device=device)
        opt.zero_grad(set_to_none=True); out=model(xt[ii]); loss=torch.nn.functional.cross_entropy(out,yt[ii]); loss.backward()
        g=torch.cat([p.grad.detach().flatten() for p in model.parameters()]); j=float(torch.dot(v,g).detach().cpu())
        if ctrl: ctrl.observe(j)
        opt.step(); losses.append(float(loss.detach().cpu()))
    with torch.no_grad(): acc=float((model(xval).argmax(1)==yval).float().mean().cpu())
    return {'final_train_loss':float(np.mean(losses[-20:])),'test_accuracy':acc,
            'mean_batch':float(np.mean(batch_history)),
            'total_examples':int(sum(batch_history)),
            'batch_values':sorted(set(batch_history)),
            'batches_changed':int(len(set(batch_history))),
            'radius_count':len(ctrl.radii) if ctrl else 0,
            'radius_first_last':([float(ctrl.radii[0]),float(ctrl.radii[-1])] if ctrl and ctrl.radii else [])}

def main():
    try: device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    except Exception: device=torch.device('cpu')
    try:
        check=math_check(); base=train(False,device,steps=1102); idea=train(True,device)
    except Exception as e:
        if device.type=='cuda':
            device=torch.device('cpu'); check=math_check(); base=train(False,device,steps=1102); idea=train(True,device)
        else: raise
    out={'device':str(device),'math_check':check,'baseline':base,'idea':idea}
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()