"""Standalone loader for the UCDP mini event encoder. Requires torch>=2.0."""
import json, torch, torch.nn as nn

class UcdpEncoder(nn.Module):
    def __init__(self, vsm, fields, d, n_heads, n_layers, dropout=0.1):
        super().__init__()
        self.fields = list(fields); self.embed_dim = d
        self.emb = nn.ModuleDict({f: nn.Embedding(vsm[f], d, padding_idx=0) for f in fields})
        self.field_proj = nn.Linear(len(fields) * d, d)
        self.proj_dropout = nn.Dropout(dropout)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, d))
        layer = nn.TransformerEncoderLayer(d_model=d, nhead=n_heads, dim_feedforward=d * 4,
                                           dropout=dropout, batch_first=True)
        self.transformer = nn.TransformerEncoder(layer, num_layers=n_layers,
                                                 enable_nested_tensor=False)
        self.ln = nn.LayerNorm(d)
        self.heads = nn.ModuleDict({f: nn.Linear(d, vsm[f] - 1) for f in fields})

    def encode_sequence(self, seq, mask):
        B = seq.size(0)
        e = torch.cat([self.emb[f](seq[:, :, i]) for i, f in enumerate(self.fields)], dim=-1)
        e = self.proj_dropout(self.field_proj(e))
        e = torch.cat([self.cls_token.expand(B, -1, -1), e], dim=1)
        cm = torch.ones(B, 1, dtype=torch.bool, device=mask.device)
        out = self.transformer(e, src_key_padding_mask=~torch.cat([cm, mask], dim=1))
        z = self.ln(out[:, 0, :])
        pad = (mask.sum(dim=1) == 0).unsqueeze(1)
        return torch.where(pad, torch.zeros_like(z), z)

def load_encoder(release_dir="."):
    cfg = json.load(open(f"{release_dir}/config.json"))
    vsm = {f: s + 1 for f, s in cfg["vocab_sizes_no_mask"].items()}
    m = UcdpEncoder(vsm, cfg["fields"], cfg["embed_dim"], cfg["n_heads"], cfg["n_layers"])
    m.load_state_dict(torch.load(f"{release_dir}/pytorch_model.pt", map_location="cpu",
                                 weights_only=True))
    m.eval()
    return m, cfg