"""Convert an MSST / lucidrains Mel-Band RoFormer .ckpt into a self-describing safetensors file
(model kwargs + inference defaults in the metadata) that toolkit.audio.melbandroformer loads.

    python scripts/convert_melbandroformer.py                       # KimberleyJSN vocals model -> MODELS_PATH/checkpoints
    python scripts/convert_melbandroformer.py --ckpt x.ckpt --config msst_config.yaml --out x.safetensors
"""
import argparse
import inspect
import json
import os
import sys

TOOLKIT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, TOOLKIT_ROOT)

from dotenv import load_dotenv

load_dotenv(os.path.join(TOOLKIT_ROOT, ".env"))  # MODELS_PATH etc., before toolkit.paths is imported

import torch
import yaml
from safetensors.torch import save_file

from toolkit.paths import MODELS_PATH
from toolkit.audio.melbandroformer.model import MelBandRoformer
from toolkit.audio.melbandroformer.separate import DEFAULT_WEIGHTS

# https://huggingface.co/KimberleyJSN/melbandroformer (MIT), config from
# MSST configs/KimberleyJensen/config_vocals_mel_band_roformer_kj.yaml
KJ_VOCALS = dict(
    repo_id="KimberleyJSN/melbandroformer",
    filename="MelBandRoformer.ckpt",
    license="MIT",
    model=dict(
        dim=384,
        depth=6,
        stereo=True,
        num_stems=1,
        time_transformer_depth=1,
        freq_transformer_depth=1,
        num_bands=60,
        dim_head=64,
        heads=8,
        attn_dropout=0,
        ff_dropout=0,
        sample_rate=44100,
        stft_n_fft=2048,
        stft_hop_length=441,
        stft_win_length=2048,
        stft_normalized=False,
        mask_estimator_depth=2,
    ),
    inference=dict(chunk_size=352800, num_overlap=2),
    stems=["vocals"],
)


def config_from_msst_yaml(path):
    with open(path) as f:
        cfg = yaml.load(f, Loader=yaml.FullLoader)  # MSST configs carry !!python/tuple
    allowed = set(inspect.signature(MelBandRoformer.__init__).parameters) - {"self"}
    model = {k: v for k, v in cfg["model"].items() if k in allowed}
    dropped = sorted(set(cfg["model"]) - allowed)
    if dropped:
        print(f"ignoring non-inference model keys: {dropped}")
    assert cfg["model"].get("linear_transformer_depth", 0) == 0, "linear transformer layers are not vendored"
    inference = dict(
        chunk_size=cfg.get("inference", {}).get("chunk_size", cfg["audio"]["chunk_size"]),
        num_overlap=cfg["inference"]["num_overlap"],
    )
    training = cfg.get("training", {})
    stems = [training["target_instrument"]] if training.get("target_instrument") else training.get("instruments", [])
    return model, inference, stems


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--ckpt", default=None, help="local .ckpt (default: download KimberleyJSN/melbandroformer)")
    parser.add_argument("--config", default=None, help="MSST yaml config (default: KJ vocals config)")
    parser.add_argument("--out", default=os.path.join(MODELS_PATH, "checkpoints", DEFAULT_WEIGHTS))
    parser.add_argument("--dtype", default="fp32", choices=["fp32", "fp16", "bf16"])
    args = parser.parse_args()

    if args.config:
        model_cfg, inference_cfg, stems = config_from_msst_yaml(args.config)
        source = args.ckpt
        license_ = ""
    else:
        model_cfg, inference_cfg, stems = KJ_VOCALS["model"], KJ_VOCALS["inference"], KJ_VOCALS["stems"]
        source = f"{KJ_VOCALS['repo_id']}/{KJ_VOCALS['filename']}"
        license_ = KJ_VOCALS["license"]

    ckpt = args.ckpt
    if ckpt is None:
        from huggingface_hub import hf_hub_download
        ckpt = hf_hub_download(KJ_VOCALS["repo_id"], KJ_VOCALS["filename"])

    print(f"loading {ckpt}")
    state_dict = torch.load(ckpt, map_location="cpu", weights_only=True)
    if "state_dict" in state_dict:
        state_dict = state_dict["state_dict"]
    state_dict = {k[len("module."):] if k.startswith("module.") else k: v for k, v in state_dict.items()}

    # validate against the vendored module before writing anything
    model = MelBandRoformer(**model_cfg)
    model.load_state_dict(state_dict, strict=True)

    dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[args.dtype]
    # shared rotary freqs alias one storage in the ckpt; safetensors needs distinct tensors
    tensors = {k: v.detach().to(dtype).clone().contiguous() for k, v in state_dict.items()}

    metadata = dict(
        klass="MelBandRoformer",
        config=json.dumps(model_cfg),
        inference=json.dumps(inference_cfg),
        stems=json.dumps(stems),
        source=str(source),
        license=license_,
    )
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    save_file(tensors, args.out, metadata=metadata)
    n_params = sum(v.numel() for v in tensors.values())
    print(f"wrote {args.out} ({n_params / 1e6:.1f}M params, {args.dtype}, {os.path.getsize(args.out) / 1e6:.0f} MB)")


if __name__ == "__main__":
    main()