"""Hierarchical CLEAN-style head (shared by training and inference) and offline ESM-2 loader.

v1 used a flat softmax. HIT-EC's central idea is that an EC number is a path, not a label, so this
head predicts all four levels from one trunk. That keeps level-3 evidence usable when the exact
fourth digit is rare, which is precisely the remote-homology case the benchmark tests.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F


def prefix(ec, level):
    return ".".join(ec.split(".")[:level])


class HierNet(nn.Module):
    def __init__(self, d_in, sizes, d_hid=1024, d_proj=256, drop=0.3):
        """sizes = [number of L1, L2, L3, L4 classes]."""
        super().__init__()
        self.body = nn.Sequential(nn.LayerNorm(d_in), nn.Linear(d_in, d_hid), nn.GELU(), nn.Dropout(drop),
                                  nn.Linear(d_hid, d_hid), nn.GELU(), nn.Dropout(drop))
        self.proj = nn.Linear(d_hid, d_proj)
        self.cls = nn.ModuleList(nn.Linear(d_hid, n) for n in sizes)

    def forward(self, x):
        h = self.body(x)
        return F.normalize(self.proj(h), dim=-1), [c(h) for c in self.cls]


def load_head(path):
    ck = torch.load(path, map_location="cpu", weights_only=False)
    net = HierNet(ck["d_in"], ck["sizes"])
    net.load_state_dict(ck["state"])
    return net.eval(), ck["ecs"], ck["l3"]


@torch.no_grad()
def evidence(net, x, C, ecs, l3, l3_of, Zdb, db_ecs, cands, topk=8, knn=20, cen_temp=0.05, batch=2048):
    """Per-query {ec: (cls, l3, cen, knn)} for the head's top ECs plus the given candidates.

    `cands[i]` are the ECs homology already proposed: they are always scored exactly, where v1
    could only score an EC that also happened to reach the head's own top 15.
    """
    index = {e: j for j, e in enumerate(ecs)}
    out = []
    for s in range(0, len(x), batch):
        z, logits = net(x[s:s + batch])
        p4 = torch.softmax(logits[3], -1)
        p3 = torch.softmax(logits[2], -1)
        pc = torch.softmax(z @ C.T / cen_temp, -1)
        sim = z @ Zdb.T
        near = torch.topk(sim, knn, -1)
        top = torch.topk(p4 + pc, topk, -1).indices.tolist()
        for r in range(len(z)):
            nb = {}
            for v, j in zip(near.values[r].tolist(), near.indices[r].tolist()):
                for e in db_ecs[j]:
                    nb[e] = nb.get(e, 0.0) + max(v, 0.0) ** 8 / len(db_ecs[j])
            tot = sum(nb.values()) or 1.0
            d = {}
            for e in set(cands[s + r]) | {ecs[j] for j in top[r]} | set(sorted(nb, key=nb.get)[-3:]):
                j = index.get(e)
                k = l3.get(prefix(e, 3))
                d[e] = (float(p4[r, j]) if j is not None else 0.0,
                        float(p3[r, k]) if k is not None else 0.0,
                        float(pc[r, j]) if j is not None else 0.0,
                        nb.get(e, 0.0) / tot)
            out.append(d)
    return out


def load_esm(path):
    """Load an ESM-2 checkpoint from a local file (no network, contact head not needed)."""
    import warnings

    import esm
    data = torch.load(str(path), map_location="cpu", weights_only=False)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        model, alphabet = esm.pretrained.load_model_and_alphabet_core(path.stem, data, None)
    return model.eval(), alphabet