# "Hopfield Pooling": multiple-instance learning on MNIST bags. A bag of N digits is positive if it contains exactly one
# "9" (a needle in a haystack), negative if it contains none. Same encoder, four ways to pool the bag:
#   mean, max, attention-based MIL (Ilse et al., 2018), HopfieldPooling (Ramsauer et al., 2021; one learned query),
#   and for the comparison post: gated attention-MIL (Ilse et al.), Set Transformer PMA (Lee et al., 2019),
#   log-sum-exp pooling with a learned temperature, generalized mean (GeM) with a learned exponent.
#   --nines k: positive bags contain k nines instead of one.
#   python mil_mnist.py --pool hopfield --bag 50 --seed 0 [--hflayers PATH]  ->  mil_<pool>_<bag>_<seed>.json (+ .pt)
#   --init mil_hopfield_10_0.pt starts from a model trained on smaller bags (a curriculum) -> mil_<pool>_<bag>_<seed>_init.json
# Needs MNIST in ./MNIST (torchvision layout) and the hopfield-layers package (pip install git+https://github.com/ml-jku/hopfield-layers).
import argparse, json, math, sys, time, numpy as np, torch, torch.nn as nn, torch.nn.functional as F, torchvision
ap = argparse.ArgumentParser(); ap.add_argument("--pool", default="hopfield"); ap.add_argument("--bag", type=int, default=50)
ap.add_argument("--seed", type=int, default=0); ap.add_argument("--epochs", type=int, default=15); ap.add_argument("--train_bags", type=int, default=500)
ap.add_argument("--hflayers", default=None); ap.add_argument("--init", default=None); ap.add_argument("--nines", type=int, default=1); A = ap.parse_args()
if A.hflayers: sys.path.insert(0, A.hflayers)
from hflayers import HopfieldPooling
torch.manual_seed(A.seed); rng = np.random.default_rng(A.seed)

def load(train):
    d = torchvision.datasets.MNIST(".", train=train, download=False)
    return (d.data.float() / 255.0 - 0.1307) / 0.3081, d.targets.numpy()
Xtr, Ytr = load(True); Xte, Yte = load(False)
def make_bags(Y, n_bags, N, rng, k=1):
    nines, others = np.where(Y == 9)[0], np.where(Y != 9)[0]; bags, labels, pos = [], [], []
    for b in range(n_bags):
        idx = rng.choice(others, N, replace=False); lab = b % 2
        p = -1
        if lab and k == 1: p = int(rng.integers(N)); idx[p] = rng.choice(nines)       # exactly one 9 (the random draws of post 4)
        elif lab: ps = rng.choice(N, k, replace=False); idx[ps] = rng.choice(nines, k, replace=False); p = int(ps[0])   # k nines
        bags.append(idx); labels.append(lab); pos.append(p)
    return np.array(bags), np.array(labels), np.array(pos)
test_seed = 1000 + A.bag if A.nines == 1 else 1000 + A.bag + 7 * A.nines
test_bags, test_y, test_pos = make_bags(Yte, 500, A.bag, np.random.default_rng(test_seed), A.nines)   # the same test bags for every run

class Net(nn.Module):
    def __init__(s, pool):
        super().__init__(); s.pool = pool; d = 64
        s.enc = nn.Sequential(nn.Conv2d(1, 20, 5), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(20, 50, 5), nn.ReLU(), nn.MaxPool2d(2),
                              nn.Flatten(), nn.Linear(800, d), nn.ReLU())
        if pool in ("attention", "gated"): s.V = nn.Linear(d, 32); s.w = nn.Linear(32, 1)
        if pool == "gated": s.U = nn.Linear(d, 32)
        if pool == "hopfield": s.hp = HopfieldPooling(input_size=d, hidden_size=32, quantity=1)
        if pool == "hopfield_nonorm":                # the same layer with its three input LayerNorms switched off
            off = {f"normalize_{k}{a}": False for k in ("stored_pattern", "state_pattern", "pattern_projection") for a in ("", "_affine")}
            s.hp = HopfieldPooling(input_size=d, hidden_size=32, quantity=1, **off)
        if pool == "pma":                            # Set Transformer: H = LN(S + MHA(S, X, X)), Z = LN(H + rFF(H)), one learned seed S
            s.seed = nn.Parameter(torch.randn(1, 1, d) * 0.02); s.mha = nn.MultiheadAttention(d, 4, batch_first=True)
            s.ln1, s.ln2 = nn.LayerNorm(d), nn.LayerNorm(d); s.ff = nn.Sequential(nn.Linear(d, d), nn.ReLU(), nn.Linear(d, d))
        if pool == "lse": s.log_r = nn.Parameter(torch.zeros(1))          # temperature r = exp(log_r), starts at 1
        if pool == "gem": s.p = nn.Parameter(torch.tensor(3.0))           # exponent p, starts at 3 as in Radenovic et al.
        s.out = nn.Linear(d, 1)
    def forward(s, x):                       # x: (B, N, 28, 28) -> logits (B,), weights (B, N) or None
        B, N = x.shape[:2]; h = s.enc(x.reshape(B * N, 1, 28, 28)).reshape(B, N, -1); w = None
        if s.pool == "mean": z = h.mean(1)
        elif s.pool == "max": z = h.max(1).values
        elif s.pool == "attention": w = torch.softmax(s.w(torch.tanh(s.V(h))).squeeze(-1), 1); z = (w.unsqueeze(-1) * h).sum(1)
        elif s.pool == "gated": w = torch.softmax(s.w(torch.tanh(s.V(h)) * torch.sigmoid(s.U(h))).squeeze(-1), 1); z = (w.unsqueeze(-1) * h).sum(1)
        elif s.pool == "lse": r = s.log_r.exp(); z = (torch.logsumexp(r * h, 1) - math.log(N)) / r
        elif s.pool == "gem": p = s.p.clamp(min=1.0); z = h.clamp(min=1e-6).pow(p).mean(1).pow(1.0 / p)
        elif s.pool == "pma":
            S = s.seed.expand(B, 1, -1); att, w = s.mha(S, h, h, need_weights=True, average_attn_weights=True)
            H = s.ln1(S + att); z = s.ln2(H + s.ff(H)).squeeze(1); w = w[:, 0]
        else: z = s.hp(h); w = s.hp.get_association_matrix(h)[:, 0, 0] if not s.training else None   # hopfield, hopfield_nonorm
        return s.out(z).squeeze(-1), w

net = Net(A.pool)
if A.init: net.load_state_dict(torch.load(A.init))          # the encoder and pooling do not depend on the bag size
opt = torch.optim.Adam(net.parameters(), lr=5e-4, weight_decay=1e-4); log = []
def evaluate():
    net.eval(); correct, wpos, rank1 = 0, [], 0
    with torch.no_grad():
        for i in range(0, len(test_y), 25):
            x = Xte[test_bags[i:i + 25]]; lg, w = net(x)
            correct += ((lg > 0).numpy() == test_y[i:i + 25]).sum()
            if w is not None:
                for j in range(len(lg)):
                    if test_y[i + j]:                     # total weight on the 9s of the bag; is the largest weight on a 9?
                        is9 = torch.tensor(Yte[test_bags[i + j]] == 9); wpos.append(float(w[j][is9].sum())); rank1 += int(bool(is9[w[j].argmax()]))
    npos = int(test_y.sum())
    return correct / len(test_y), (float(np.mean(wpos)) if wpos else None), (rank1 / npos if wpos else None)
t0 = time.time()
for ep in range(A.epochs):
    bags, ys, _ = make_bags(Ytr, A.train_bags, A.bag, rng, A.nines); net.train(); perm = rng.permutation(len(ys)); tot = 0
    bs = max(1, 400 // A.bag)
    for i in range(0, len(perm), bs):
        j = perm[i:i + bs]; lg, _ = net(Xtr[bags[j]]); loss = F.binary_cross_entropy_with_logits(lg, torch.tensor(ys[j], dtype=torch.float32))
        opt.zero_grad(); loss.backward(); opt.step(); tot += loss.item() * len(j)
    acc, wpos, r1 = evaluate(); log.append({"epoch": ep + 1, "loss": tot / len(ys), "test_acc": acc, "weight_on_9": wpos, "top1_is_9": r1})
    print(json.dumps(log[-1]), flush=True)
out = {"args": vars(A), "test_acc": log[-1]["test_acc"], "weight_on_9": log[-1]["weight_on_9"], "top1_is_9": log[-1]["top1_is_9"], "log": log, "seconds": round(time.time() - t0)}
tag = f"mil_{A.pool}_{A.bag}_{A.seed}" + ("_init" if A.init else "") + (f"_k{A.nines}" if A.nines != 1 else "")
json.dump(out, open(f"{tag}.json", "w")); torch.save(net.state_dict(), f"{tag}.pt")
print("final", out["test_acc"], out["weight_on_9"], out["top1_is_9"], out["seconds"], "s")