import json, math, random, time
import numpy as np
import torch
from torch import nn

SEED = 1254
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(min(8, torch.get_num_threads()))
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
try:
    if DEVICE == "cuda":
        torch.cuda.empty_cache()
        torch.zeros(1, device=DEVICE)
except Exception:
    DEVICE = "cpu"

def exact_incidence_check():
    rng = np.random.default_rng(SEED)
    A = rng.integers(0, 2, size=(17, 23)).astype(np.float32)
    A[A.sum(1) == 0, 0] = 1
    s = rng.normal(size=23).astype(np.float32)
    dense = A @ s
    # Same operation expressed as row-wise sparse entries.
    sparse = np.array([s[row.astype(bool)].sum() for row in A])
    return {"max_abs_error": float(np.max(np.abs(dense-sparse))),
            "all_nonempty_rows": bool(np.all(A.sum(1) > 0))}

def normalization_sweep():
    # For s_m=c, a k-atom action has logit c*k^(1-alpha).
    c, eps = 2.0, 1e-12
    ks = np.arange(1, 33, dtype=float)
    rows = []
    for alpha in (0.0, 0.5, 1.0):
        observed = c * ks / (ks + eps) ** alpha
        predicted = c * ks ** (1-alpha)
        slope = float(np.polyfit(np.log(ks), np.log(observed), 1)[0])
        rel = float(np.max(np.abs(observed-predicted) / np.maximum(1e-9, np.abs(predicted))))
        rows.append({"alpha": alpha, "predicted_loglog_slope": 1-alpha,
                     "observed_loglog_slope": slope, "max_relative_error": rel,
                     "logit_k1": float(observed[0]), "logit_k32": float(observed[-1])})
    return rows

def softmax_bias_sweep():
    # Equal atom scores: raw sum's odds between k and 1 atom are exp((k-1)c).
    c = 0.25
    out=[]
    for alpha in (0., .5, 1.):
        vals=[]
        for k in (1,2,4,8,16):
            l1=c
            lk=c*k/(k**alpha)
            vals.append(math.exp(lk-l1))
        pred_last=math.exp(c*(16**(1-alpha)-1))
        out.append({"alpha":alpha, "odds_k16_vs_k1_observed":vals[-1],
                    "odds_k16_vs_k1_predicted":pred_last,
                    "odds_all_sizes":vals})
    return out

def make_batch(n, max_atoms, d, max_actions, rng):
    # Each target action is the additive utility of its distinct atomic entities.
    xs=[]; inc=[]; labels=[]
    true_w=np.array([1.1,-.8,.6,.35], dtype=np.float32)[:d]
    for _ in range(n):
        na=int(rng.integers(3,max_atoms+1)); pa=int(rng.integers(3,max_actions+1))
        x=rng.normal(size=(na,d)).astype(np.float32)
        A=np.zeros((pa,na),np.float32)
        for p in range(pa):
            cnt=int(rng.integers(1,min(na,5)+1)); A[p,rng.choice(na,cnt,replace=False)]=1
        util=A @ (x@true_w) + .08*rng.normal(size=pa)
        xs.append(x); inc.append(A); labels.append(int(np.argmax(util)))
    return xs,inc,labels

class Structured(nn.Module):
    def __init__(self,d):
        super().__init__(); self.scorer=nn.Sequential(nn.Linear(d,16),nn.ReLU(),nn.Linear(16,1))
    def forward(self,x,A): return A @ self.scorer(x).squeeze(-1)

class PaddedMLP(nn.Module):
    def __init__(self,d,max_atoms,max_actions):
        super().__init__(); self.na=max_atoms; self.pa=max_actions
        self.net=nn.Sequential(nn.Linear(max_atoms*d,32),nn.ReLU(),nn.Linear(32,max_actions))
    def forward(self,x,A):
        # Standard padded action head receives a fixed-size flattened graph and
        # emits a fixed global action vector; absent actions are masked.
        z=torch.zeros(self.na,x.shape[1],device=x.device); z[:x.shape[0]]=x
        logits=self.net(z.reshape(-1)).unsqueeze(0).expand(A.shape[0],-1)
        # This baseline cannot represent instance-specific action composition;
        # use one shared graph score plus a fixed action index bias and mask.
        return logits[:, :A.shape[0]].squeeze(0)

def mini_experiment():
    rng=np.random.default_rng(SEED+1); d=4; train_n=700; test_n=250
    tr=make_batch(train_n,8,d,12,rng); te=make_batch(test_n,16,d,32,rng)
    # Baseline is intentionally a conventional padded fixed-action MLP.
    models=[("structured",Structured(d).to(DEVICE)),("padded_mlp",PaddedMLP(d,16,32).to(DEVICE))]
    results={}
    for name,model in models:
        opt=torch.optim.Adam(model.parameters(),lr=3e-3)
        t0=time.time()
        for epoch in range(35):
            order=rng.permutation(train_n)
            for ii in order:
                x=torch.tensor(tr[0][ii],device=DEVICE); A=torch.tensor(tr[1][ii],device=DEVICE)
                y=torch.tensor(tr[2][ii],device=DEVICE)
                opt.zero_grad(); logits=model(x,A).reshape(1,-1)
                loss=nn.functional.cross_entropy(logits,y.reshape(1)); loss.backward(); opt.step()
        correct=0; total=0; losses=[]
        with torch.no_grad():
            for x0,A0,y0 in zip(*te):
                x=torch.tensor(x0,device=DEVICE); A=torch.tensor(A0,device=DEVICE)
                logits=model(x,A); y=torch.tensor(y0,device=DEVICE)
                losses.append(float(nn.functional.cross_entropy(logits.reshape(1,-1),y.reshape(1))))
                correct += int(logits.argmax().item()==y0); total += 1
        results[name]={"unseen_schema_accuracy":correct/total,
                       "unseen_schema_ce":float(np.mean(losses)),
                       "parameters":sum(p.numel() for p in model.parameters()),
                       "seconds":time.time()-t0}
    # Activation storage comparison for a batch of heterogeneous schemas.
    results["padding_elements_per_instance"] = 16*32*d
    results["incidence_elements_per_instance_example"] = int(np.mean([a.size for a in te[1]]))
    return results

def main():
    result={"device":DEVICE,"exact_incidence":exact_incidence_check(),
            "prediction_1_exact_additivity": "l_p equals sum of incident s_m; numerical max error reported above",
            "prediction_2_normalization_scaling":normalization_sweep(),
            "prediction_3_equal-score_size_bias":softmax_bias_sweep(),
            "mini_experiment":mini_experiment()}
    with open("results.json","w") as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))
if __name__ == "__main__": main()