on_unknown_ch: Literal['raise', 'warn', 'ignore']
# Behaviour when pre-trained weights are loaded,
# and channels that were not used to pre-train are passed 

# %% ######################################################
# Imports constants and setup
#######################################################
import sys
import torch
import re
from pathlib import Path
import difflib

import torch
from torch import nn
from braindecode.models.labram import Labram, LABRAM_CHANNEL_ORDER

# Add LaBraM repo to path
# > cd ~/Projects/
# > git clone [email protected]:935963004/LaBraM.git
labram_path = Path("~/Projects/LaBraM").expanduser()
sys.path.insert(0, str(labram_path))

from timm.models import create_model
import modeling_finetune  # Register the models to timm

DROPPED_KEYS = ["mask_token", "lm_head.weight", "lm_head.bias"]

# Load checkpoint from GitHub
print("Loading checkpoint from GitHub...")
checkpoint_url = (
    "https://github.com/935963004/LaBraM/raw/main/checkpoints/labram-base.pth"
)
checkpoint = torch.hub.load_state_dict_from_url(
    checkpoint_url,
    map_location="cpu",
    progress=True,
    check_hash=True,
    weights_only=False,
)

# Extract state dict
state_dict = {
    key[8:]: value
    for key, value in checkpoint["model"].items()
    if key.startswith("student.")
}
# Drop unused keys
for k in DROPPED_KEYS:
    del state_dict[k]


# %% ######################################################
# Instantiate the original LaBraM model with pre-trained weights.
#######################################################

# Create model
print("Creating model...")
model = create_model(
    "labram_base_patch200_200",
    # arguments form
    # https://github.com/935963004/LaBraM/blob/c431221e6cfd23dbfa9950e0180682fb322b0548/run_class_finetuning.py#L196C1-L208C6
    # and readme
    pretrained=False,
    num_classes=0,
    drop_rate=0,
    drop_path_rate=0.1,
    attn_drop_rate=0.0,
    drop_block_rate=None,
    use_mean_pooling=False,
    init_scale=0.001,
    use_rel_pos_bias=False,
    use_abs_pos_emb=True,
    init_values=0.1,
    qkv_bias=False,
)
model.eval()

# Load pre-trained weights
model.load_state_dict(state_dict, strict=True)


# %% ######################################################
# Instantiate the Braindecode implementation of LaBraM, rename the weights and load the pre-trained weights.
######################################################

# Instantiate model
chs_info = [{"ch_name": ch_name} for ch_name in LABRAM_CHANNEL_ORDER]
model_braindecode = Labram(
    n_times=3000,
    chs_info=chs_info,
    n_outputs=0,
)
_ = model_braindecode.eval()


# Map keys from original LaBraM to Braindecode implementation
braindecode_state_dict = {}
key_mappings = {}
for k, v in state_dict.items():
    new_k = k

    # simple renames
    if new_k == "pos_embed":
        new_k = "position_embedding"
    elif new_k == "time_embed":
        new_k = "temporal_embedding"

    # blocks mlp fc1/fc2 -> mlp.0 / mlp.2
    new_k = re.sub(r"blocks\.(\d+)\.mlp\.fc1", r"blocks.\1.mlp.0", new_k)
    new_k = re.sub(r"blocks\.(\d+)\.mlp\.fc2", r"blocks.\1.mlp.2", new_k)

    # patch_embed conv/norm mappings:
    m = re.match(r"patch_embed\.conv(\d+)(\..+)$", k)
    if m:
        n = int(m.group(1))
        suffix = m.group(2)
        # map convs to temporal_conv.conv{n}
        new_k = f"patch_embed.temporal_conv.conv{n}{suffix}"

    m2 = re.match(r"patch_embed\.norm(\d+)(\..+)$", k)
    if m2:
        n = int(m2.group(1))
        suffix = m2.group(2)
        new_k = f"patch_embed.temporal_conv.norm{n}{suffix}"

    # fallback: use the simple/new_k computed above
    braindecode_state_dict[new_k] = v
    if new_k != k:
        key_mappings[k] = new_k

print("Key mappings:")
max_len = max(len(old_key) for old_key in key_mappings.keys())
for old_key, new_key in key_mappings.items():
    print(f"{old_key.ljust(max_len)} -> {new_key}")

# Load mapped state dict into model
model_braindecode.load_state_dict(braindecode_state_dict, strict=True)


# %%########################################################
# Compare the outputs of the original and Braindecode implementations for the same input.
##########################################################

rng = torch.Generator().manual_seed(42)
x = torch.randn(1, 128, 3000, generator=rng)

with torch.no_grad():
    y1 = model(x.view(1, 128, 15, 200), return_all_tokens=True)
    y2 = model_braindecode(x, return_all_tokens=True)  # return_patch_tokens=True)

assert y1.shape == y2.shape
assert (y1 == y2).all()

# %%########################################################
# Push weights to hub
#########################################################

model_braindecode.push_to_hub(
    repo_id="braindecode/labram-pretrained", commit_message="Several fixes (see PR#931)"
)


# %%########################################################
# Now load form the newly pushed HF checkpoint
#########################################################

model_hf = Labram.from_pretrained(
    "braindecode/labram-pretrained",
    cache_dir=Path("~/.cache/huggingface/").expanduser(),
)


# %%########################################################
# Finally, compare the outputs of the original and HF model for the same input.
##########################################################

rng = torch.Generator().manual_seed(42)
x = torch.randn(1, 128, 3000, generator=rng)

with torch.no_grad():
    y1 = model(x.view(1, 128, 15, 200), return_all_tokens=True)
    y2 = model_hf(x, return_all_tokens=True)  # return_patch_tokens=True)

assert y1.shape == y2.shape
assert (y1 == y2).all()


# %%########################################################
# Sandbox
#########################################################

def format_model_str(model):
    return list(
        map(lambda line: line.replace(" ", ""), repr(model).splitlines(keepends=True))
    )


model_str = format_model_str(model)
model_braindecode_str = format_model_str(model_braindecode)

# diff = difflib.unified_diff(
diff = difflib.ndiff(
    model_str,
    model_braindecode_str,
    # fromfile="Original LaBraM",
    # tofile="Braindecode LaBraM",
)
print("".join(diff))