"""FastPDN-PII FP32 ONNX -> model.safetensors for BertForTokenClassification (plan Task 4a Step 2, Review Gate 14 A).

Named initializers keep their torch names. Each linear weight is the MatMul input that feeds the Add or BiasGelu taking
its named bias, transposed from ONNX's [in, out] to torch's [out, in]. Every weight is used once, the state dict loads
strictly, and no unused float tensor is left in the graph; anything else stops. Weights are copied, never changed.
Usage: python convert.py    # writes model.safetensors and convert.json here; create-only
"""
from pathlib import Path
import hashlib, json

import numpy as np
from onnx import numpy_helper

HERE = Path(__file__).resolve().parent
REPO, REVISION = 'ArkadiuszPawlak/fastpdn-ner-polish-pii', '636c57fe77b0afad8f065e7c76243dc051437035'
SNAPSHOT = Path.home()/'.cache/huggingface/hub/models--ArkadiuszPawlak--fastpdn-ner-polish-pii/snapshots'/REVISION
ONNX_SHA256 = '11924bfd20eb478ee614e5e1184536825db4747753d341af7c04b940efa84645'  # model.onnx, LFS pin at REVISION
CONFIG_SHA256 = '3cef51cd7d184478adbdec9f68730db279386fe145cb32473a0622edea923578'
NAMED = ('bert.', 'classifier.')
BUFFERS = ('bert.embeddings.position_ids',)  # non-persistent in transformers: checked, not saved


def sha(path):
    h = hashlib.sha256()
    with open(path, 'rb') as f:
        for block in iter(lambda: f.read(1 << 20), b''):
            h.update(block)
    return h.hexdigest()


def weights(graph):
    """({torch name: array}, used initializer names)."""
    init = {i.name: numpy_helper.to_array(i) for i in graph.initializer}
    producer = {o: n for n in graph.node for o in n.output}
    state = {k: v for k, v in init.items() if k.startswith(NAMED)}
    used = set(state)
    for n in graph.node:
        bias = [x for x in n.input if x.endswith('.bias') and 'LayerNorm' not in x and x in init]
        if n.op_type not in ('Add', 'BiasGelu') or not bias:
            continue
        (b,), (other,) = bias, [x for x in n.input if x not in bias]
        matmul = producer.get(other)
        assert matmul is not None and matmul.op_type == 'MatMul' and matmul.input[1] in init, b
        w, key = matmul.input[1], b[:-len('bias')] + 'weight'
        assert w not in used and key not in state, (w, key)
        state[key] = init[w].T.copy()
        used.add(w)
    return state, used


def leftovers(init, used):
    """Unused float tensors with more than one value: a weight the mapping missed. Scalars and int shapes pass."""
    return sorted(k for k, v in init.items() if k not in used and np.asarray(v).dtype.kind == 'f' and np.size(v) > 1)


def main():
    import onnx, torch
    from safetensors.torch import save_file
    from transformers import BertConfig, BertForTokenClassification
    out = HERE/'model.safetensors'
    assert not out.exists(), 'create-only'
    assert sha(SNAPSHOT/'model.onnx') == ONNX_SHA256 and sha(SNAPSHOT/'config.json') == CONFIG_SHA256
    graph = onnx.load(SNAPSHOT/'model.onnx').graph
    init = {i.name: numpy_helper.to_array(i) for i in graph.initializer}
    state, used = weights(graph)
    left = leftovers(init, used)
    assert not left, f'{len(left)} unused float tensors'
    for k in BUFFERS:
        v = state.pop(k)
        assert np.array_equal(v, np.arange(v.shape[-1]).reshape(v.shape)), k
    tensors = {k: torch.from_numpy(np.ascontiguousarray(v)) for k, v in sorted(state.items())}
    model = BertForTokenClassification(BertConfig.from_pretrained(SNAPSHOT))
    expected = model.state_dict()
    assert {k: tuple(v.shape) for k, v in tensors.items()} == {k: tuple(v.shape) for k, v in expected.items()}, 'keys/shapes'
    assert all(t.dtype == torch.float32 for t in tensors.values())
    model.load_state_dict(tensors, strict=True)
    save_file(tensors, out, metadata=dict(source=f'{REPO}@{REVISION}/model.onnx', onnx_sha256=ONNX_SHA256,
                                          note='converted from ONNX, weights unchanged'))
    report = dict(repo=REPO, revision=REVISION, onnx_sha256=ONNX_SHA256, config_sha256=CONFIG_SHA256,
                  initializers=len(init), used=len(used), tensors=len(tensors),
                  transposed=len(used) - sum(k.startswith(NAMED) for k in used),
                  unused_constants=len(init) - len(used), safetensors_sha256=sha(out), convert_sha256=sha(__file__))
    (HERE/'convert.json').write_text(json.dumps(report, indent=1) + '\n')
    return report


if __name__ == '__main__':
    print(json.dumps(main()))