"""PeSTo geometric transformer, single-channel BindCraft-targetability head.

Vendored verbatim from the PeSTo repo (`src/model_operations.py` + `model/model.py`)
so the Space has no dependency on the training tree.  Only the pieces the forward
pass touches are kept; the cross-state layers used by other PeSTo heads are not.

One deliberate change: `StateUpdateLayer.forward` skips gradient checkpointing when
grad is disabled.  Checkpointing exists to trade compute for memory during training;
under `torch.no_grad()` it only costs time, and the arithmetic is identical.

`config_model` is a copy of `model/config_bindcraft.py::config_model`, i.e. the
architecture `weights.pt` was trained with.  It must match the checkpoint --
`load_state_dict` is called strict, so a mismatch raises rather than silently
producing garbage scores.
"""

import torch as pt
from torch.utils.checkpoint import checkpoint


# architecture of the shipped checkpoint -- must match weights.pt exactly
config_model = {
    "em": {'N0': 30, 'N1': 32},
    "sum": (
        [{'Ns': 32, 'Nh': 2, 'Nk': 3, 'nn': 8}] * 8
        + [{'Ns': 32, 'Nh': 2, 'Nk': 3, 'nn': 16}] * 8
        + [{'Ns': 32, 'Nh': 2, 'Nk': 3, 'nn': 32}] * 8
        + [{'Ns': 32, 'Nh': 2, 'Nk': 3, 'nn': 64}] * 8
    ),
    "spl": {'N0': 32, 'N1': 32, 'Nh': 4},
    "dm": {'N0': 32, 'N1': 32, 'N2': 1},  # single targetability score
}


# >> UTILS
def unpack_state_features(X, ids_topk, q):
    # compute displacement vectors
    R_nn = X[ids_topk - 1] - X.unsqueeze(1)
    # compute distance matrix
    D_nn = pt.norm(R_nn, dim=2)
    # mask distances
    D_nn = D_nn + pt.max(D_nn) * (D_nn < 1e-2).float()
    # normalize displacement vectors
    R_nn = R_nn / D_nn.unsqueeze(2)

    # prepare sink
    q = pt.cat([pt.zeros((1, q.shape[1]), device=q.device), q], dim=0)
    ids_topk = pt.cat([pt.zeros((1, ids_topk.shape[1]), dtype=pt.long, device=ids_topk.device), ids_topk], dim=0)
    D_nn = pt.cat([pt.zeros((1, D_nn.shape[1]), device=D_nn.device), D_nn], dim=0)
    R_nn = pt.cat([pt.zeros((1, R_nn.shape[1], R_nn.shape[2]), device=R_nn.device), R_nn], dim=0)

    return q, ids_topk, D_nn, R_nn


# >>> OPERATIONS
class StateUpdate(pt.nn.Module):
    def __init__(self, Ns, Nh, Nk):
        super(StateUpdate, self).__init__()
        # operation parameters
        self.Ns = Ns
        self.Nh = Nh
        self.Nk = Nk

        # node query model
        self.nqm = pt.nn.Sequential(
            pt.nn.Linear(2 * Ns, Ns),
            pt.nn.ELU(),
            pt.nn.Linear(Ns, Ns),
            pt.nn.ELU(),
            pt.nn.Linear(Ns, 2 * Nk * Nh),
        )

        # edges scalar keys model
        self.eqkm = pt.nn.Sequential(
            pt.nn.Linear(6 * Ns + 1, Ns),
            pt.nn.ELU(),
            pt.nn.Linear(Ns, Ns),
            pt.nn.ELU(),
            pt.nn.Linear(Ns, Nk),
        )

        # edges vector keys model
        self.epkm = pt.nn.Sequential(
            pt.nn.Linear(6 * Ns + 1, Ns),
            pt.nn.ELU(),
            pt.nn.Linear(Ns, Ns),
            pt.nn.ELU(),
            pt.nn.Linear(Ns, 3 * Nk),
        )

        # edges value model
        self.evm = pt.nn.Sequential(
            pt.nn.Linear(6 * Ns + 1, 2 * Ns),
            pt.nn.ELU(),
            pt.nn.Linear(2 * Ns, 2 * Ns),
            pt.nn.ELU(),
            pt.nn.Linear(2 * Ns, 2 * Ns),
        )

        # scalar projection model
        self.qpm = pt.nn.Sequential(
            pt.nn.Linear(Nh * Ns, Ns),
            pt.nn.ELU(),
            pt.nn.Linear(Ns, Ns),
            pt.nn.ELU(),
            pt.nn.Linear(Ns, Ns),
        )

        # vector projection model
        self.ppm = pt.nn.Sequential(
            pt.nn.Linear(Nh * Ns, Ns, bias=False),
        )

        # scaling factor for attention
        self.sdk = pt.nn.Parameter(pt.sqrt(pt.tensor(Nk).float()), requires_grad=False)

    def forward(self, q, p, q_nn, p_nn, d_nn, r_nn):
        # q: [N, S]        p: [N, 3, S]
        # q_nn: [N, n, S]  p_nn: [N, n, 3, S]
        # d_nn: [N, n]     r_nn: [N, n, 3]

        # get dimensions
        N, n, S = q_nn.shape

        # node inputs packing
        X_n = pt.cat([
            q,
            pt.norm(p, dim=1),
        ], dim=1)  # [N, 2*S]

        # edge inputs packing
        X_e = pt.cat([
            d_nn.unsqueeze(2),                                  # distance
            X_n.unsqueeze(1).repeat(1, n, 1),                   # centered state
            q_nn,                                               # neighbors states
            pt.norm(p_nn, dim=2),                               # neighbors vector states norms
            pt.sum(p.unsqueeze(1) * r_nn.unsqueeze(3), dim=2),  # centered vector state projections
            pt.sum(p_nn * r_nn.unsqueeze(3), dim=2),            # neighbors vector states projections
        ], dim=2)  # [N, n, 6*S+1]

        # node queries
        Q = self.nqm.forward(X_n).view(N, 2, self.Nh, self.Nk)  # [N, 2*S] -> [N, 2, Nh, Nk]

        # scalar edges keys while keeping interaction order inveriance
        Kq = self.eqkm.forward(X_e).view(N, n, self.Nk).transpose(1, 2)  # [N, n, 6*S+1] -> [N, Nk, n]

        # vector edges keys while keeping bond order inveriance
        Kp = pt.cat(pt.split(self.epkm.forward(X_e), self.Nk, dim=2), dim=1).transpose(1, 2)

        # edges values while keeping interaction order inveriance
        V = self.evm.forward(X_e).view(N, n, 2, S).transpose(1, 2)  # [N, n, 6*S+1] -> [N, 2, n, S]

        # vectorial inputs packing
        Vp = pt.cat([
            V[:, 1].unsqueeze(2) * r_nn.unsqueeze(3),
            p.unsqueeze(1).repeat(1, n, 1, 1),
            p_nn,
        ], dim=1).transpose(1, 2)  # [N, 3, 3*n, S]

        # queries and keys collapse
        Mq = pt.nn.functional.softmax(pt.matmul(Q[:, 0], Kq) / self.sdk, dim=2)  # [N, Nh, n]
        Mp = pt.nn.functional.softmax(pt.matmul(Q[:, 1], Kp) / self.sdk, dim=2)  # [N, Nh, 3*n]

        # scalar state attention mask and values collapse
        Zq = pt.matmul(Mq, V[:, 0]).view(N, self.Nh * self.Ns)  # [N, Nh*S]
        Zp = pt.matmul(Mp.unsqueeze(1), Vp).view(N, 3, self.Nh * self.Ns)  # [N, 3, Nh*S]

        # decode outputs
        qh = self.qpm.forward(Zq)
        ph = self.ppm.forward(Zp)

        # update state with residual
        qz = q + qh
        pz = p + ph

        return qz, pz


class StatePoolLayer(pt.nn.Module):
    def __init__(self, N0, N1, Nh):
        super(StatePoolLayer, self).__init__()
        # state attention model
        self.sam = pt.nn.Sequential(
            pt.nn.Linear(2 * N0, N0),
            pt.nn.ELU(),
            pt.nn.Linear(N0, N0),
            pt.nn.ELU(),
            pt.nn.Linear(N0, 2 * Nh),
        )

        # attention heads decoding
        self.zdm = pt.nn.Sequential(
            pt.nn.Linear(Nh * N0, N0),
            pt.nn.ELU(),
            pt.nn.Linear(N0, N0),
            pt.nn.ELU(),
            pt.nn.Linear(N0, N1),
        )

        # vector attention heads decoding
        self.zdm_vec = pt.nn.Sequential(
            pt.nn.Linear(Nh * N0, N1, bias=False)
        )

    def forward(self, q, p, M):
        # create filter for softmax
        F = (1.0 - M + 1e-6) / (M - 1e-6)

        # pack features
        z = pt.cat([q, pt.norm(p, dim=1)], dim=1)

        # multiple attention pool on state
        Ms = pt.nn.functional.softmax(self.sam.forward(z).unsqueeze(1) + F.unsqueeze(2), dim=0).view(M.shape[0], M.shape[1], -1, 2)
        qh = pt.matmul(pt.transpose(q, 0, 1), pt.transpose(Ms[:, :, :, 0], 0, 1))
        ph = pt.matmul(pt.transpose(pt.transpose(p, 0, 2), 0, 1), pt.transpose(Ms[:, :, :, 1], 0, 1).unsqueeze(1))

        # attention heads decoding
        qr = self.zdm.forward(qh.view(Ms.shape[1], -1))
        pr = self.zdm_vec.forward(ph.view(Ms.shape[1], p.shape[1], -1))

        return qr, pr


# >>> LAYERS
class StateUpdateLayer(pt.nn.Module):
    def __init__(self, layer_params):
        super(StateUpdateLayer, self).__init__()
        # define operation
        self.su = StateUpdate(*[layer_params[k] for k in ['Ns', 'Nh', 'Nk']])
        # store number of nearest neighbors
        self.m_nn = pt.nn.Parameter(pt.arange(layer_params['nn'], dtype=pt.int64), requires_grad=False)

    def forward(self, Z):
        # unpack input
        q, p, ids_topk, D_topk, R_topk = Z

        # gather neighbor states
        ids_nn = ids_topk[:, self.m_nn]

        if pt.is_grad_enabled():
            # training path: recompute activations in the backward pass instead of storing them
            q = q.requires_grad_()
            p = p.requires_grad_()
            q, p = checkpoint(self.su.forward, q, p, q[ids_nn], p[ids_nn], D_topk[:, self.m_nn], R_topk[:, self.m_nn])
        else:
            # inference path: identical arithmetic, no checkpoint bookkeeping
            q, p = self.su.forward(q, p, q[ids_nn], p[ids_nn], D_topk[:, self.m_nn], R_topk[:, self.m_nn])

        # sink
        q[0] = q[0] * 0.0
        p[0] = p[0] * 0.0

        return q, p, ids_topk, D_topk, R_topk


class Model(pt.nn.Module):
    def __init__(self, config):
        super(Model, self).__init__()
        # features encoding models for structures and library
        self.em = pt.nn.Sequential(
            pt.nn.Linear(config['em']['N0'], config['em']['N1']),
            pt.nn.ELU(),
            pt.nn.Linear(config['em']['N1'], config['em']['N1']),
            pt.nn.ELU(),
            pt.nn.Linear(config['em']['N1'], config['em']['N1']),
        )
        # atomic level state update model
        self.sum = pt.nn.Sequential(*[StateUpdateLayer(layer_params) for layer_params in config['sum']])

        # atomic to residue reduction layer
        self.spl = StatePoolLayer(config['spl']['N0'], config['spl']['N1'], config['spl']['Nh'])

        # decoding mlp
        self.dm = pt.nn.Sequential(
            pt.nn.Linear(2 * config['dm']['N0'], config['dm']['N1']),
            pt.nn.ELU(),
            pt.nn.Linear(config['dm']['N1'], config['dm']['N1']),
            pt.nn.ELU(),
            pt.nn.Linear(config['dm']['N1'], config['dm']['N2']),
        )

    def forward(self, X, ids_topk, q0, M):
        # encode features
        q = self.em.forward(q0)

        # initial state vectors
        p0 = pt.zeros((q.shape[0] + 1, X.shape[1], q.shape[1]), device=X.device)

        # unpack state features with sink
        q, ids_topk, D_nn, R_nn = unpack_state_features(X, ids_topk, q)

        # atomic tsa layers
        qa, pa, _, _, _ = self.sum.forward((q, p0, ids_topk, D_nn, R_nn))

        # atomic to residue attention pool (without sink)
        qr, pr = self.spl.forward(qa[1:], pa[1:], M)

        # decode state
        zr = pt.cat([qr, pt.norm(pr, dim=1)], dim=1)
        z = self.dm.forward(zr)

        return z