"""Convert the released MEG-GPT + EphysTokenizer weights into one braindecode MEGGPT state dict.

Inputs (OHBA-analysis/MEG-GPT @78ff9439, models/):
  meg-gpt/model_state.pt        MEGGPT (meg_gpt/models/meg_gpt.py) state dict
  tokenizer/model_state.pt      EphysTokenizer state dict
  tokenizer/vocab.pkl           vocabulary ({"label_map": codebook entry -> token id, ...})

Dropped on purpose (documented in the PR): the Cam-CAN session embedding
(input_embedding_layer.extra_embeds.0.*), the next-token prediction head, the
attention-mask buffers (rebuilt from the config, checked equal in parity.py),
the tokenizer decoder (token_basis_layer, only used to reconstruct signals) and
the annealing temperature buffer.

usage: python convert.py <models dir> <out dir>

<models dir> is ``models/`` of the original repository; <out dir> receives the
braindecode export (``MEGGPT.save_pretrained``: config.json, model.safetensors,
pytorch_model.bin) that ``MEGGPT.from_pretrained(<out dir or Hub repo>)`` loads.
The classification head (n_outputs=2) is not part of the release: it is
randomly initialised (seed 0); pass ``n_outputs`` to ``from_pretrained`` to
replace it.
"""

import pickle
import re
import sys
from pathlib import Path

import torch

_GPT_RULES = [
    (r"input_embedding_layer\.token_embed\.base_module\.weight", "token_embedding.weight"),
    (r"input_embedding_layer\.pos_embed\.base_module\.position_embeddings", "time_embedding"),
    (r"input_embedding_layer\.channel_embed\.base_module\.position_embeddings", "channel_embedding"),
    (r"transformer_decoder\.layers\.(\d+)\.norm_time1\.norm_layer\.", r"layers.\1.norm1."),
    (r"transformer_decoder\.layers\.(\d+)\.norm_time2\.norm_layer\.", r"layers.\1.norm2."),
    (r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.patch_projection_k\.", r"layers.\1.attn.patch_k."),
    (r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.patch_projection_v\.", r"layers.\1.attn.patch_v."),
    (r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.patched_projection_k\.", r"layers.\1.attn.k_proj."),
    (r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.patched_projection_v\.", r"layers.\1.attn.v_proj."),
    (r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.query_projection\.", r"layers.\1.attn.q_proj."),
    (r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.output_projection\.", r"layers.\1.attn.out_proj."),
    (r"transformer_decoder\.layers\.(\d+)\.ff_time\.net\.0\.", r"layers.\1.mlp.0."),
    (r"transformer_decoder\.layers\.(\d+)\.ff_time\.net\.2\.", r"layers.\1.mlp.3."),
]
_GPT_DROP = re.compile(
    r"input_embedding_layer\.extra_embeds\.|prediction_head\.|.*\.gasp_layer\.attention_mask$"
)
_TOK_RULES = [
    (r"encoder_layer\.rnn_layers\.0\.", "tokenizer.rnn."),
    (r"token_weights_layer\.dense_layer\.", "tokenizer.proj."),
    (r"token_weights_layer\.norm_layer\.", "tokenizer.norm."),
]
_TOK_DROP = re.compile(r"decoder_layer\.token_basis_layer\.|token_weights_layer\._temperature$")


def _rename(state, rules, drop):
    out, dropped = {}, []
    for key, value in state.items():
        if drop.match(key):
            dropped.append(key)
            continue
        for pattern, repl in rules:
            new, n = re.subn("^" + pattern, repl, key)
            if n:
                out[new] = value
                break
        else:
            raise KeyError(f"no conversion rule for {key}")
    return out, dropped


def convert(gpt_state, tok_state, vocab):
    """Return (braindecode state dict, list of dropped original keys)."""
    gpt, dropped_gpt = _rename(gpt_state, _GPT_RULES, _GPT_DROP)
    tok, dropped_tok = _rename(tok_state, _TOK_RULES, _TOK_DROP)
    state = {**tok, **gpt}
    state["tokenizer.vocab"] = torch.as_tensor(vocab["label_map"], dtype=torch.long)
    return state, dropped_gpt + dropped_tok


def load_released(models_dir):
    models_dir = Path(models_dir)
    gpt = torch.load(models_dir / "meg-gpt/model_state.pt", map_location="cpu", weights_only=True)
    tok = torch.load(models_dir / "tokenizer/model_state.pt", map_location="cpu", weights_only=True)
    with open(models_dir / "tokenizer/vocab.pkl", "rb") as f:
        vocab = pickle.load(f)
    return gpt, tok, vocab


def export(models_dir, out_dir):
    from braindecode.models import MEGGPT

    state, dropped = convert(*load_released(models_dir))
    torch.manual_seed(0)
    model = MEGGPT(n_chans=52, n_outputs=2, sfreq=250).eval()
    res = model.load_state_dict(state, strict=False)
    assert sorted(res.missing_keys) == ["final_layer.1.bias", "final_layer.1.weight"], res
    assert not res.unexpected_keys, res
    model.save_pretrained(out_dir)
    return state, dropped


if __name__ == "__main__":
    state, dropped = export(sys.argv[1], sys.argv[2])
    print(f"exported {len(state)} converted tensors to {sys.argv[2]}; dropped {len(dropped)}: {dropped}")