# Authors: Julien Gadonneix <juliengado.2001@gmail.com>
# License: BSD (3-clause)
"""Convert, check and publish the CORTEG checkpoints on CPU.

Run from an editable Braindecode checkout containing CORTEG, with the ``hub``
extra and ``timm`` and ``einops`` (for the reference) installed::

    git clone https://github.com/LiuyinYang1101/CORTEG && git -C CORTEG checkout caafb78
    curl -L -o checkpoint-300.pth \\
        https://github.com/LiuyinYang1101/STEEGFormer/releases/download/ST-EEGFormer-small/checkpoint-300.pth
    python convert_corteg_checkpoint.py --reference-dir CORTEG \\
        --backbone checkpoint-300.pth --output-dir /tmp/corteg-converted

``--reference-dir`` is a clone of the official repository at ``REVISION``, which
ships the Stanford adapter (``ADAPTER``); ``--backbone`` is the ST-EEGFormer
small release CORTEG starts from. The SHA-256 values below must match. The port
takes the backbone from its Braindecode conversion (``BACKBONE_REPO`` at
``BACKBONE_REVISION``). Two models are written (``--variant`` picks one):

- ``pretrained``: the backbone, with the high-gamma patch embedding warm-started
  from the low-frequency one by linear interpolation, as the authors'
  ``build_model`` does. The spatial adapter, gate and head are left out:
  Braindecode initializes them on load, the adapter from the electrode
  positions. The warm start must match the reference's within ``TOLERANCE``.
- ``stanford``: that model with the released adapter, which holds the trainable
  parameters of the authors' Stanford run (the spatial adapter, LoRA factors on
  blocks 4-7, the high-gamma patch embedding, those blocks' LayerNorms and the
  readout). The LoRA factors are folded into the fused qkv, proj, fc1 and fc2
  weights, so the model is a plain ``CORTEG``. The authors' ``load_corteg``
  builds the reference from the two release files; the outputs must agree
  within ``TOLERANCE``.

Download this script from either Braindecode CORTEG model repository on Hugging
Face; it is not bundled in the Braindecode source tree. Add ``--push-to
braindecode`` to upload each converted model, this script, the MIT licences, a
report and a model card in one commit.
"""

import argparse
import hashlib
import json
import shutil
import sys
import tempfile
from contextlib import contextmanager
from pathlib import Path
from types import SimpleNamespace

import numpy as np
import torch
from huggingface_hub import HfApi, hf_hub_download
from safetensors.torch import load_file
from torch import nn
from torch.nn import functional as F

from braindecode.models import CORTEG

REPOSITORY = "https://github.com/LiuyinYang1101/CORTEG"
REVISION = "caafb78fbebdd9448814f96d7e61083a3eb62663"
ADAPTER = "checkpoints/corteg_stanford_pooled.pt"
ADAPTER_SHA256 = "77d63a214e0e0f2c790e90878a0ef6b525f6e24e1453fcd94e118c63ff7e8396"
ADAPTER_PARAMS = 297_236  # the release manifest's trainable_params
BACKBONE_URL = (
    "https://github.com/LiuyinYang1101/STEEGFormer/releases/download/"
    "ST-EEGFormer-small/checkpoint-300.pth"
)
BACKBONE_SHA256 = "98d2848abdaea7a143d399c07f67f8782e07dd66e7c309ddb058c7158a8a8b02"
BACKBONE_REPO = "braindecode/STEEGFormer-small"
BACKBONE_REVISION = "633c9c0d76c6e26459e1e0d1337c91858578f13d"
REPOS = {"pretrained": "corteg-small-pretrained", "stanford": "corteg-stanford-pooled"}
# The paper's setting (the Stanford manifest's architecture_args): 1 s windows of
# 128 low-frequency and 200 high-gamma samples, the "average" fusion before the
# last four blocks. The Stanford head has one output per finger.
CONFIG = dict(
    n_times=128,
    sfreq=128.0,
    merge_strategy="average",
    n_merge_blocks=4,
    hi_patch_size=25,
    knn_k=8,
    drop_path=0.1,
)
N_OUTPUTS = {"pretrained": 2, "stanford": 5}
# Left out of the pretrained files: initialized on load (the adapter from chs_info).
NOT_PRETRAINED = ("channel_adapter", "final_layer")
LORA_SCALE = 16 / 4  # alpha / r
# Reference layer of a block -> its Braindecode modules; the fused qkv splits in three.
LAYERS = {
    "norm1": ["0.fn.0"],
    "norm2": ["1.fn.0"],
    "attn.qkv": ["0.fn.1.queries", "0.fn.1.keys", "0.fn.1.values"],
    "attn.proj": ["0.fn.1.projection"],
    "mlp.fc1": ["1.fn.1.0"],
    "mlp.fc2": ["1.fn.1.3"],
}
TOLERANCE = 1e-4
FINGERS = ["thumb", "index", "middle", "ring", "pinky"]
MIT = """MIT License

Copyright (c) {year} {holder}

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""
STEEGFORMER_LICENSE = (
    "The ST-EEGFormer small backbone (release ST-EEGFormer-small of\n"
    "https://github.com/LiuyinYang1101/STEEGFormer):\n\n"
    + MIT.format(year=2025, holder="LIUYIN YANG")
)
CORTEG_LICENSE = (
    f"The CORTEG adapter ({ADAPTER} of\n{REPOSITORY} at {REVISION}):\n\n"
    + MIT.format(year=2026, holder="Liuyin Yang")
)
LICENSES = {
    "pretrained": STEEGFORMER_LICENSE,
    "stanford": CORTEG_LICENSE + "\n\n" + STEEGFORMER_LICENSE,
}
CITATION = """```bibtex
@inproceedings{yang2026corteg,
  title     = {{CORTEG}: Foundation Models Enable Cross-Modality Representation Transfer
               from Scalp to Intracranial Brain Recordings},
  author    = {Yang, Liuyin and Sun, Qiang and Van Dyck, Bob and Calvo Merino, Eva and
               Van Hulle, Marc M.},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year      = {2026}
}
```"""


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


def montages():
    """Electrode positions in metres: the Braindecode test montage and 46 random contacts."""
    rng = np.random.default_rng(0)
    random = rng.uniform([-0.07, -0.10, -0.04], [0.07, 0.07, 0.08], size=(46, 3))
    return [
        torch.tensor(
            [[-0.04 + 0.004 * i, -0.02 + 0.003 * i, 0.06 - 0.002 * i] for i in range(8)]
        ),
        torch.tensor(random, dtype=torch.float32),
    ]


def load_backbone(model):
    """ST-EEGFormer small into ``model``, the high-gamma embedding warm-started.

    Returns the keys the backbone leaves to the adapter.
    """
    path = hf_hub_download(BACKBONE_REPO, "model.safetensors", revision=BACKBONE_REVISION)
    backbone = {k: v for k, v in load_file(path).items() if not k.startswith("final_layer.")}
    missing, unexpected = model.load_state_dict(backbone, strict=False)
    expected = {"channel_adapter", "final_layer", "patch_embed_hi"}
    if unexpected or {k.split(".")[0] for k in missing} != expected:
        raise SystemExit(f"backbone does not fit CORTEG: {missing=}, {unexpected=}")
    lo, hi = model.patch_embed.proj, model.patch_embed_hi.proj
    with torch.no_grad():
        hi.weight.copy_(F.interpolate(lo.weight[None], size=hi.in_features, mode="linear")[0])
        hi.bias.copy_(lo.bias)
    return missing


def fold_adapter(adapter, model, missing):
    """Load the released adapter into ``model``, the LoRA factors folded in."""
    state = model.state_dict()
    filled = set()
    for key, value in adapter.items():
        key = key.removeprefix("backbone.").replace("head.head.", "final_layer.")
        if not key.startswith("blocks."):
            if key not in state or state[key].shape != value.shape:
                raise SystemExit(f"adapter tensor {key} has no home in CORTEG")
            state[key] = value
            filled.add(key)
            continue
        _, block, rest = key.split(".", 2)
        layer, _, param = rest.rpartition(".")
        if layer.endswith(".B"):  # folded with its A factor
            continue
        if layer.endswith(".A"):
            layer = layer.removesuffix(".A")
            lora_b = adapter[f"backbone.blocks.{block}.{layer}.B.weight"]
            delta = (LORA_SCALE * lora_b @ value).chunk(len(LAYERS[layer]))
            for name, part in zip(LAYERS[layer], delta):
                state[f"encoder.{block}.{name}.weight"] += part
        else:
            target = f"encoder.{block}.{LAYERS[layer][0]}.{param}"
            state[target] = value
            filled.add(target)
    unfilled = {k for k in missing if k not in filled}
    if unfilled:
        raise SystemExit(f"the adapter leaves {sorted(unfilled)} at their initialization")
    model.load_state_dict(state, strict=True)


def build(variant, adapter, n_chans):
    model = CORTEG(n_chans=n_chans, n_outputs=N_OUTPUTS[variant], **CONFIG).eval()
    missing = load_backbone(model)
    if variant == "stanford":
        fold_adapter(adapter, model, missing)
    return model


@contextmanager
def reference(reference_dir, backbone):
    """The authors' code, with a backbone configuration pointing at ``backbone``."""
    sys.path.insert(0, str(reference_dir))
    config = json.loads((reference_dir / "configs/steegformer_small.json").read_text())
    config["pretrained"]["path"] = str(backbone.resolve())
    with tempfile.TemporaryDirectory() as tmp:
        path = Path(tmp) / "steegformer_small.json"
        path.write_text(json.dumps(config))
        yield path


def check_pretrained(model, reference_dir, backbone, pos):
    """Largest difference between the warm-started high-gamma embeddings."""
    manifest = json.loads((reference_dir / ADAPTER).with_suffix(".json").read_text())
    with reference(reference_dir, backbone) as config:
        from experiments.run_regression_hilo_clean import build_model

        overrides = dict(model_kwargs_json=str(config), no_pretrained=False)
        args = SimpleNamespace(**(manifest["build_args"] | overrides))
        ref = build_model(args, C_in=len(pos), T_in=128, ecog_xyz_m=1000 * pos.numpy())
    state = ref.state_dict()
    proj = model.patch_embed_hi.proj
    return max(
        (proj.weight - state["backbone.patch_embed_hi.proj.weight"]).abs().max().item(),
        (proj.bias - state["backbone.patch_embed_hi.proj.bias"]).abs().max().item(),
    )


def check_stanford(model, reference_dir, backbone, pos):
    """Largest output difference with the authors' ``load_corteg``."""
    n_chans = len(pos)
    g = torch.Generator().manual_seed(0)
    x = torch.randn(4, n_chans, 128, generator=g)
    x_hi = torch.randn(4, n_chans, 200, generator=g)
    with torch.no_grad():
        out = model(x, x_hi, pos=pos)
    with reference(reference_dir, backbone) as config:
        from load_corteg import load_corteg, predict

        ref = load_corteg(
            C_in=n_chans,
            T_in=128,
            ecog_xyz_mm=1000 * pos.numpy(),
            d_out=N_OUTPUTS["stanford"],
            backbone_config=str(config),
        )
    expected = torch.from_numpy(predict(ref, x.numpy(), x_hi.numpy()))
    return (out - expected).abs().max().item()


@contextmanager
def without(model, names):
    """Hide the ``names`` submodules so that their tensors are not serialized."""
    hidden = {name: getattr(model, name) for name in names}
    for name in names:
        setattr(model, name, nn.Identity())
    try:
        yield
    finally:
        for name, module in hidden.items():
            setattr(model, name, module)


def pretrained_card(report):
    return f"""---
library_name: braindecode
license: mit
pipeline_tag: feature-extraction
base_model: {BACKBONE_REPO}
tags:
- braindecode
- CORTEG
- ieeg
- ecog
- foundation-model
- pytorch_model_hub_mixin
---

# CORTEG, ST-EEGFormer small backbone

The starting point of CORTEG ([Yang et al., 2026](https://arxiv.org/abs/2605.10337)),
ST-EEGFormer small adapted to intracranial recordings, for
[`braindecode.models.CORTEG`](https://braindecode.org/stable/generated/braindecode.models.CORTEG.html):
the ST-EEGFormer small encoder ([`checkpoint-300.pth`]({BACKBONE_URL}),
sha256 `{BACKBONE_SHA256}`, through its conversion
[{BACKBONE_REPO}](https://huggingface.co/{BACKBONE_REPO}) at revision
`{BACKBONE_REVISION[:8]}`), {report["n_params"] / 1e6:.1f}M parameters. The
high-gamma patch embedding, which the scalp model does not have, is warm-started
from the low-frequency one, linearly interpolated from 16 to 25 samples, as the
authors' `build_model` does ([LiuyinYang1101/CORTEG]({REPOSITORY}/tree/{REVISION}),
`experiments/run_regression_hilo_clean.py`); it matches theirs to
{report["max_abs_diff"]:.1e} in float32 on CPU. `convert_corteg_checkpoint.py` in this
repository reproduces the conversion.

## Usage

```python
from braindecode.models import CORTEG

model = CORTEG.from_pretrained(
    "braindecode/{REPOS["pretrained"]}", n_outputs=5, n_times=128, chs_info=raw.info["chs"]
)
```

The spatial adapter, the head and, with `merge_strategy="layerwise_gate"`, the
gate are not in these files: Braindecode initializes them on load, the adapter
from the electrode positions of `chs_info` (metres). The paper then trains the
adapter, the high-gamma embedding, LoRA (r=4, alpha=16) on the last four blocks
with their LayerNorms, and the head, with the rest frozen. See
[`braindecode/{REPOS["stanford"]}`](https://huggingface.co/braindecode/{REPOS["stanford"]})
for the authors' trained Stanford model and the input streams.

## Limitations

These files hold no trained intracranial component, so the model needs
training before prediction. The check above covers the float32 CPU tensors of
the warm start.

## Citation

{CITATION}

## License

The ST-EEGFormer weights are released by Liuyin Yang under the MIT licence
(`LICENSE`). Braindecode's code is BSD-3-Clause.
"""


def stanford_card(report):
    return f"""---
library_name: braindecode
license: mit
base_model: braindecode/{REPOS["pretrained"]}
tags:
- braindecode
- CORTEG
- ieeg
- ecog
- finger-flexion
- pytorch_model_hub_mixin
---

# CORTEG, Stanford finger flexion (pooled)

The released CORTEG model of [Yang et al. (2026)](https://arxiv.org/abs/2605.10337),
ST-EEGFormer small adapted to ECoG and trained on the finger-flexion recordings
of the 9 Stanford patients pooled (Table 1, "CORTEG (pooled)", Pearson r 0.554),
{report["n_params"] / 1e6:.1f}M parameters, converted to
[`braindecode.models.CORTEG`](https://braindecode.org/stable/generated/braindecode.models.CORTEG.html).
It decodes the flexion of the five fingers ({", ".join(FINGERS)}) from 1 s of ECoG.

Sources:

- the adapter, `{ADAPTER}` of [LiuyinYang1101/CORTEG]({REPOSITORY}/tree/{REVISION}),
  sha256 `{ADAPTER_SHA256}`: the spatial adapter, LoRA factors (r=4, alpha=16) on
  blocks 4-7, the high-gamma patch embedding, those blocks' LayerNorms and the
  readout, {ADAPTER_PARAMS:,} parameters;
- the frozen backbone it was trained on, the ST-EEGFormer small release
  ([`checkpoint-300.pth`]({BACKBONE_URL}), sha256 `{BACKBONE_SHA256}`), here
  through [braindecode/{REPOS["pretrained"]}](https://huggingface.co/braindecode/{REPOS["pretrained"]}).

Conversion folds the LoRA factors into the fused qkv, proj, fc1 and fc2 weights
(scale alpha/r = 4), so the model is a plain `CORTEG`;
`convert_corteg_checkpoint.py` in this repository reproduces it. Outputs match
the authors' `load_corteg` on the two release files to {report["max_abs_diff"]:.1e}
in float32 on CPU, for montages of {" and ".join(str(n) for n in report["n_chans_checked"])}
contacts.

## Usage

The model reads two streams of the same 1 s window, as the paper's Stanford
pipeline (`data/stanford_preprocessing/` of the official repository): after a
notch at 60, 120 and 180 Hz and a common average reference, `x` is the signal
resampled to 128 Hz and `x_hi` its 70-200 Hz Hilbert envelope at 200 Hz, each
z-scored per channel with the statistics of the training windows. Electrode
positions are in metres, in `chs_info` or `pos`.

```python
import torch
from braindecode.models import CORTEG

model = CORTEG.from_pretrained(
    "braindecode/{REPOS["stanford"]}", chs_info=raw.info["chs"]
).eval()
flexion = model(x, x_hi)  # x: (batch, n_chans, 128), x_hi: (batch, n_chans, 200)
```

The output, `(batch, 5)`, is the flexion of each finger, z-scored per finger
with the training statistics of the patient, at the time 40 ms before the
window's end. The saved geometry is {report["n_chans"]} channels; the model
reads any montage, so pass your own `chs_info` (or `n_chans` and `pos` at each
call) when loading.

## Limitations

The readout is at its random initialization, as released: the paper run built
it after the optimizer, so it was never trained, and the adapter learned around
it (see the release manifest, `checkpoints/corteg_stanford_pooled.json`).
The model was trained on the training split of the 9 Stanford patients (the
first 400 s of each recording, two thirds of one under 600 s); other patients
or tasks need fine-tuning, as in the paper's leave-one-subject-out protocol.
The check above covers float32 CPU outputs, not decoding accuracy, GPU kernels
or mixed precision; the authors report that this checkpoint reproduces the
cohort mean of 0.554 (manifest, `paper.verified`).

## Citation

{CITATION}

## License

The adapter and the ST-EEGFormer backbone are both released by Liuyin Yang
under the MIT licence (`LICENSE`). Braindecode's code is BSD-3-Clause.
"""


CARDS = {"pretrained": pretrained_card, "stanford": stanford_card}


def convert(variant, adapter, args):
    """Convert, check and write one model; push it if asked."""
    check = check_pretrained if variant == "pretrained" else check_stanford
    max_abs_diff, n_chans_checked = 0.0, []
    for pos in montages():
        model = build(variant, adapter, len(pos))
        diff = check(model, args.reference_dir, args.backbone, pos)
        max_abs_diff = max(max_abs_diff, diff)
        n_chans_checked.append(len(pos))
    if max_abs_diff > TOLERANCE:
        raise SystemExit(f"{variant}: differs by {max_abs_diff:.2e} > {TOLERANCE}")

    # Users pass their own montage when loading; the weights do not depend on it.
    n_chans = n_chans_checked[0]
    model = build(variant, adapter, n_chans)
    hidden = NOT_PRETRAINED if variant == "pretrained" else ()
    with without(model, hidden):
        n_params = sum(t.numel() for t in model.state_dict().values())
    report = {
        "backbone": f"{BACKBONE_URL} via {BACKBONE_REPO}@{BACKBONE_REVISION}",
        "backbone_sha256": BACKBONE_SHA256,
        "n_params": n_params,
        "max_abs_diff": max_abs_diff,
        "n_chans_checked": n_chans_checked,
        "n_chans": n_chans,
        "n_times": CONFIG["n_times"],
        "torch_version": torch.__version__,
    }
    if variant == "pretrained":
        report["not_saved"] = list(NOT_PRETRAINED)
    else:
        report |= {
            "source": f"{REPOSITORY}@{REVISION}:{ADAPTER}",
            "sha256": ADAPTER_SHA256,
            "adapter_tensors": len(adapter),
            "adapter_params": ADAPTER_PARAMS,
            "lora_scale": LORA_SCALE,
        }
    print(variant, json.dumps(report, indent=2))
    destination = args.output_dir / REPOS[variant]
    destination.mkdir(parents=True, exist_ok=True)
    with without(model, hidden):
        model.save_pretrained(destination)
    # The Hub rejects NaN and Infinity, which Python's json writes by default.
    json.loads(
        (destination / "config.json").read_text(),
        parse_constant=lambda c: (_ for _ in ()).throw(ValueError(f"{c} in config.json")),
    )
    shutil.copyfile(Path(__file__).resolve(), destination / Path(__file__).name)
    (destination / "LICENSE").write_text(LICENSES[variant])
    (destination / "report.json").write_text(json.dumps(report, indent=2))
    # Written last: save_pretrained leaves a placeholder card behind.
    (destination / "README.md").write_text(CARDS[variant](report))
    if args.push_to:
        repo_id = f"{args.push_to}/{REPOS[variant]}"
        api = HfApi()
        api.create_repo(repo_id, exist_ok=True)
        api.upload_folder(
            folder_path=str(destination),
            repo_id=repo_id,
            commit_message=f"CORTEG {variant} model, converted from the official releases",
        )
        print(f"published https://huggingface.co/{repo_id}")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--reference-dir", type=Path, required=True)
    parser.add_argument("--backbone", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--variant", choices=sorted(REPOS), action="append")
    parser.add_argument("--push-to", help="Hub namespace to publish the models to")
    args = parser.parse_args()
    torch.set_num_threads(1)
    adapter_path = args.reference_dir / ADAPTER
    for path, expected in ((adapter_path, ADAPTER_SHA256), (args.backbone, BACKBONE_SHA256)):
        if sha256(path) != expected:
            raise SystemExit(f"{path} is not the released file (sha256 mismatch)")
    adapter = torch.load(adapter_path, map_location="cpu", weights_only=True)
    n_adapter = sum(t.numel() for t in adapter.values())
    if n_adapter != ADAPTER_PARAMS:
        raise SystemExit(f"adapter holds {n_adapter} parameters, not {ADAPTER_PARAMS}")
    for variant in args.variant or REPOS:
        torch.manual_seed(0)
        convert(variant, adapter, args)


if __name__ == "__main__":
    main()