"""Convert the released BBL Lightning checkpoint (CamCAN Pretrainer of
neural-processing-lab/the-brains-bitter-lesson) to braindecode's ``BBL``.

usage: python convert.py <bbl.ckpt> <out_dir>

Keeps the backbone only (the SSL projector and the band/amp/phase pretext heads
are training-only), renames the keys to braindecode's, and fills the
``trained_subjects`` buffer with the subject rows the optimizer updated
(non-zero Adam second moment), since the original keeps that set in a Python
attribute that is not saved.

``<out_dir>`` gets ``config.json``, ``model.safetensors`` and
``pytorch_model.bin`` from ``BBL.save_pretrained`` for
``BBL(n_chans=306, n_outputs=2, n_times=125, sfreq=250.0)``. The linear head
(``final_layer``) is not part of the release: it is PyTorch's default init
under ``torch.manual_seed(0)``; ``from_pretrained(..., n_outputs=k)`` replaces it.
"""

import sys

import torch

RENAMES = [
    ("backbone.dataset_gating.projector.", "projector."),
    ("backbone.encoder.model.", "encoder."),
    ("backbone.subject_embedding.", "subject_embedding."),
    (".conv.conv.weight_g", ".conv.parametrizations.weight.original0"),
    (".conv.conv.weight_v", ".conv.parametrizations.weight.original1"),
    (".conv.conv.bias", ".conv.bias"),
]


def convert_state_dict(state_dict):
    out = {}
    for key, value in state_dict.items():
        if not key.startswith("backbone."):
            continue
        for old, new in RENAMES:
            key = key.replace(old, new)
        out[key] = value
    return out


def trained_subject_mask(checkpoint):
    # Adam state 0 is the subject embedding, the first parameter of the Pretrainer.
    state = checkpoint["optimizer_states"][0]["state"][0]["exp_avg_sq"]
    assert state.shape == checkpoint["state_dict"]["backbone.subject_embedding.weight"].shape
    return (state.abs().sum(dim=1) > 0).float()


def convert(checkpoint):
    out = convert_state_dict(checkpoint["state_dict"])
    out["trained_subjects"] = trained_subject_mask(checkpoint)
    return out


if __name__ == "__main__":
    from braindecode.models import BBL

    ckpt_path, out_dir = sys.argv[1:3]
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    state = convert(ckpt)
    torch.manual_seed(0)
    model = BBL(n_chans=306, n_outputs=2, n_times=125, sfreq=250.0).eval()
    result = model.load_state_dict(state, strict=False)
    assert sorted(result.missing_keys) == ["final_layer.bias", "final_layer.weight"], result
    assert not result.unexpected_keys, result
    model.save_pretrained(out_dir)
    print("hyper_parameters", ckpt["hyper_parameters"])
    print("n_keys", len(state), "trained subjects", int(state["trained_subjects"].sum()))