import json, math, random
import numpy as np
import torch
from torch import nn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

SEED = 371
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device.type == 'cuda':
        torch.cuda.get_device_properties(0)
except Exception:
    device = torch.device('cpu')


def binomial_probs(x, n):
    """Enumerated P(Y=y|X=x), with a differentiable torch implementation."""
    x = x.clamp(1e-6, 1-1e-6)
    y = torch.arange(n + 1, device=x.device, dtype=x.dtype)
    logc = torch.lgamma(torch.tensor(float(n + 1), device=x.device, dtype=x.dtype)) \
        - torch.lgamma(y + 1) - torch.lgamma(torch.tensor(float(n), device=x.device, dtype=x.dtype) - y + 1)
    return torch.exp(logc + x[..., None] * 0 + y * torch.log(x[..., None])
                     + (n-y) * torch.log1p(-x[..., None]))


def beta_binomial_prior(n):
    y = np.arange(n + 1, dtype=float)
    # comb(n,y) B(y+1/2,n-y+1/2)/B(1/2,1/2), evaluated stably
    logq = (math.lgamma(n+1) - np.array([math.lgamma(v+1) for v in y])
            - np.array([math.lgamma(n-v+1) for v in y])
            + np.array([math.lgamma(v+.5) + math.lgamma(n-v+.5) - math.lgamma(n+1)
                        for v in y]) - math.log(math.pi))
    q = np.exp(logq); return q / q.sum()


def entropy(p):
    return -(p.clamp_min(1e-12) * p.clamp_min(1e-12).log()).sum(-1)


def bottleneck_terms(x, n):
    p = binomial_probs(x, n)
    q = p.mean(0)
    I = (entropy(q) - entropy(p).mean(0)).mean()
    qr = torch.tensor(beta_binomial_prior(n), device=x.device, dtype=x.dtype)
    kl = (q.clamp_min(1e-12) * (q.clamp_min(1e-12).log() - qr.log())).sum(-1).mean()
    return p, I, kl


def derivative_check():
    # Check Lemma 7 for several smooth functions and interior probabilities.
    n, x = 11, .37
    vals = []
    for f in [lambda y: float(y*y), lambda y: math.exp(.08*y), lambda y: math.sin(.4*y)]:
        def expectation(z, trials):
            return sum(math.comb(trials, y)*z**y*(1-z)**(trials-y)*f(y) for y in range(trials+1))
        h = 1e-5
        lhs = (expectation(x+h,n)-expectation(x-h,n))/(2*h)
        rhs = n * sum(math.comb(n-1,y)*x**y*(1-x)**(n-1-y)*(f(y+1)-f(y)) for y in range(n))
        vals.append(abs(lhs-rhs))
    q = beta_binomial_prior(n)
    return {"max_abs_derivative_error": float(max(vals)), "prior_sum": float(q.sum()),
            "prior_endpoint_mass": float(q[0]+q[-1]), "prior": q.tolist()}


class Net(nn.Module):
    def __init__(self, idea=False, n=8):
        super().__init__(); self.idea=idea; self.n=n
        self.enc=nn.Sequential(nn.Linear(64,32), nn.ReLU(), nn.Linear(32,8))
        self.head=nn.Sequential(nn.Linear(8,16), nn.ReLU(), nn.Linear(16,10))
    def forward(self, z):
        x=torch.sigmoid(self.enc(z))
        if self.idea:
            p,I,KL=bottleneck_terms(x,self.n)
            return self.head(x), I, KL, x, p
        return self.head(x), None, None, x, None


def run(idea, epochs=45, n=8):
    X,y=load_digits(return_X_y=True)
    X=X.astype('float32')/16.0
    xtr,xte,ytr,yte=train_test_split(X,y,test_size=.25,random_state=SEED,stratify=y)
    tr=torch.tensor(xtr,device=device); ty=torch.tensor(ytr,device=device,dtype=torch.long)
    te=torch.tensor(xte,device=device); ey=torch.tensor(yte,device=device,dtype=torch.long)
    model=Net(idea,n).to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3,weight_decay=1e-4)
    bs=96; history=[]
    for ep in range(epochs):
        model.train(); perm=torch.randperm(len(tr),device=device); total=0
        for ix in perm.split(bs):
            out,I,KL,_,_=model(tr[ix]); ce=nn.functional.cross_entropy(out,ty[ix])
            # modest shaping: maximize exact batch MI and match the prescribed prior
            loss=ce + ((.04*KL - .01*I) if idea else 0)
            opt.zero_grad(); loss.backward(); opt.step(); total += ce.item()*len(ix)
        if ep in (epochs//2,epochs-1): history.append(total/len(tr))
    model.eval()
    with torch.no_grad():
        out,I,KL,z,p=model(te); pred=out.argmax(1); acc=(pred==ey).float().mean().item()
        # evaluate information and prior divergence over the full test set in batches
        ps=[]
        for ix in torch.arange(len(te),device=device).split(bs): ps.append(binomial_probs(z[ix],n))
        pp=torch.cat(ps); q=pp.mean(0); qr=torch.tensor(beta_binomial_prior(n),device=device,dtype=pp.dtype)
        info=(entropy(q)-entropy(pp).mean(0)).mean().item(); kl=(q*(q.clamp_min(1e-12).log()-qr.log())).sum(-1).mean().item()
        endpoint=(q[:,0]+q[:,-1]).mean().item()
        sampled_acc = None
        if idea:
            # Discrete deployment path: sample one count per latent coordinate.
            ys = torch.multinomial(pp.reshape(-1, n+1), 1).reshape(len(te), 8).float() / n
            sampled_acc = (model.head(ys).argmax(1) == ey).float().mean().item()
    return {"accuracy":acc,"sampled_count_accuracy":sampled_acc,
            "test_I_nats":info,"test_KL_to_beta_binomial":kl,
            "test_endpoint_mass":endpoint,"train_ce_checkpoints":history,
            "activation_storage_bits_per_coordinate":(math.log2(n+1) if idea else 32.0)}


if __name__ == '__main__':
    # Keep the typo-prone implementation path above explicit and fail-safe by running checks first.
    result={"device":str(device),"math_check":derivative_check()}
    result["baseline_continuous"] = run(False)
    result["capacity_shaped_binomial"] = run(True)
    print(json.dumps(result, indent=2))