"""Convert the released MEG-XL Lightning checkpoint to braindecode's MEGXL.

usage: python convert.py meg-xl-med-v2.ckpt <out_dir> <fif with the default chs_info>

The checkpoint holds the frozen BioCodec tokenizer, so ``biocodec_ckpt.pt`` is
not needed (its tensors equal the copy inside the checkpoint, max-abs 0.0).
The keys are renamed to braindecode's (``RENAME``). Dropped: the BioCodec
decoder, the duplicated ``biocodec_embedding_*`` codebooks, ``mask_token``,
``output_head`` (pretraining only) and the complex RoPE cache. The port's
RoPE table (``rotate``, as many positions as the model's buffer: 625, the
150 s pretraining window) is recomputed from ``freqs`` with ``torch.polar``
on the CPU, as the authors' loader recomputes it, and stored as ``(cos, sin)``
pairs.

``<out_dir>`` gets ``config.json`` from ``save_pretrained`` (the 306 MEG
channels of the given FIF file as the default ``chs_info``, 150 s at 50 Hz)
and ``model.safetensors`` / ``pytorch_model.bin`` without the classification
head: its shape depends on the number of channels and it is not pretrained,
so ``from_pretrained`` builds a new one for the caller's ``chs_info``.
"""

import sys

import torch

DROP = ("tokenizer.decoder.", "biocodec_embedding_", "mask_token", "output_head.")
RENAME = (
    ("tokenizer.encoder.model.", "encoder."),
    (".conv.conv.", ".conv."),
    ("tokenizer.quantizer.vq.", "quantizer."),
    ("criss_cross_transformer.layers.", "blocks."),
    (".ff.layer.0.", ".ff.0."),
    (".ff.layer.2.", ".ff.3."),
    ("position_fourier_emb.B", "position_freqs"),
    ("orientation_fourier_emb.B", "orientation_freqs"),
)


def convert(state_dict, model):
    """Map ``state_dict`` onto ``model``'s keys."""
    out = {}
    for key, value in state_dict.items():
        if key.startswith(DROP) or key.endswith("rope_embedding_layer.rotate"):
            continue
        for old, new in RENAME:
            key = key.replace(old, new)
        out[key] = value
    for key, value in model.state_dict().items():
        if key.endswith("rope_embedding_layer.rotate"):
            freqs = out[key.replace("rotate", "freqs")]
            angles = torch.outer(torch.arange(value.shape[0]).type_as(freqs), freqs)
            out[key] = torch.view_as_real(torch.polar(torch.ones_like(angles), angles))
    return out


if __name__ == "__main__":
    import mne
    from safetensors.torch import save_file

    from braindecode.models import MEGXL

    ckpt_path, out_dir, fif = sys.argv[1:4]
    raw = mne.io.read_raw_fif(fif, verbose=False)
    raw.pick(mne.pick_types(raw.info, meg=True, ref_meg=False, exclude=[]))
    torch.manual_seed(0)
    model = MEGXL(chs_info=raw.info["chs"], n_outputs=2, n_times=7500, sfreq=50.0)
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    state = convert(ckpt["state_dict"], model)
    missing, unexpected = model.load_state_dict(state, strict=False)
    assert unexpected == [] and set(missing) == {"final_layer.weight", "final_layer.bias"}
    model.save_pretrained(out_dir)
    weights = {
        k: v.contiguous()
        for k, v in model.state_dict().items()
        if not k.startswith("final_layer.")
    }
    save_file(weights, f"{out_dir}/model.safetensors", metadata={"format": "pt"})
    torch.save(weights, f"{out_dir}/pytorch_model.bin")