"""ESM-2 and hierarchical-head scoring, run as a separate process.

Isolation is deliberate. PyTorch and LightGBM each bring their own OpenMP runtime, and loading both
into one interpreter aborts the process on some platforms. A child process that never imports
LightGBM removes the conflict, and frees the model's memory when it exits.

Reads a JSON job from argv[1] and writes {query: {ec: [cls, l3, cen, knn]}} to argv[2]. Exits 0 with
whatever it finished when the deadline passes, so a slow machine degrades rather than fails.
"""
import json
import sys
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F

sys.path.insert(0, str(Path(__file__).parent))
from embmodel import evidence, load_esm, load_head  # noqa: E402


def main():
    job = json.loads(Path(sys.argv[1]).read_text())
    out_path = Path(sys.argv[2])
    cfg, weights = job["cfg"], Path(job["weights"])
    seqs, order, cands = job["seqs"], job["order"], job["cands"]
    budget = job["budget"]
    start = time.time()

    torch.set_num_threads(job.get("threads", 2))
    model, alphabet = load_esm(weights / cfg["esm_file"])
    net, ecs, l3 = load_head(weights / cfg["head_file"])
    C = torch.tensor(np.load(weights / cfg["centroids_file"]))
    Z = F.normalize(torch.tensor(np.load(weights / cfg["z_file"]).astype(np.float32)), dim=-1)
    ref_ecs = json.loads((weights / cfg["z_ecs_file"]).read_text())

    bc = alphabet.get_batch_converter()
    layer = model.num_layers
    max_len, max_tok = cfg["max_len"], cfg["max_tokens"]

    out = {}
    p = 0
    with torch.inference_mode():
        while p < len(order) and time.time() - start < budget:
            q = p + 1
            top = min(len(seqs[str(order[p])]), max_len)
            while q < len(order):
                top2 = max(top, min(len(seqs[str(order[q])]), max_len))
                if (q - p + 1) * top2 > max_tok:
                    break
                top = top2
                q += 1
            batch = order[p:q]
            _, _, toks = bc([(str(i), seqs[str(i)][:max_len]) for i in batch])
            rep = model(toks, repr_layers=[layer])["representations"][layer]
            mask = (toks != alphabet.padding_idx) & (toks != alphabet.cls_idx) & (toks != alphabet.eos_idx)
            x = (rep * mask.unsqueeze(-1)).sum(1) / mask.sum(1, keepdim=True)
            ev = evidence(net, x, C, ecs, l3, None, Z, ref_ecs, [cands[str(i)] for i in batch])
            for i, d in zip(batch, ev):
                out[str(i)] = d
            p = q
    out_path.write_text(json.dumps(out))
    print(f"embedded {len(out)}/{len(order)} queries in {time.time() - start:.1f}s", file=sys.stderr)


if __name__ == "__main__":
    main()