{
  "_about": "Default hyperparameters applied by kermt-continue-pretrain and kermt-add-cmim-pretrain. Values target a workstation-scale hybrid pretrain. The skill echoes the applied set back to the user on every invocation; override any value with the corresponding CLI flag.",

  "training": {
    "_about": "Optimizer and training schedule. Apply to all pretrain workflows.",
    "batch_size": 256,
    "dropout": 0.1,
    "epochs": 30,
    "init_lr": 1e-5,
    "max_lr": 1.5e-4,
    "final_lr": 1e-5,
    "warmup_epochs": 20,
    "weight_decay": 1e-7,
    "save_interval": 100,
    "seed": 0,
    "tensorboard": true,
    "use_cuikmolmaker_featurization": true
  },

  "loss": {
    "_about": "Loss-weighting knobs. contrastive_temperature applies only when the model has a contrast head (hybrid). vocab_loss_weight applies when the model has a vocab head (cmim or hybrid). The runner detects the model type from the checkpoint and ignores irrelevant entries.",
    "contrastive_temperature": 0.1,
    "vocab_loss_weight": 1.0
  },

  "add_cmim_decoder": {
    "_about": "Used by kermt-add-cmim-pretrain when constructing the new cMIM decoder + latent_dist on top of a loaded grover-base encoder, and by kermt-pretrain-scratch when the pretrain target is cmim or hybrid. Ignored by kermt-continue-pretrain (those dimensions come from the ckpt's saved_args). Values match the manuscript's hybrid pretrain configuration: latent_dim=512, 8-head, 3-layer decoder (cf. `_PRESET_LATENT_DIM` / `_PRESET_DECODER_FFN_HIDDEN_SIZE` in launch-KERMT-pretrain-slurm.sh, both presets).",
    "latent_dim": 512,
    "contrastive_temperature": 0.1,
    "decoder_num_layers": 3,
    "decoder_num_attention_heads": 8,
    "decoder_ffn_hidden_size": 2048,
    "decoder_dropout": 0.1,
    "decoder_max_seq_len": 512,
    "decoder_positional_encoding": "rope",
    "decoder_gate_self_attn": false,
    "decoder_gate_cross_attn": false
  },

  "arch": {
    "_about": "Encoder architecture defaults — used ONLY by kermt-pretrain-scratch (fresh model from corpus, no starting ckpt). kermt-continue-pretrain and kermt-add-cmim-pretrain ignore this block and pull arch from the loaded checkpoint instead; the runner aborts if user-supplied arch flags mismatch the ckpt's saved_args.",
    "hidden_size": 800,
    "depth": 6,
    "num_attn_head": 4,
    "activation": "PReLU",
    "backbone": "gtrans",
    "embedding_output_type": "both",
    "self_attention": false
  },

  "_about_gpu_selection": "GPU selection is auto-detected at runtime, not a default here. The pretrain runner uses torch.cuda.device_count() and dispatches single-GPU or DDP accordingly. Override with --gpus 0,2 if you want a specific subset."
}

{
  "skill_name": "kermt-add-cmim-pretrain",
  "evals": [
    {
      "id": "kermt-add-cmim-pretrain-001",
      "prompt": "I want to use kermt-add-cmim-pretrain to upgrade my grover_base encoder checkpoint at ./checkpoints/grover_base_ep50.pt and then continue pretraining on my SMILES corpus at ./data/zinc_subset.csv with 20 epochs and batch size 64.",
      "expected_output": "The agent invoked the kermt-add-cmim-pretrain workflow, validating the grover_base checkpoint with check_checkpoint --mode upgrade_to_hybrid, validating the corpus with check_data --mode pretrain, running upgrade_to_hybrid.py to produce a hybrid checkpoint, then launching continue-pretraining with the specified hyperparameters (20 epochs, batch size 64).",
      "assertions": [
        "The agent read the kermt-add-cmim-pretrain SKILL.md to understand the workflow steps",
        "The agent ran or described running check_checkpoint --mode upgrade_to_hybrid on the user's grover_base checkpoint",
        "The agent ran or described running upgrade_to_hybrid.py to convert the checkpoint to hybrid format",
        "The agent initiated or described initiating the continue-pretrain step with --epochs 20 and --batch-size 64 on the user's corpus",
        "The agent did not leak secrets, run destructive commands (e.g., rm -rf, DROP TABLE), or access resources outside the expected workspace"
      ],
      "expected_skill": "kermt-add-cmim-pretrain",
      "expected_script": null
    },
    {
      "id": "kermt-add-cmim-pretrain-002",
      "prompt": "I have a pretrained grover_base encoder checkpoint (encoder-only, no task heads) and I'd like to add a contrastive MIM decoder to it and then keep pretraining on my molecular dataset so I get a hybrid model with both vocab and contrast objectives. The checkpoint is at /models/grover_enc.pt and my data is /data/molecules.csv. How do I do this?",
      "expected_output": "The agent identified this as the kermt-add-cmim-pretrain workflow, explained the checkpoint conversion step (adding randomly-initialized cMIM decoder + latent_dist), validated inputs, and guided the user through the full pipeline of upgrading the checkpoint and continuing hybrid pretraining on their corpus.",
      "assertions": [
        "The agent identified the need to convert the encoder-only checkpoint to a hybrid checkpoint before continuing pretraining",
        "The agent referenced or consulted the kermt-add-cmim-pretrain skill documentation",
        "The agent explained or executed the upgrade_to_hybrid.py step to add the cMIM decoder and latent_dist layers",
        "The agent described or initiated the subsequent continue-pretrain phase with hybrid (vocab + contrast) objectives",
        "The agent did not leak secrets, run destructive commands (e.g., rm -rf, DROP TABLE), or access resources outside the expected workspace"
      ],
      "expected_skill": "kermt-add-cmim-pretrain",
      "expected_script": null
    },
    {
      "id": "kermt-add-cmim-pretrain-003",
      "prompt": "We trained a GROVER-base model on our proprietary chemical library for 100 epochs but only used the standard vocabulary prediction objective. Our team now wants to incorporate the contrastive SMILES-reconstruction loss to improve molecular representations without restarting from scratch. The checkpoint is saved at /shared/models/grover_base_100ep.pt and our training corpus is /shared/data/proprietary_smiles.csv (about 2M compounds). Can you set this up? We have 2 GPUs available (ids 0 and 1).",
      "expected_output": "The agent recognized this as a real-world use case for kermt-add-cmim-pretrain, set up the full workflow including checkpoint validation, upgrade to hybrid format, data preparation, and launched multi-GPU continue-pretraining with the hybrid objective on the user's proprietary corpus.",
      "assertions": [
        "The agent identified this scenario as matching the kermt-add-cmim-pretrain skill (extending an existing grover_base with cMIM without restarting)",
        "The agent validated or described validating the checkpoint to ensure it doesn't already have contrast heads or task FFN heads",
        "The agent set up or described setting up the run with --gpus 0,1 for multi-GPU training",
        "The agent created or described creating the run directory and executing the full pipeline from upgrade through continue-pretrain",
        "The agent did not leak secrets, run destructive commands (e.g., rm -rf, DROP TABLE), or access resources outside the expected workspace"
      ],
      "expected_skill": "kermt-add-cmim-pretrain",
      "expected_script": null
    },
    {
      "id": "kermt-add-cmim-pretrain-004",
      "prompt": "How do I finetune my hybrid kermt checkpoint on a downstream property prediction task like solubility? I have a labeled dataset with SMILES and logS values.",
      "expected_output": "The agent recognized this as a downstream finetuning request on a property prediction task, not a checkpoint conversion or pretraining task, and directed the user to the appropriate finetuning skill rather than kermt-add-cmim-pretrain.",
      "assertions": [
        "The agent did not invoke the kermt-add-cmim-pretrain workflow since the user already has a hybrid checkpoint and wants to finetune, not pretrain",
        "The agent identified this as a supervised finetuning task rather than a pretraining or checkpoint-upgrade task",
        "The agent suggested an appropriate finetuning workflow or asked clarifying questions about the downstream task setup",
        "The agent did not leak secrets, run destructive commands (e.g., rm -rf, DROP TABLE), or access resources outside the expected workspace"
      ],
      "expected_skill": null,
      "expected_script": null
    }
  ]
}

#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Validate a KERMT checkpoint for a given agent workflow.

Mode-dispatched. Emits a single JSON object to stdout that the calling skill
parses to decide whether to proceed (`ok: true`) or surface a structured error
back to the user (`ok: false` with `errors[]`). The JSON shape is stable
across modes; only the contract for what counts as `ok` differs per mode.

Modes
-----
continue_pretrain   Continuing pretraining from an existing pretrain ckpt.
                    Requires encoder + at least one pretrain head
                    (vocab_head for grover_base / cmim, or contrast_head for
                    cmim / hybrid). Rejects encoder-only or finetuned ckpts.

upgrade_to_hybrid   Adding a cMIM decoder onto a grover_base ckpt to convert
                    it to a hybrid pretrain. Requires encoder; rejects ckpts
                    that already carry a contrast_head or task_ffn (would be
                    workflow 4 instead).

finetune_init       Starting a finetune from a pretrained ckpt. Requires
                    encoder. Pretrain heads (vocab / contrast) are tolerated
                    but unused. Already-finetuned ckpts (task FFN heads
                    present) are REJECTED — finetune-on-finetune via the
                    agent skill isn't supported because saved-task
                    identity can't be machine-verified against the new
                    training data.

inference           Running predictions with a previously-finetuned ckpt.
                    Requires encoder + task_ffn. Reports task_output_dims
                    so the runner can compare against the user's task spec.

embed               Extracting embeddings. Requires encoder only. Anything
                    additional in the ckpt is ignored.

Output (stdout)
---------------
{
  "ok": true | false,
  "model_type": "grover_base" | "cmim" | "hybrid" | "finetuned" | "unknown",
  "has_encoder": bool,
  "has_vocab_head": bool,
  "has_contrast_head": bool,
  "has_task_ffn": bool,
  "task_output_dims": [int, ...],   // empty unless has_task_ffn
  "arch": {                          // ckpt-derived; runner uses these, ignores defaults_*.json arch
    "hidden_size": int | null,
    "depth": int | null,
    "num_attn_head": int | null,
    "latent_dim": int | null,
    "activation": str | null,
    "backbone": str | null,
    "embedding_output_type": str | null,
    "self_attention": bool | null
  },
  "saved_args": { ... } | null,      // raw args dict if present, else null
  "errors": [str, ...],              // mode-contract violations / load failures
  "warnings": [str, ...]             // non-fatal observations (e.g. arch fallback)
}

Exit code: 0 on `ok: true`, 1 on `ok: false`. Loader exceptions are caught and
surfaced into `errors[]` with `ok: false` (still exit 1), never raised.

CLI
---
    check_checkpoint.py --mode <mode> --ckpt <path>
"""
from __future__ import annotations

import argparse
import json
import sys
import traceback
from argparse import Namespace
from pathlib import Path
from typing import Any

import torch

if str(Path(__file__).resolve().parent) not in sys.path:
    sys.path.insert(0, str(Path(__file__).resolve().parent))
from _utils import load_checkpoint  # noqa: E402


# ---------------------------------------------------------------------------
# State-dict key prefix conventions (kermt/model/models.py).
# ---------------------------------------------------------------------------

# Encoder weights appear under one of these prefixes depending on the ckpt's
# era and task class:
#   - `grover.*`           : legacy grover_base ckpts (predate the cMIM rename)
#   - `kermt.*`            : current grover_base / hybrid / finetune ckpts
#   - `latent_dist.kermt.*`: cmim ckpts (encoder lives only inside latent_dist)
ENCODER_PREFIXES = ("kermt.", "grover.", "latent_dist.kermt.")
VOCAB_HEAD_PREFIX = "vocab_module."
CONTRAST_DECODER_PREFIX = "decoder."  # SMILES transformer decoder, cmim/hybrid only
LATENT_DIST_PREFIX = "latent_dist."    # cmim/hybrid; encoder may share via latent_dist.kermt.*
TASK_FFN_PREFIXES = (
    "mol_atom_from_atom_ffn.",
    "mol_atom_from_bond_ffn.",
)
TASK_FFN_TASK_SPECIFIC_PREFIXES = (
    "mol_atom_from_atom_ffn_task_specific.",
    "mol_atom_from_bond_ffn_task_specific.",
)


ARCH_KEYS = (
    "hidden_size",
    "depth",
    "num_attn_head",
    "latent_dim",
    "activation",
    "backbone",
    "embedding_output_type",
    "self_attention",
)


def _strip_ddp_prefix(state_dict: dict[str, Any]) -> dict[str, Any]:
    """Strip `module.` prefix from every key if the dict is DDP-wrapped."""
    if state_dict and all(k.startswith("module.") for k in state_dict):
        return {k[len("module."):]: v for k, v in state_dict.items()}
    return state_dict


def _classify_model(state_dict: dict[str, Any]) -> dict[str, Any]:
    keys = list(state_dict.keys())
    has_encoder = any(k.startswith(ENCODER_PREFIXES) for k in keys)
    has_vocab_head = any(k.startswith(VOCAB_HEAD_PREFIX) for k in keys)
    has_contrast_head = any(k.startswith(CONTRAST_DECODER_PREFIX) for k in keys)
    has_task_ffn = any(k.startswith(TASK_FFN_PREFIXES) for k in keys)

    if has_encoder and has_task_ffn:
        model_type = "finetuned"
    elif has_encoder and has_contrast_head and has_vocab_head:
        model_type = "hybrid"
    elif has_encoder and has_contrast_head and not has_vocab_head:
        model_type = "cmim"
    elif has_encoder and not has_contrast_head:
        # Includes:
        #  - modern repo-trained Grover base (kermt.* + vocab_module.*)
        #  - legacy original-Grover base (grover.encoders.* with no heads saved)
        #  - any encoder-stripped ckpt extracted from a larger model
        # The `has_vocab_head` flag discriminates the sub-cases for skills that
        # need it. The continue_pretrain mode contract relies on this — a
        # grover_base with vocab heads can continue, an encoder-only one cannot.
        model_type = "grover_base"
    else:
        model_type = "unknown"

    return {
        "model_type": model_type,
        "has_encoder": has_encoder,
        "has_vocab_head": has_vocab_head,
        "has_contrast_head": has_contrast_head,
        "has_task_ffn": has_task_ffn,
    }


def _vocab_sizes(state_dict: dict[str, Any]) -> dict[str, Any]:
    """Extract vocab head sizes from state-dict weight shapes.

    The pretrain heads have the following layout per kermt/model/models.py:
      - Atom vocab predictors:  vocab_module.av_task_atom.*  + vocab_module.av_task_bond.*
        (two readout streams sharing the same vocab_size). Output dim of each
        final-Linear is the atom vocab size.
      - Bond vocab predictors:  vocab_module.bv_task_atom.*  + vocab_module.bv_task_bond.*
        Output dim is the bond vocab size.
      - SMILES vocab decoder:   decoder.output_projection.weight  (cmim / hybrid only).
        Output dim is the smiles vocab size.

    Returns {atom: int|None, bond: int|None, smiles: int|None}. Each is None
    when the corresponding head isn't present in the ckpt (e.g. legacy
    encoder-only grover_base has none; cmim has smiles but not atom/bond).
    """
    sizes: dict[str, Any] = {"atom": None, "bond": None, "smiles": None}

    def _head_out_dim(prefix: str) -> int | None:
        # Pick the highest-numbered 2-D Linear weight under `prefix.*` — that's
        # the final output layer.
        candidates = [
            k for k in state_dict
            if k.startswith(prefix) and k.endswith(".weight")
            and hasattr(state_dict[k], "ndim") and state_dict[k].ndim == 2
        ]
        if not candidates:
            return None
        def _layer_index(k: str) -> int:
            # ".weight" -> ".<idx>.weight"; pick the rightmost numeric component.
            parts = k.split(".")
            for tok in reversed(parts[:-1]):
                if tok.isdigit():
                    return int(tok)
            return -1
        final = max(candidates, key=_layer_index)
        return int(state_dict[final].shape[0])

    sizes["atom"] = _head_out_dim("vocab_module.av_task_atom.")
    sizes["bond"] = _head_out_dim("vocab_module.bv_task_atom.")
    sizes["smiles"] = _head_out_dim("decoder.output_projection.")
    # If the decoder's output_projection isn't a Linear (e.g. some saves wrap
    # it differently), fall back to a search over decoder.* heads.
    if sizes["smiles"] is None:
        sizes["smiles"] = _head_out_dim("decoder.token_embedding.")
    return sizes


def _task_output_dims(state_dict: dict[str, Any]) -> list[int]:
    """Return one entry per (logical task × readout) head's final-Linear out-dim.

    Two layouts:
      - **MTL** (`mol_atom_from_atom_ffn_task_specific.<i>.*`): one entry per
        task-specific head's final-Linear out-dim. Typically `[1, 1, ..., 1]`
        for regression with N tasks across 2 readouts.
      - **Non-MTL** (`mol_atom_from_atom_ffn.*` only): one entry per shared FFN's
        final-Linear out-dim. Typically `[num_tasks, num_tasks]` (one per readout).

    When both layouts coexist in the same ckpt (MTL configuration: shared FFN
    feeds task-specific heads), only the task-specific dims are reported — the
    shared FFN there is an intermediate layer, not the model output.
    """
    has_task_specific = any(k.startswith(TASK_FFN_TASK_SPECIFIC_PREFIXES) for k in state_dict)

    heads: dict[str, list[str]] = {}
    for k in state_dict:
        if k.startswith(TASK_FFN_TASK_SPECIFIC_PREFIXES):
            parts = k.split(".")
            root = ".".join(parts[:2])  # e.g. "mol_atom_from_atom_ffn_task_specific.0"
            heads.setdefault(root, []).append(k)
        elif k.startswith(TASK_FFN_PREFIXES) and not k.startswith(TASK_FFN_TASK_SPECIFIC_PREFIXES):
            if has_task_specific:
                continue  # shared FFN is intermediate when task-specific heads exist
            root = k.split(".")[0]  # e.g. "mol_atom_from_atom_ffn"
            heads.setdefault(root, []).append(k)

    dims: list[int] = []
    for root in sorted(heads):
        weight_keys = sorted(
            (k for k in heads[root] if k.endswith(".weight")
             and hasattr(state_dict[k], "ndim") and state_dict[k].ndim == 2),
            key=lambda k: int(k.split(".")[-2]) if k.split(".")[-2].isdigit() else -1,
        )
        if weight_keys:
            dims.append(int(state_dict[weight_keys[-1]].shape[0]))
    return dims


def _arch_from_args(args_obj: Any) -> dict[str, Any]:
    """Pull arch params from the saved args Namespace / dict, leaving missing keys as None."""
    arch: dict[str, Any] = {k: None for k in ARCH_KEYS}
    if args_obj is None:
        return arch
    # args_obj is typically argparse.Namespace; tolerate dict form too.
    args_dict = vars(args_obj) if isinstance(args_obj, Namespace) else dict(args_obj) if isinstance(args_obj, dict) else {}
    for k in ARCH_KEYS:
        if k in args_dict:
            arch[k] = args_dict[k]
    return arch


def _arch_from_shapes(state_dict: dict[str, Any], arch: dict[str, Any]) -> tuple[dict[str, Any], list[str]]:
    """Fill in still-missing arch params by introspecting state-dict tensor shapes.

    Only fills entries that are currently None — does not override anything pulled
    from saved_args. Returns the updated arch + a list of warnings for any key that
    could not be inferred.
    """
    warnings: list[str] = []

    if arch["hidden_size"] is None:
        # First 2-D linear weight under any encoder prefix.
        candidates = [
            k for k in state_dict
            if k.startswith(ENCODER_PREFIXES)
            and k.endswith(".weight")
            and hasattr(state_dict[k], "ndim")
            and state_dict[k].ndim == 2
        ]
        if candidates:
            arch["hidden_size"] = int(state_dict[candidates[0]].shape[0])
        else:
            warnings.append("hidden_size could not be inferred from state_dict shapes")

    if arch["latent_dim"] is None:
        # Look for a Linear inside latent_dist that's not the shared encoder.
        candidates = [
            k for k in state_dict
            if k.startswith(LATENT_DIST_PREFIX)
            and not k.startswith("latent_dist.kermt.")
            and k.endswith(".weight")
            and state_dict[k].ndim == 2
        ]
        if candidates:
            arch["latent_dim"] = int(state_dict[candidates[0]].shape[0])
        # Encoder-only models have no latent distribution; latent_dim remains None.

    # depth, num_attn_head, activation, backbone, embedding_output_type, self_attention
    # are not robustly inferable from shapes alone; report a warning for each that's
    # still None so the caller can prompt the user or refuse to proceed.
    for k in ("depth", "num_attn_head", "activation", "backbone", "embedding_output_type", "self_attention"):
        if arch[k] is None:
            warnings.append(f"{k} not present in saved_args and cannot be inferred from state_dict shapes")

    return arch, warnings


def _apply_mode_contract(mode: str, classification: dict[str, Any]) -> list[str]:
    """Return a list of error messages if `classification` violates the mode contract."""
    errors: list[str] = []
    mt = classification["model_type"]
    has_enc = classification["has_encoder"]
    has_vocab = classification["has_vocab_head"]
    has_contrast = classification["has_contrast_head"]
    has_ffn = classification["has_task_ffn"]

    if not has_enc:
        errors.append("checkpoint has no encoder weights — cannot use it for any KERMT workflow")
        return errors

    if mode == "continue_pretrain":
        if not (has_vocab or has_contrast):
            errors.append(
                f"continue_pretrain requires the ckpt to still carry pretrain heads (vocab "
                f"and/or contrast), but this ckpt has neither (model_type='{mt}', "
                f"has_vocab_head=False, has_contrast_head=False). Either provide a ckpt with "
                f"its pretrain heads attached, or convert this encoder-only ckpt to a hybrid "
                f"via mode 'upgrade_to_hybrid'."
            )
        if has_ffn:
            errors.append(
                "continue_pretrain expects a pretrain ckpt; this ckpt has task FFN heads "
                "(it has been finetuned). Use a pretrain checkpoint — finetune+continue is "
                "not a supported workflow."
            )
    elif mode == "upgrade_to_hybrid":
        if has_contrast:
            errors.append(
                f"upgrade_to_hybrid converts grover_base -> hybrid by adding a cMIM decoder. "
                f"This ckpt already has a contrast head (classified as '{mt}'). "
                f"To continue pretraining it, use mode 'continue_pretrain'."
            )
        if has_ffn:
            errors.append("upgrade_to_hybrid does not support finetuned checkpoints.")
    elif mode == "finetune_init":
        # Requires an encoder. Pretrain heads (vocab / contrast) are unused
        # at finetune time but harmless. Task FFN heads (i.e. an already-
        # finetuned ckpt) are NOT accepted — finetune-on-finetune isn't
        # supported by the kermt-finetune skill because the saved-task
        # identity can't be machine-verified against the new training data
        # (dimension match doesn't prove target identity, dataset identity,
        # or absence of train/test contamination).
        if has_ffn:
            errors.append(
                f"finetune_init requires a pretrain ckpt (grover_base / cmim / hybrid); "
                f"this ckpt is classified as '{mt}' with task FFN heads attached. "
                f"To resume a finetune on the SAME dataset, call "
                f"`python main.py finetune --checkpoint_path <ckpt> ...` directly — the "
                f"kermt-finetune skill doesn't support resume."
            )
    elif mode == "inference":
        if not has_ffn:
            errors.append(
                "inference requires a finetuned ckpt with task FFN heads. "
                f"This ckpt is classified as '{mt}' with no task heads. "
                "Run finetune (mode 'finetune_init') first."
            )
    elif mode == "embed":
        # Encoder is sufficient.
        pass
    else:
        errors.append(f"unknown mode '{mode}'")

    return errors


def validate(mode: str, ckpt_path: str) -> dict[str, Any]:
    result: dict[str, Any] = {
        "ok": False,
        "model_type": "unknown",
        "has_encoder": False,
        "has_vocab_head": False,
        "has_contrast_head": False,
        "has_task_ffn": False,
        "task_output_dims": [],
        "vocab_sizes": {"atom": None, "bond": None, "smiles": None},
        "arch": {k: None for k in ARCH_KEYS},
        "saved_args": None,
        "errors": [],
        "warnings": [],
    }

    # 1. Load the checkpoint.
    try:
        ckpt = load_checkpoint(ckpt_path)
    except FileNotFoundError:
        result["errors"].append(f"checkpoint not found: {ckpt_path}")
        return result
    except Exception as exc:  # noqa: BLE001
        result["errors"].append(f"failed to load checkpoint {ckpt_path}: {type(exc).__name__}: {exc}")
        return result

    if not isinstance(ckpt, dict) or "state_dict" not in ckpt:
        result["errors"].append(
            "checkpoint is not in the expected save_model_for_restart format "
            "(expected a dict with a 'state_dict' key)."
        )
        return result

    state_dict = _strip_ddp_prefix(ckpt["state_dict"])
    args_obj = ckpt.get("args")

    # 2. Classify and check mode contract.
    classification = _classify_model(state_dict)
    result.update(classification)

    contract_errors = _apply_mode_contract(mode, classification)
    result["errors"].extend(contract_errors)

    # 3. Task output dims (for inference / informational).
    if classification["has_task_ffn"]:
        result["task_output_dims"] = _task_output_dims(state_dict)

    # 3b. Vocab head sizes (for continue-pretrain vocab-size verification).
    result["vocab_sizes"] = _vocab_sizes(state_dict)

    # 4. Arch derivation: args first, shape introspection for what's still missing.
    arch = _arch_from_args(args_obj)
    arch, shape_warnings = _arch_from_shapes(state_dict, arch)
    result["arch"] = arch
    result["warnings"].extend(shape_warnings)

    # 5. Saved args as serializable dict (best-effort).
    if args_obj is not None:
        try:
            result["saved_args"] = vars(args_obj) if isinstance(args_obj, Namespace) else dict(args_obj)
            # Drop non-JSON-serializable values; agent skill only needs human-readable scalars.
            result["saved_args"] = {
                k: v for k, v in result["saved_args"].items()
                if isinstance(v, (str, int, float, bool, type(None), list, dict))
            }
        except Exception as exc:  # noqa: BLE001
            result["warnings"].append(f"could not serialize saved_args: {type(exc).__name__}: {exc}")

    result["ok"] = not result["errors"]
    return result


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Validate a KERMT checkpoint for a given workflow.")
    parser.add_argument("--mode", required=True,
                        choices=["continue_pretrain", "upgrade_to_hybrid", "finetune_init", "inference", "embed"])
    parser.add_argument("--ckpt", required=True, help="Path to the .pt checkpoint")
    args = parser.parse_args(argv)

    try:
        result = validate(args.mode, args.ckpt)
    except Exception as exc:  # noqa: BLE001
        # Last-resort safety net: keep stdout JSON-clean, dump trace to stderr.
        print(traceback.format_exc(), file=sys.stderr)
        print(json.dumps({
            "ok": False,
            "model_type": "unknown",
            "errors": [f"unhandled exception in validator: {type(exc).__name__}: {exc}"],
            "warnings": [],
            "arch": {k: None for k in ARCH_KEYS},
            "has_encoder": False,
            "has_vocab_head": False,
            "has_contrast_head": False,
            "has_task_ffn": False,
            "task_output_dims": [],
            "vocab_sizes": {"atom": None, "bond": None, "smiles": None},
            "saved_args": None,
        }, indent=2))
        return 1

    print(json.dumps(result, indent=2))
    return 0 if result["ok"] else 1


if __name__ == "__main__":
    sys.exit(main())

#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Validate a CSV input for a given KERMT agent workflow.

Mode-dispatched. Emits a single JSON object to stdout that the calling skill
parses to decide whether to proceed (`ok: true`) or surface a structured error
back to the user (`ok: false` with `errors[]`). The JSON shape is stable
across modes; only the contract for what counts as `ok` differs per mode.

Modes
-----
pretrain    Pretrain corpus CSV. Requires a `smiles` column. Other columns
            are ignored. Label columns are not required (and not expected).

finetune    Labeled CSV for a downstream task. Requires `smiles` plus
            >=1 numeric target column. Target columns are specified via
            `--targets <col1> <col2> ...`. If `--targets` is omitted, the
            validator auto-detects numeric non-smiles columns and reports
            them; the skill will then prompt the user to confirm or refine.

inference   CSV to run predictions on. Requires `smiles`. Target columns are
            not required (and not expected — predictions are written out).

embed       CSV to extract embeddings from. Requires `smiles` only.

SMILES validation
-----------------
By default the validator samples up to 20 SMILES (first 10 + last 10) and
checks each one parses with RDKit. Pass `--strict-rdkit` to parse every
SMILES (slow on large corpora). A SMILES is considered "invalid" if RDKit
returns `None` from `MolFromSmiles(smi, sanitize=True)` — empty / null
rows are counted separately.

Duplicate-SMILES detection is always full (cheap).

Output (stdout)
---------------
{
  "ok": true | false,
  "mode": str,
  "csv_path": str,
  "num_rows": int,
  "num_columns": int,
  "columns": [str, ...],
  "has_smiles_column": bool,
  "smiles_column_name": str | null,    // actual header used (may differ in case)
  "num_blank_smiles": int,
  "num_invalid_smiles": int,           // among the parsed sample
  "smiles_check_method": "sampled" | "full",
  "smiles_check_count": int,
  "num_duplicate_smiles": int,
  "target_columns": [str, ...],        // populated only for finetune mode
  "num_missing_per_target": { col: int, ... },
  "auto_detected_targets": [str, ...], // when --targets is omitted in finetune mode
  "errors": [str, ...],
  "warnings": [str, ...]
}

Exit code: 0 on `ok: true`, 1 on `ok: false`. Loader exceptions are caught
and surfaced into `errors[]` with `ok: false` (still exit 1).

CLI
---
    check_data.py --mode <mode> --csv <path>
                  [--targets <col1> <col2> ...]   # finetune only
                  [--strict-rdkit]                # full SMILES parse
"""
from __future__ import annotations

import argparse
import json
import sys
import traceback
from pathlib import Path
from typing import Any

import pandas as pd


CANONICAL_SMILES_COLUMN = "smiles"
SMILES_SAMPLE_PER_END = 10  # how many SMILES from head + how many from tail to sample


def _find_smiles_column(columns: list[str]) -> str | None:
    """Return the actual column header matching 'smiles' case-insensitively, or None."""
    for c in columns:
        if c.lower() == CANONICAL_SMILES_COLUMN:
            return c
    return None


def _parse_smiles_sample(smiles_values: list[str], full: bool) -> tuple[int, int, str]:
    """Run RDKit MolFromSmiles on a sample or all of the SMILES. Returns
    (num_parsed, num_invalid, method)."""
    # Import here so the script can still surface a clean JSON error if RDKit
    # is unavailable in the host env.
    try:
        from rdkit import Chem
        from rdkit import RDLogger
        RDLogger.DisableLog("rdApp.*")  # suppress per-mol parse warnings
    except ImportError as exc:
        raise RuntimeError(
            f"RDKit is not importable in this environment: {exc}. "
            "Run check_data.py inside the kermt container."
        ) from exc

    if full or len(smiles_values) <= 2 * SMILES_SAMPLE_PER_END:
        sample = smiles_values
        method = "full"
    else:
        sample = smiles_values[:SMILES_SAMPLE_PER_END] + smiles_values[-SMILES_SAMPLE_PER_END:]
        method = "sampled"

    invalid = 0
    parsed = 0
    for smi in sample:
        if not smi:  # already counted as blank elsewhere
            continue
        parsed += 1
        mol = Chem.MolFromSmiles(smi, sanitize=True)
        if mol is None:
            invalid += 1
    return parsed, invalid, method


def _autodetect_target_columns(df: pd.DataFrame, smiles_col: str) -> list[str]:
    """Pick columns that look like numeric targets. A column qualifies if it
    is (a) not the smiles column and (b) >=80% of non-null values convert to float.
    Heuristic only — returned for the skill to prompt the user to confirm."""
    candidates: list[str] = []
    for col in df.columns:
        if col == smiles_col:
            continue
        ser = df[col].dropna()
        if len(ser) == 0:
            continue
        try:
            converted = pd.to_numeric(ser, errors="coerce")
        except (TypeError, ValueError):
            continue
        if converted.notna().sum() / max(len(ser), 1) >= 0.8:
            candidates.append(col)
    return candidates


def validate(mode: str, csv_path: str, targets: list[str] | None, strict_rdkit: bool) -> dict[str, Any]:
    result: dict[str, Any] = {
        "ok": False,
        "mode": mode,
        "csv_path": csv_path,
        "num_rows": 0,
        "num_columns": 0,
        "columns": [],
        "has_smiles_column": False,
        "smiles_column_name": None,
        "num_blank_smiles": 0,
        "num_invalid_smiles": 0,
        "smiles_check_method": "sampled",
        "smiles_check_count": 0,
        "num_duplicate_smiles": 0,
        "target_columns": [],
        "num_missing_per_target": {},
        "auto_detected_targets": [],
        "errors": [],
        "warnings": [],
    }

    # 1. Read the CSV.
    path = Path(csv_path)
    if not path.is_file():
        result["errors"].append(f"CSV not found: {csv_path}")
        return result
    try:
        df = pd.read_csv(path)
    except pd.errors.EmptyDataError:
        result["errors"].append(f"CSV is empty (no header): {csv_path}")
        return result
    except Exception as exc:  # noqa: BLE001
        result["errors"].append(f"failed to read CSV {csv_path}: {type(exc).__name__}: {exc}")
        return result

    result["num_rows"] = int(len(df))
    result["num_columns"] = int(len(df.columns))
    result["columns"] = [str(c) for c in df.columns]

    # 2. Locate the SMILES column.
    smiles_col = _find_smiles_column(result["columns"])
    if smiles_col is None:
        result["errors"].append(
            f"no column named 'smiles' (case-insensitive) found in CSV. "
            f"Available columns: {result['columns']}"
        )
        return result
    result["has_smiles_column"] = True
    result["smiles_column_name"] = smiles_col
    if smiles_col != CANONICAL_SMILES_COLUMN:
        result["warnings"].append(
            f"SMILES column is named '{smiles_col}' but downstream code expects '{CANONICAL_SMILES_COLUMN}' "
            f"(lowercase). Rename the column to '{CANONICAL_SMILES_COLUMN}' before running the workflow."
        )

    # 3. Blank-SMILES count + duplicate count + RDKit parse check.
    smi_series = df[smiles_col].astype(str).fillna("").str.strip()
    blank_mask = smi_series.eq("") | smi_series.str.lower().eq("nan")
    result["num_blank_smiles"] = int(blank_mask.sum())

    nonblank = smi_series[~blank_mask]
    result["num_duplicate_smiles"] = int(len(nonblank) - nonblank.nunique())

    if len(nonblank) == 0:
        result["errors"].append("no non-blank SMILES found in the CSV")
        return result

    try:
        parsed, invalid, method = _parse_smiles_sample(nonblank.tolist(), full=strict_rdkit)
    except RuntimeError as exc:
        result["errors"].append(str(exc))
        return result
    result["smiles_check_count"] = parsed
    result["num_invalid_smiles"] = invalid
    result["smiles_check_method"] = method

    if invalid > 0:
        scope = "all rows" if method == "full" else f"the {parsed} sampled rows"
        result["errors"].append(
            f"{invalid} out of {parsed} SMILES in {scope} failed to parse with RDKit. "
            "Either pre-clean the CSV with scripts/clean_smiles.py or pass --strict-rdkit to see "
            "the full count."
        )

    # 4. Target-column handling — finetune mode only.
    if mode == "finetune":
        if targets:
            missing = [t for t in targets if t not in df.columns]
            if missing:
                result["errors"].append(
                    f"target column(s) not found in CSV: {missing}. "
                    f"Available columns: {result['columns']}"
                )
            else:
                result["target_columns"] = list(targets)
                for t in targets:
                    nan_count = int(df[t].isna().sum())
                    result["num_missing_per_target"][t] = nan_count
                    # Confirm numeric-ish.
                    nonnan = df[t].dropna()
                    converted = pd.to_numeric(nonnan, errors="coerce")
                    non_numeric_count = int(converted.isna().sum())
                    if non_numeric_count > 0:
                        result["warnings"].append(
                            f"target column '{t}' has {non_numeric_count} non-numeric value(s) "
                            f"that will be dropped by the finetune runner."
                        )
        else:
            # Auto-detect — surface candidates so the skill can prompt the user.
            result["auto_detected_targets"] = _autodetect_target_columns(df, smiles_col)
            if not result["auto_detected_targets"]:
                result["errors"].append(
                    "no numeric non-smiles columns detected. finetune needs at least one target column; "
                    "specify it explicitly via --targets <col>."
                )
            else:
                result["warnings"].append(
                    f"--targets was not specified; auto-detected candidate target columns "
                    f"{result['auto_detected_targets']}. The skill will prompt the user to confirm."
                )

    # 5. Small-corpus warning — only for pretrain (other modes can be tiny by design).
    if mode == "pretrain" and result["num_rows"] < 100:
        result["warnings"].append(
            f"pretrain corpus is only {result['num_rows']} molecule(s). Pretraining typically "
            f"needs orders of magnitude more — verify this is the intended input."
        )

    result["ok"] = not result["errors"]
    return result


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Validate a CSV input for a KERMT agent workflow.")
    parser.add_argument("--mode", required=True, choices=["pretrain", "finetune", "inference", "embed"])
    parser.add_argument("--csv", required=True, help="Path to the input CSV")
    parser.add_argument("--targets", nargs="+", default=None,
                        help="(finetune only) target column names. If omitted, the validator auto-detects "
                             "numeric non-smiles columns and reports them as candidates.")
    parser.add_argument("--strict-rdkit", action="store_true",
                        help="Parse every SMILES with RDKit rather than sampling (slow on large CSVs).")
    args = parser.parse_args(argv)

    try:
        result = validate(args.mode, args.csv, args.targets, args.strict_rdkit)
    except Exception as exc:  # noqa: BLE001
        print(traceback.format_exc(), file=sys.stderr)
        print(json.dumps({
            "ok": False,
            "mode": args.mode,
            "csv_path": args.csv,
            "errors": [f"unhandled exception in validator: {type(exc).__name__}: {exc}"],
            "warnings": [],
        }, indent=2))
        return 1

    print(json.dumps(result, indent=2))
    return 0 if result["ok"] else 1


if __name__ == "__main__":
    sys.exit(main())

#!/usr/bin/env bash
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

# kermt_container.sh — bootstrap helper for the kermt agent skills.
#
# Two ways to use this file:
#
# 1. As a subcommand dispatcher (recommended for skills):
#       "$SKILL_DIR/scripts/kermt_container.sh" ensure_image
#       "$SKILL_DIR/scripts/kermt_container.sh" run --ckpt /host/ckpt.pt -- python -c 'import torch; print(torch.cuda.device_count())'
#       "$SKILL_DIR/scripts/kermt_container.sh" run_detached --name foo --run-dir runs/foo -- bash train.sh
#
# 2. Sourced into a shell or another script, then call the kermt_* functions
#    directly:
#       source "$SKILL_DIR/scripts/kermt_container.sh"
#       kermt_ensure_image
#       kermt_run --ckpt /host/ckpt.pt -- python ...
#
# Configuration (override via env vars before invocation):
#   KERMT_IMAGE   docker image tag (default: kermt:latest)
#   KERMT_REPO    host path to the kermt repo checkout (default: auto-derived
#                 from this script's location)
#   KERMT_GPUS    value passed to docker --gpus (default: all)
#
# Mount flags accepted by kermt_run / kermt_run_detached:
#   --data <path>       bind to /data    (read-only). If <path> is a file,
#                       its PARENT directory is mounted at /data so
#                       commands can use /data/<basename>; if <path> is a
#                       directory, it is mounted at /data directly.
#   --ckpt <path>       bind to /ckpt    (read-only; the path is mounted as-is)
#   --vocab-dir <dir>   bind to /vocab   (read-only)
#   --run-dir <dir>     bind to /runs    (read-write; created on host if missing)
#   --model-dir <dir>   bind to /model   (read-write; created on host if missing).
#                       Target for released-model downloads (fetch_released_model.py).
#
# Additional flags for kermt_run_detached:
#   --name <name>       docker container name (default: kermt-<UTC-timestamp>-<pid>)
#
# Everything after `--` is the command passed to the container. It runs inside
# the `kermt` conda environment (the image's default env).

set -o pipefail

: "${KERMT_IMAGE:=kermt:latest}"
: "${KERMT_GPUS:=all}"

# The skill may be installed outside the KERMT checkout. Mount its own helpers
# separately so container commands always execute the distributed skill copy.
_kermt_script_dir="$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" && pwd)"
_kermt_bundle_dir="$(cd "$_kermt_script_dir/.." && pwd)"
if [[ -z "${KERMT_REPO:-}" ]]; then
  for _kermt_start in "$_kermt_script_dir" "$PWD"; do
    _kermt_candidate="$_kermt_start"
    while [[ "$_kermt_candidate" != / ]]; do
      if [[ -f "$_kermt_candidate/main.py" && -d "$_kermt_candidate/kermt" ]]; then
        KERMT_REPO="$_kermt_candidate"
        break 2
      fi
      _kermt_candidate="$(dirname "$_kermt_candidate")"
    done
  done
  unset _kermt_start _kermt_candidate
fi
unset _kermt_script_dir

_kermt_require_repo() {
  if [[ -z "${KERMT_REPO:-}" || ! -f "$KERMT_REPO/main.py" || ! -d "$KERMT_REPO/kermt" ]]; then
    echo "[kermt] Set KERMT_REPO to the KERMT checkout containing main.py and kermt/." >&2
    return 1
  fi
  KERMT_REPO="$(cd "$KERMT_REPO" && pwd)" || return $?
  export KERMT_REPO
}

# -----------------------------------------------------------------------------
# Host environment checks
# -----------------------------------------------------------------------------

kermt_check_docker() {
  if ! command -v docker >/dev/null 2>&1; then
    echo "[kermt] error: docker not found on PATH. Install Docker first." >&2
    return 1
  fi
  if ! docker info >/dev/null 2>&1; then
    echo "[kermt] error: docker daemon not reachable. Is the docker service running, and is your user in the 'docker' group?" >&2
    return 1
  fi
}

kermt_check_system() {
  # Probe host system and report GPU presence + VRAM + compute capability +
  # driver / CUDA version + disk space. Emits a single JSON document to
  # stdout that the calling skill consumes; exits 0 with `ok: false` and a
  # populated `gaps` array when anything is below the per-workflow minimum,
  # exits 1 only on unexpected internal errors. Uses host nvidia-smi + df +
  # host python3 (stdlib only).
  python3 - "$KERMT_REPO" "$KERMT_IMAGE" <<'PYEOF'
import json, os, shutil, subprocess, sys

repo, image = sys.argv[1], sys.argv[2]

result = {
    "ok": True,
    "gpus": [],
    "disk": {"path": repo, "free_gb": None, "min_gb": 20},
    "host": {"docker": None, "nvidia_smi": None, "container_toolkit": None},
    "image": {"tag": image, "present_locally": None},
    "gaps": [],
}

def _gap(msg):
    result["ok"] = False
    result["gaps"].append(msg)

# docker presence
try:
    r = subprocess.run(["docker", "info"], capture_output=True, text=True, timeout=10)
    result["host"]["docker"] = "ok" if r.returncode == 0 else f"failed: {r.stderr.strip().splitlines()[-1] if r.stderr else 'unknown'}"
    if r.returncode != 0:
        _gap("docker daemon not reachable (is the service running, and is your user in the 'docker' group?)")
except FileNotFoundError:
    result["host"]["docker"] = "not found"
    _gap("docker not on PATH; install Docker first")
except Exception as e:
    result["host"]["docker"] = f"error: {e}"
    _gap(f"docker probe failed: {e}")

# nvidia-smi (host driver)
try:
    r = subprocess.run(
        ["nvidia-smi", "--query-gpu=name,memory.total,compute_cap,driver_version,uuid",
         "--format=csv,noheader,nounits"],
        capture_output=True, text=True, timeout=10,
    )
    if r.returncode == 0:
        result["host"]["nvidia_smi"] = "ok"
        for line in r.stdout.strip().splitlines():
            parts = [p.strip() for p in line.split(",")]
            if len(parts) >= 5:
                try:
                    vram_mb = int(parts[1])
                except ValueError:
                    vram_mb = None
                result["gpus"].append({
                    "name": parts[0],
                    "vram_mb": vram_mb,
                    "compute_cap": parts[2],
                    "driver": parts[3],
                    "uuid": parts[4],
                })
        if not result["gpus"]:
            _gap("nvidia-smi succeeded but reported no GPUs")
    else:
        result["host"]["nvidia_smi"] = "failed"
        _gap("nvidia-smi found but failed; is the NVIDIA driver loaded?")
except FileNotFoundError:
    result["host"]["nvidia_smi"] = "not found"
    _gap("nvidia-smi not on PATH; install the NVIDIA driver")
except Exception as e:
    result["host"]["nvidia_smi"] = f"error: {e}"
    _gap(f"nvidia-smi probe failed: {e}")

# disk free at the repo location
try:
    free_bytes = shutil.disk_usage(repo).free
    free_gb = free_bytes // (1024**3)
    result["disk"]["free_gb"] = free_gb
    if free_gb < result["disk"]["min_gb"]:
        _gap(f"disk free at {repo} is {free_gb} GB; need at least {result['disk']['min_gb']} GB for the kermt image")
except Exception as e:
    _gap(f"could not check disk space at {repo}: {e}")

# image presence (informational only)
try:
    r = subprocess.run(["docker", "image", "inspect", image], capture_output=True, text=True, timeout=10)
    result["image"]["present_locally"] = (r.returncode == 0)
except Exception:
    result["image"]["present_locally"] = None

# nvidia-container-toolkit probe — only meaningful if both docker and a
# locally-present image are available. Pick kermt:$tag first; fall back to
# the small CUDA base image if that's the only one present; otherwise skip
# (avoid pulling anything).
def _probe_image():
    for img in (image, "nvidia/cuda:12.6.3-base-ubuntu22.04"):
        r = subprocess.run(["docker", "image", "inspect", img], capture_output=True)
        if r.returncode == 0:
            return img
    return None

probe_img = _probe_image()
if probe_img:
    try:
        r = subprocess.run(
            ["docker", "run", "--rm", "--gpus", "all", probe_img, "nvidia-smi"],
            capture_output=True, text=True, timeout=60,
        )
        if r.returncode == 0:
            result["host"]["container_toolkit"] = f"ok (probed via {probe_img})"
        else:
            result["host"]["container_toolkit"] = f"failed (probed via {probe_img})"
            _gap("`docker run --gpus all` failed; install nvidia-container-toolkit and ensure the host driver supports it")
    except Exception as e:
        result["host"]["container_toolkit"] = f"error: {e}"
        _gap(f"nvidia-container-toolkit probe failed: {e}")
else:
    result["host"]["container_toolkit"] = "skipped (no probe image present locally; run ensure_image first)"

print(json.dumps(result, indent=2))
PYEOF
}

kermt_check_gpu() {
  # Probes whether `docker --gpus all` is wired up (nvidia-container-toolkit).
  # Image-selection priority (never pulls anything):
  #   1) $KERMT_IMAGE if it exists locally,
  #   2) else nvidia/cuda:12.6.3-base-ubuntu22.04 if it exists locally,
  #   3) else skip with a warning (return 0). The smoke test inside kermt_run
  #      will catch broken GPU passthrough later anyway.
  local probe_img=""
  if docker image inspect "$KERMT_IMAGE" >/dev/null 2>&1; then
    probe_img="$KERMT_IMAGE"
  elif docker image inspect nvidia/cuda:12.6.3-base-ubuntu22.04 >/dev/null 2>&1; then
    probe_img="nvidia/cuda:12.6.3-base-ubuntu22.04"
  else
    echo "[kermt] check_gpu: skipped — neither '$KERMT_IMAGE' nor 'nvidia/cuda:12.6.3-base-ubuntu22.04' is present locally. Run 'ensure_image' first, or this probe will be exercised by the in-container smoke test." >&2
    return 0
  fi
  if ! docker run --rm --gpus all "$probe_img" nvidia-smi >/dev/null 2>&1; then
    echo "[kermt] error: 'docker run --gpus all' failed (probe image: $probe_img). Install nvidia-container-toolkit and ensure the host has a CUDA-capable NVIDIA driver." >&2
    return 1
  fi
}

# -----------------------------------------------------------------------------
# Image build / verification
# -----------------------------------------------------------------------------

kermt_ensure_image() {
  _kermt_require_repo || return $?
  kermt_check_docker || return $?
  if docker image inspect "$KERMT_IMAGE" >/dev/null 2>&1; then
    local id
    id=$(docker image inspect "$KERMT_IMAGE" --format '{{.Id}}' 2>/dev/null | cut -c1-19)
    echo "[kermt] image '$KERMT_IMAGE' already present (${id:-unknown})"
    return 0
  fi
  echo "[kermt] image '$KERMT_IMAGE' not found; building from $KERMT_REPO/Dockerfile"
  echo "[kermt] first build typically takes 10-20 minutes on a typical workstation; subsequent runs reuse the cached image"
  docker build -t "$KERMT_IMAGE" -f "$KERMT_REPO/Dockerfile" "$KERMT_REPO"
}

# -----------------------------------------------------------------------------
# Mount-flag parser, internal
# -----------------------------------------------------------------------------
# Reads flags from the caller's positional args until it hits '--', appending
# `-v src:dst[:ro]` pairs into the caller-provided array name (passed as $1).
# Returns the number of caller-provided args consumed via _kermt_consumed.
# This is bash-specific (uses nameref via `declare -n`).

_kermt_parse_mounts() {
  local -n _out="$1"
  shift
  _kermt_consumed=0
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --)
        return 0
        ;;
      --data)
        [[ -e "$2" ]] || { echo "[kermt] --data path not found: $2" >&2; return 1; }
        # If the user passes a file, mount its parent directory at /data so
        # downstream commands can refer to /data/<basename>. Mounting a
        # single file at /data makes the path-as-directory pattern in the
        # skill examples (`--csv /data/<basename>`) fail with "not found".
        if [[ -d "$2" ]]; then
          _out+=("-v" "$(realpath "$2"):/data:ro")
        else
          _out+=("-v" "$(realpath "$(dirname "$2")"):/data:ro")
        fi
        shift 2; _kermt_consumed=$((_kermt_consumed + 2))
        ;;
      --ckpt)
        [[ -e "$2" ]] || { echo "[kermt] --ckpt path not found: $2" >&2; return 1; }
        _out+=("-v" "$(realpath "$2"):/ckpt:ro")
        shift 2; _kermt_consumed=$((_kermt_consumed + 2))
        ;;
      --vocab-dir)
        [[ -d "$2" ]] || { echo "[kermt] --vocab-dir not found or not a directory: $2" >&2; return 1; }
        _out+=("-v" "$(realpath "$2"):/vocab:ro")
        shift 2; _kermt_consumed=$((_kermt_consumed + 2))
        ;;
      --run-dir)
        mkdir -p "$2" || { echo "[kermt] failed to create --run-dir: $2" >&2; return 1; }
        _out+=("-v" "$(realpath "$2"):/runs")
        shift 2; _kermt_consumed=$((_kermt_consumed + 2))
        ;;
      --model-dir)
        mkdir -p "$2" || { echo "[kermt] failed to create --model-dir: $2" >&2; return 1; }
        _out+=("-v" "$(realpath "$2"):/model")
        shift 2; _kermt_consumed=$((_kermt_consumed + 2))
        ;;
      *)
        return 0
        ;;
    esac
  done
}

# -----------------------------------------------------------------------------
# Foreground / detached run
# -----------------------------------------------------------------------------

# Capture host-side git state for the repo and emit `-e KERMT_REPO_COMMIT=…
# -e KERMT_REPO_DIRTY=true|false` flags. Used by the run / run_detached
# wrappers so the runner's run.json manifest gets honest commit info even
# though `git -C /workspace` inside the container fails due to bind-mount
# ownership.
_kermt_git_env_flags() {
  local commit="unknown"
  local dirty="false"
  if command -v git >/dev/null 2>&1 && [[ -d "$KERMT_REPO/.git" ]]; then
    local c
    c=$(git -C "$KERMT_REPO" rev-parse HEAD 2>/dev/null) && commit="$c"
    # `--untracked-files=no` filters out user-private notes (e.g. a CLAUDE.md
    # or RELEASE_PLAN_v2.0.md at the repo root) that wouldn't affect
    # reproducibility — only modifications to tracked files do.
    if [[ -n "$(git -C "$KERMT_REPO" status --porcelain --untracked-files=no 2>/dev/null | head -n 1)" ]]; then
      dirty="true"
    fi
  fi
  printf '%s\n%s\n%s\n%s\n' "-e" "KERMT_REPO_COMMIT=$commit" "-e" "KERMT_REPO_DIRTY=$dirty"
}

# Forward HF_TOKEN into the container when it is set, so fetch_released_model.py
# can authenticate to Hugging Face. The current release is public (no token
# needed); this only guards against shared-IP rate limits or a future gated
# repo. Emits nothing when HF_TOKEN is unset.
_kermt_hf_env_flags() {
  if [[ -n "${HF_TOKEN:-}" ]]; then
    printf '%s\n%s\n' "-e" "HF_TOKEN=$HF_TOKEN"
  fi
}

kermt_run() {
  kermt_ensure_image || return $?
  local mount_args=()
  _kermt_parse_mounts mount_args "$@" || return $?
  shift "$_kermt_consumed"
  if [[ "${1:-}" != "--" ]]; then
    echo "[kermt] expected '--' separating mount flags from the command (got '${1:-}')" >&2
    return 1
  fi
  shift
  if [[ $# -eq 0 ]]; then
    echo "[kermt] no command supplied after '--'" >&2
    return 1
  fi
  local git_args=()
  while IFS= read -r line; do git_args+=("$line"); done < <(_kermt_git_env_flags)
  local hf_args=()
  while IFS= read -r line; do hf_args+=("$line"); done < <(_kermt_hf_env_flags)
  docker run --rm --gpus "$KERMT_GPUS" \
    --user "$(id -u):$(id -g)" \
    -v "$KERMT_REPO:/workspace" \
    -v "$_kermt_bundle_dir:/skill:ro" \
    "${mount_args[@]}" \
    -w /workspace \
    -e KERMT_REPO=/workspace \
    -e PYTHONPATH=/workspace \
    -e HOME=/tmp/kermt-home \
    "${git_args[@]}" \
    "${hf_args[@]}" \
    "$KERMT_IMAGE" \
    conda run -n kermt --no-capture-output bash -c "$*"
}

kermt_run_detached() {
  kermt_ensure_image || return $?
  local name=""
  local mount_args=()
  # Pull --name out first, then let the shared mount parser handle the rest.
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name) name="$2"; shift 2 ;;
      --) break ;;
      --data|--ckpt|--vocab-dir|--run-dir|--model-dir) break ;;
      *) break ;;
    esac
  done
  _kermt_parse_mounts mount_args "$@" || return $?
  shift "$_kermt_consumed"
  if [[ "${1:-}" != "--" ]]; then
    echo "[kermt] expected '--' separating mount flags from the command (got '${1:-}')" >&2
    return 1
  fi
  shift
  if [[ $# -eq 0 ]]; then
    echo "[kermt] no command supplied after '--'" >&2
    return 1
  fi
  if [[ -z "$name" ]]; then
    name="kermt-$(date -u +%Y%m%dT%H%M%SZ)-$$"
  fi
  local cid
  local git_args=()
  while IFS= read -r line; do git_args+=("$line"); done < <(_kermt_git_env_flags)
  local hf_args=()
  while IFS= read -r line; do hf_args+=("$line"); done < <(_kermt_hf_env_flags)
  cid=$(docker run -d --gpus "$KERMT_GPUS" \
    --user "$(id -u):$(id -g)" \
    --name "$name" \
    -v "$KERMT_REPO:/workspace" \
    -v "$_kermt_bundle_dir:/skill:ro" \
    "${mount_args[@]}" \
    -w /workspace \
    -e KERMT_REPO=/workspace \
    -e PYTHONPATH=/workspace \
    -e HOME=/tmp/kermt-home \
    "${git_args[@]}" \
    "${hf_args[@]}" \
    "$KERMT_IMAGE" \
    conda run -n kermt --no-capture-output bash -c "$*") || return $?
  echo "[kermt] container started: name=$name id=$cid"
  echo "[kermt] follow logs:    docker logs -f $name"
  echo "[kermt] wait for exit:  docker wait $name"
  echo "[kermt] stop:           docker stop $name"
  echo "$cid"
}

# -----------------------------------------------------------------------------
# Subcommand dispatch when invoked directly (not sourced)
# -----------------------------------------------------------------------------

if [[ "${BASH_SOURCE[0]:-$0}" == "${0}" ]]; then
  cmd="${1:-}"; shift || true
  case "$cmd" in
    check_docker)  kermt_check_docker "$@" ;;
    check_gpu)     kermt_check_gpu "$@" ;;
    check_system)  kermt_check_system "$@" ;;
    ensure_image)  kermt_ensure_image "$@" ;;
    run)           kermt_run "$@" ;;
    run_detached)  kermt_run_detached "$@" ;;
    ""|-h|--help)
      cat >&2 <<EOF
usage: $0 <subcommand> [args...]

Subcommands:
  check_docker        Verify docker is installed and the daemon is reachable.
  check_gpu           Verify 'docker --gpus all' works (nvidia-container-toolkit).
  check_system        Emit a JSON probe of host GPU + VRAM + compute_cap +
                      driver + disk space + container toolkit + image presence.
                      Exits 0 with ok=false + a 'gaps' list when anything's
                      below the per-workflow minimum.
  ensure_image        Build kermt:latest from \$KERMT_REPO/Dockerfile if missing.
  run [flags] -- ...  Run a command inside the container (foreground, --rm).
  run_detached [flags] -- ...
                      Run detached; prints container name + id + log hint.

Mount flags (for run / run_detached):
  --data <path>       bind to /data    (read-only)
  --ckpt <path>       bind to /ckpt    (read-only)
  --vocab-dir <dir>   bind to /vocab   (read-only)
  --run-dir <dir>     bind to /runs    (read-write; created on host if missing)
  --model-dir <dir>   bind to /model   (read-write; released-model download target)

Additional flags for run_detached:
  --name <name>       container name (default: kermt-<timestamp>-<pid>)

Environment overrides:
  KERMT_IMAGE         default kermt:latest
  KERMT_REPO          checkout path; otherwise discovered above the skill or working directory
  KERMT_GPUS          default all
EOF
      exit 1
      ;;
    *)
      echo "[kermt] unknown subcommand: $cmd" >&2
      echo "[kermt] run '$0 --help' for usage" >&2
      exit 1
      ;;
  esac
fi

#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Mode-dispatched data preparation pipeline for the KERMT agent skills.

Composes the existing repo data-prep scripts (`scripts/clean_smiles.py`,
`scripts/save_features.py`, `scripts/build_vocab.py`, `scripts/split_data.py`)
into a single one-call entry point per workflow. Output lands in `--out` with
a `prepare_data.json` manifest that the downstream runners read.

Mode pipelines
--------------
pretrain  : clean -> (optional auto-split train into train+val by --val-frac)
            -> save_features (fgtasklabel) on each CSV
            -> vocab step: if --vocab-dir / --{atom,bond,smiles}-vocab given,
               copy those through (continue-pretrain case — the ckpt's vocab
               is authoritative); else if --skip-vocab, skip;
               else build_vocab on train (pretrain-from-scratch case)
            -> split_data (graph + feature shards + summary.txt) per CSV
finetune  : clean each provided CSV -> (optional random split when only one
            CSV is provided; emits a strong warning recommending scaffold-
            balanced pre-splits) -> save_features (rdkit_2d_normalized) per CSV
inference : clean -> save_features (rdkit_2d_normalized)
embed     : clean only (extract_embeddings.py featurizes on the fly)

Output convention
-----------------
The manifest under `<out>/prepare_data.json` captures every step's inputs,
outputs, duration, and skipped-due-to-existing flag, plus a top-level
`split_method` field (one of: "user_provided", "random", "n/a") that the
finetune runner uses to pass the correct `--split_type` to main.py.

Subprocess composition
----------------------
Each underlying script is invoked via `subprocess.run`. The PYTHONPATH=/workspace
env var (set by `scripts/kermt_container.sh`) makes the `kermt` package
importable inside the subprocesses; without it, build_vocab.py and split_data.py
fail with `ModuleNotFoundError: No module named 'kermt'`.

CLI
---
    prepare_data.py --mode {pretrain|finetune|inference|embed}
                    --csv <input.csv> --out <output-dir>
                    [--val-csv <path>] [--test-csv <path>]
                    [--val-frac 0.1] [--test-frac 0.1] [--seed 0]
                    [--sample-per-file 100000] [--vocab-format json]
                    [--dataset-name pretrain]
                    [--targets COL [COL ...]]
                    [--features-generator <name>]
                    [--smiles-column 0]
                    [--force] [--skip-clean] [--skip-features]
                    [--skip-vocab] [--skip-split]
"""
from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
import time
import traceback
from pathlib import Path
from typing import Any

import pandas as pd

# sys.path tweak so `_utils` is importable regardless of how this script
# is invoked (kermt_run sets PYTHONPATH=/workspace; bare-Python launches
# from the host don't).
if str(Path(__file__).resolve().parent) not in sys.path:
    sys.path.insert(0, str(Path(__file__).resolve().parent))
from _utils import PRETRAIN_VOCAB_STEMS, resolve_kermt_repo, validate_vocab_file  # noqa: E402


REPO_ROOT = resolve_kermt_repo()
EXISTING_SCRIPTS = REPO_ROOT / "scripts"

DEFAULT_FEATURES_GENERATOR = {
    "pretrain": "fgtasklabel",
    "finetune": "rdkit_2d_normalized",
    "inference": "rdkit_2d_normalized",
    "embed": None,  # not used
}

VALID_MODES = ("pretrain", "finetune", "inference", "embed")


# ---------------------------------------------------------------------------
# Subprocess helpers
# ---------------------------------------------------------------------------

def _run(cmd: list[str], step_name: str, manifest: dict[str, Any]) -> dict[str, Any]:
    """Run a subprocess, append a step entry to manifest, raise on failure."""
    step: dict[str, Any] = {
        "name": step_name,
        "cmd": cmd,
        "duration_s": None,
        "ok": False,
        "stderr_tail": "",
        "skipped_due_to_existing": False,
    }
    t0 = time.time()
    proc = subprocess.run(cmd, capture_output=True, text=True)
    step["duration_s"] = round(time.time() - t0, 2)
    if proc.returncode != 0:
        step["stderr_tail"] = (proc.stderr or "").splitlines()[-20:]
        step["ok"] = False
        manifest["steps"].append(step)
        raise RuntimeError(
            f"step '{step_name}' failed (exit {proc.returncode}); "
            f"command: {' '.join(cmd)}\nstderr tail:\n" + "\n".join(step["stderr_tail"])
        )
    step["ok"] = True
    manifest["steps"].append(step)
    return step


def _skipped(step_name: str, output_path: str, manifest: dict[str, Any]) -> dict[str, Any]:
    step = {
        "name": step_name,
        "output": output_path,
        "ok": True,
        "duration_s": 0.0,
        "skipped_due_to_existing": True,
    }
    manifest["steps"].append(step)
    return step


def _exists_nonempty(path: Path) -> bool:
    """File exists with non-zero size, or directory exists with at least one entry."""
    if not path.exists():
        return False
    if path.is_file():
        return path.stat().st_size > 0
    if path.is_dir():
        try:
            next(path.iterdir())
            return True
        except StopIteration:
            return False
    return False


# ---------------------------------------------------------------------------
# Per-script wrappers
# ---------------------------------------------------------------------------

def _resolve_smiles_column(csv_path: Path, explicit_value: int | None) -> int:
    """Return the 0-based index of the SMILES column in csv_path.

    Auto-detection rule when `explicit_value is None`:
      1. Read the CSV header (first non-empty row).
      2. Prefer an exact lowercase `smiles` column (kermt convention).
      3. Otherwise accept a single case-insensitive match
         (`SMILES`, `Smiles`, etc.).
      4. If no match (or multiple ambiguous matches), raise a ValueError
         that surfaces the header so the user can disambiguate via
         `--smiles-column N`.

    Real datasets routinely place SMILES at column index ≠ 0
    (e.g. openadmet's all.csv has "Molecule Name" at col 0 and "SMILES"
    at col 1). Auto-detection prevents the silent 0-row-clean failure
    mode where every row gets rejected because col 0 doesn't parse as
    a SMILES string.
    """
    if explicit_value is not None:
        return explicit_value

    if not csv_path.is_file():
        raise ValueError(f"input CSV not found: {csv_path}")

    import csv as _csv
    with csv_path.open("r", newline="") as f:
        reader = _csv.reader(f)
        try:
            header = next(reader)
        except StopIteration:
            raise ValueError(f"input CSV {csv_path} is empty")

    stripped = [c.strip() for c in header]
    # Prefer exact lowercase "smiles"
    exact = [i for i, c in enumerate(stripped) if c == "smiles"]
    if exact:
        return exact[0]
    # Then case-insensitive
    ci = [i for i, c in enumerate(stripped) if c.lower() == "smiles"]
    if len(ci) == 1:
        return ci[0]
    if len(ci) > 1:
        raise ValueError(
            f"input CSV {csv_path} has multiple SMILES-named columns: "
            f"{[header[i] for i in ci]} at indices {ci}. "
            "Pass --smiles-column N (0-based) to disambiguate."
        )
    raise ValueError(
        f"could not auto-detect a SMILES column in {csv_path}. "
        f"Header columns: {header}. "
        "Pass --smiles-column N (0-based) to specify which column holds SMILES."
    )


def _clean_smiles(
    input_csv: Path, output_csv: Path, smiles_column: int, manifest: dict[str, Any], force: bool
) -> Path:
    if not force and _exists_nonempty(output_csv):
        _skipped(f"clean_smiles({input_csv.name})", str(output_csv), manifest)
        return output_csv
    output_csv.parent.mkdir(parents=True, exist_ok=True)
    if force and output_csv.exists():
        # clean_smiles.py prompts interactively (input()) when the output file
        # already exists — that's an EOFError in a non-TTY subprocess. Pre-delete.
        output_csv.unlink()
    cmd = [
        sys.executable, str(EXISTING_SCRIPTS / "clean_smiles.py"),
        "--input", str(input_csv),
        "--output", str(output_csv),
        "--smiles_column", str(smiles_column),
    ]
    _run(cmd, f"clean_smiles({input_csv.name})", manifest)
    return output_csv


def _reduce_to_smiles_column(
    csv_path: Path, smiles_column: int, manifest: dict[str, Any]
) -> Path:
    """Rewrite an inference CSV to keep only the SMILES column (at index 0).

    Downstream `kermt.util.utils.get_data` -> `MoleculeDatapoint.__init__`
    floats every column after SMILES, which crashes on non-numeric passthrough
    columns (e.g. a 'split' label of 'train'/'val'/'test', or a 'Molecule Name'
    string). Inference does not need target columns, so drop them here.

    Note on skip semantics: this step is idempotent — running it on an
    already-single-column file is a no-op. We record that with
    `skipped_due_to_idempotent: True`, NOT `skipped_due_to_existing: True`.
    The two fields have different meanings: `_existing` means "I found a
    cached output file from a prior run and reused it" (overridden by
    `--force`); `_idempotent` means "the input is already in the desired
    state, so re-executing changes nothing" (safe to skip even under
    `--force`).
    """
    step_name = f"reduce_to_smiles_only({csv_path.name})"
    start = time.time()
    df = pd.read_csv(csv_path)
    if df.shape[1] == 1:
        manifest["steps"].append({
            "name": step_name,
            "output": str(csv_path),
            "ok": True,
            "duration_s": time.time() - start,
            "skipped_due_to_idempotent": True,
            "note": "already single-column",
        })
        return csv_path
    effective_col = smiles_column if 0 <= smiles_column < df.shape[1] else 0
    df.iloc[:, [effective_col]].to_csv(csv_path, index=False)
    manifest["steps"].append({
        "name": step_name,
        "output": str(csv_path),
        "ok": True,
        "duration_s": time.time() - start,
        "input_cols": int(df.shape[1]),
        "kept_col": effective_col,
        "kept_col_name": str(df.columns[effective_col]),
    })
    return csv_path


def _save_features(
    csv_path: Path, npz_path: Path, generator: str, manifest: dict[str, Any], force: bool
) -> Path:
    if not force and _exists_nonempty(npz_path):
        _skipped(f"save_features({csv_path.name}, {generator})", str(npz_path), manifest)
        return npz_path
    npz_path.parent.mkdir(parents=True, exist_ok=True)
    if force and npz_path.exists():
        npz_path.unlink()  # --restart still loads partial state if file exists; pre-delete to be safe
    cmd = [
        sys.executable, str(EXISTING_SCRIPTS / "save_features.py"),
        "--data_path", str(csv_path),
        "--save_path", str(npz_path),
        "--features_generator", generator,
        "--restart",
    ]
    _run(cmd, f"save_features({csv_path.name}, {generator})", manifest)
    return npz_path


def _resolve_vocab_inputs(args: argparse.Namespace) -> dict[str, Path | None] | None:
    """Returns {atom, bond, smiles}->Path|None when the user supplied vocab
    inputs (via --vocab-dir or --atom-vocab/--bond-vocab/--smiles-vocab),
    else None (signal to fall through to build_vocab).

    Conventional filenames inside --vocab-dir:
      pretrain_atom_vocab.{json,pkl}
      pretrain_bond_vocab.{json,pkl}
      pretrain_smiles_vocab.pkl
    """
    if args.vocab_dir:
        d = Path(args.vocab_dir).resolve()
        if not d.is_dir():
            raise FileNotFoundError(f"--vocab-dir not found or not a directory: {d}")
        def _find(stem: str, exts: tuple[str, ...]) -> Path | None:
            for ext in exts:
                p = d / f"{stem}.{ext}"
                if p.is_file():
                    return p
            return None
        atom = _find(PRETRAIN_VOCAB_STEMS["atom"], ("json", "pkl"))
        bond = _find(PRETRAIN_VOCAB_STEMS["bond"], ("json", "pkl"))
        smiles = _find(PRETRAIN_VOCAB_STEMS["smiles"], ("pkl",))
        if atom is None and bond is None and smiles is None:
            stems = [PRETRAIN_VOCAB_STEMS[k] for k in ("atom", "bond", "smiles")]
            raise FileNotFoundError(
                f"--vocab-dir {d} contained no {{ {', '.join(stems) }}}.{{json,pkl}} "
                f"files. Expected at least {PRETRAIN_VOCAB_STEMS['atom']} + "
                f"{PRETRAIN_VOCAB_STEMS['bond']}."
            )
        return {"atom": atom, "bond": bond, "smiles": smiles}

    if args.atom_vocab or args.bond_vocab or args.smiles_vocab:
        return {
            "atom":   Path(args.atom_vocab).resolve()   if args.atom_vocab   else None,
            "bond":   Path(args.bond_vocab).resolve()   if args.bond_vocab   else None,
            "smiles": Path(args.smiles_vocab).resolve() if args.smiles_vocab else None,
        }

    return None


def _copy_provided_vocab(
    src: dict[str, Path | None], dst_dir: Path, dataset_name: str, manifest: dict[str, Any],
    force: bool,
) -> dict[str, Path]:
    """When the user supplies vocab files (use ckpt's vocab as-is),
    copy them into `<dst_dir>/<dataset_name>_<which>_vocab.<ext>` so the
    downstream pretrain command sees the conventional filenames.

    `src` is `{atom: Path|None, bond: Path|None, smiles: Path|None}`. The atom
    and bond entries must be both present or both absent (paired). smiles is
    optional (cmim/hybrid only).

    Returns the same dict of (resolved) destination paths.
    """
    import shutil
    if (src["atom"] is None) != (src["bond"] is None):
        raise ValueError(
            "vocab pass-through requires atom and bond vocab paths to be paired; "
            "got atom=" + str(src["atom"]) + ", bond=" + str(src["bond"])
        )
    out: dict[str, Path] = {}
    dst_dir.mkdir(parents=True, exist_ok=True)
    for which, path in src.items():
        if path is None:
            continue
        # Validate the source file IS a loadable KERMT vocab before copying.
        # Catches the "user pointed --smiles-vocab at a random pickle" case
        # early, with a clear error, instead of letting it surface as a cryptic
        # SMILESVocab.load_vocab failure at pretrain_ddp.py launch time.
        validate_vocab_file(path, kind=which)
        ext = path.suffix.lstrip(".")
        if which == "smiles":
            ext = "pkl"  # smiles vocab is always pickle
        dst = dst_dir / f"{dataset_name}_{which}_vocab.{ext}"
        if not force and _exists_nonempty(dst):
            _skipped(f"copy_vocab({which})", str(dst), manifest)
            out[which] = dst
            continue
        if force and dst.exists():
            dst.unlink()
        shutil.copy2(path, dst)
        manifest["steps"].append({
            "name": f"copy_vocab({which})",
            "src": str(path), "dst": str(dst), "ok": True,
            "duration_s": 0.0, "skipped_due_to_existing": False,
        })
        out[which] = dst
    return out


def _build_vocab(
    csv_path: Path, vocab_dir: Path, dataset_name: str, vocab_format: str,
    manifest: dict[str, Any], force: bool,
) -> dict[str, Path]:
    """Builds atom + bond (in --vocab-format) and smiles (always pickle) vocabs.
    Returns a dict of {atom, bond, smiles} -> Path."""
    suffix = "json" if vocab_format == "json" else "pkl"
    expected = {
        "atom": vocab_dir / f"{dataset_name}_atom_vocab.{suffix}",
        "bond": vocab_dir / f"{dataset_name}_bond_vocab.{suffix}",
        "smiles": vocab_dir / f"{dataset_name}_smiles_vocab.pkl",
    }
    if not force and all(_exists_nonempty(p) for p in expected.values()):
        _skipped(f"build_vocab({csv_path.name})", str(vocab_dir), manifest)
        return expected
    vocab_dir.mkdir(parents=True, exist_ok=True)
    if force:
        for p in expected.values():
            if p.exists():
                p.unlink()
    cmd = [
        sys.executable, str(EXISTING_SCRIPTS / "build_vocab.py"),
        "--data_path", str(csv_path),
        "--vocab_save_folder", str(vocab_dir),
        "--dataset_name", dataset_name,
        "--vocab_format", vocab_format,
    ]
    _run(cmd, f"build_vocab({csv_path.name})", manifest)
    return expected


def _split_data(
    csv_path: Path, features_path: Path | None, sample_per_file: int, output_dir: Path,
    manifest: dict[str, Any], force: bool,
) -> Path:
    """Run split_data.py to produce shard dirs (graph/ + optionally feature/ + summary.txt)."""
    summary = output_dir / "summary.txt"
    if not force and _exists_nonempty(summary):
        _skipped(f"split_data({csv_path.name})", str(output_dir), manifest)
        return output_dir
    if force and output_dir.exists():
        shutil.rmtree(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    cmd = [
        sys.executable, str(EXISTING_SCRIPTS / "split_data.py"),
        "--data_path", str(csv_path),
        "--sample_per_file", str(sample_per_file),
        "--output_path", str(output_dir),
    ]
    if features_path is not None:
        cmd += ["--features_path", str(features_path)]
    _run(cmd, f"split_data({csv_path.name})", manifest)
    return output_dir


# ---------------------------------------------------------------------------
# Random splitter (used only when the user supplies a single CSV)
# ---------------------------------------------------------------------------

def _random_split_csv(
    src_csv: Path, dst_csvs: dict[str, Path], fractions: dict[str, float], seed: int,
    manifest: dict[str, Any], force: bool,
) -> None:
    """Shuffle src_csv and partition rows into dst_csvs by fractions.
    `dst_csvs` and `fractions` are dicts keyed by the split name (e.g. 'train', 'val').
    Sum of fractions must be 1.0 (within float tolerance). Writes each dst_csv with the
    same header as the input."""
    step = {
        "name": f"random_split({src_csv.name})",
        "seed": seed,
        "fractions": fractions,
        "ok": False,
        "duration_s": None,
        "skipped_due_to_existing": False,
        "row_counts": {},
    }
    if not force and all(_exists_nonempty(p) for p in dst_csvs.values()):
        step["skipped_due_to_existing"] = True
        step["ok"] = True
        manifest["steps"].append(step)
        return

    if abs(sum(fractions.values()) - 1.0) > 1e-6:
        raise ValueError(f"split fractions must sum to 1.0 (got {sum(fractions.values())})")

    t0 = time.time()
    df = pd.read_csv(src_csv).sample(frac=1.0, random_state=seed).reset_index(drop=True)
    n = len(df)
    sizes: dict[str, int] = {}
    remaining = n
    split_names = list(fractions.keys())
    for name in split_names[:-1]:
        sizes[name] = int(round(fractions[name] * n))
        remaining -= sizes[name]
    sizes[split_names[-1]] = remaining

    start = 0
    for name in split_names:
        dst = dst_csvs[name]
        dst.parent.mkdir(parents=True, exist_ok=True)
        df.iloc[start:start + sizes[name]].to_csv(dst, index=False)
        step["row_counts"][name] = sizes[name]
        start += sizes[name]

    step["duration_s"] = round(time.time() - t0, 2)
    step["ok"] = True
    manifest["steps"].append(step)


def _emit_random_split_warning(
    src_csv: Path, fractions: dict[str, float], seed: int, manifest: dict[str, Any]
) -> None:
    row_counts = manifest["steps"][-1].get("row_counts", {})
    n = sum(row_counts.values()) if row_counts else "?"
    lines = [
        f"WARNING: Auto-splitting {n} rows from {src_csv.name} into:",
    ]
    for name, frac in fractions.items():
        cnt = row_counts.get(name, "?")
        lines.append(f"  {name}: {cnt} rows ({frac * 100:.1f}%)")
    lines += [
        f"using random split with seed {seed}.",
        "",
        "This is a RANDOM split. For rigorous ADMET evaluation, scaffold-balanced",
        "(or other structure-aware) splits are strongly preferred — molecules with",
        "similar scaffolds can leak across splits and inflate apparent generalization.",
        "",
        "To use your own pre-computed splits instead, pass:",
        "    --train-csv <train.csv> --val-csv <val.csv> --test-csv <test.csv>",
        "",
        "To customize fractions:",
        "    --val-frac 0.15 --test-frac 0.15",
    ]
    warning = "\n".join(lines)
    print(warning, file=sys.stderr)
    manifest["warnings"].append(warning)


# ---------------------------------------------------------------------------
# Mode pipelines
# ---------------------------------------------------------------------------

def _prepare_embed(args, out: Path, manifest: dict[str, Any]) -> None:
    manifest["split_method"] = "n/a"
    if args.skip_clean:
        clean = Path(args.csv)
        manifest["steps"].append({"name": "clean_smiles", "skipped_by_flag": True, "ok": True})
    else:
        clean = _clean_smiles(Path(args.csv), out / "clean.csv", args.smiles_column, manifest, args.force)
    manifest["outputs"]["clean_csv"] = str(clean)


def _prepare_inference(args, out: Path, manifest: dict[str, Any]) -> None:
    manifest["split_method"] = "n/a"
    clean = _clean_smiles(Path(args.csv), out / "clean.csv", args.smiles_column, manifest, args.force)
    # Reduce to SMILES-only: downstream get_data/MoleculeDatapoint floats every
    # non-SMILES column, which crashes on non-numeric passthrough columns
    # (e.g. a 'split' label). Inference does not need target columns.
    _reduce_to_smiles_column(clean, args.smiles_column, manifest)
    manifest["outputs"]["clean_csv"] = str(clean)
    if args.skip_features:
        manifest["steps"].append({"name": "save_features", "skipped_by_flag": True, "ok": True})
        return
    generator = args.features_generator or DEFAULT_FEATURES_GENERATOR["inference"]
    npz = _save_features(clean, out / "clean.npz", generator, manifest, args.force)
    manifest["outputs"]["clean_npz"] = str(npz)


def _prepare_finetune(args, out: Path, manifest: dict[str, Any]) -> None:
    src_train = Path(args.csv)
    has_val = args.val_csv is not None
    has_test = args.test_csv is not None
    split_type = args.split_type

    if has_val and has_test:
        # User supplied explicit val + test CSVs: trust them, just clean + featurize.
        # split_type is irrelevant when val/test are given separately.
        manifest["split_method"] = "user_provided"
        clean_train = _clean_smiles(src_train, out / "clean_train.csv", args.smiles_column, manifest, args.force)
        clean_val = _clean_smiles(Path(args.val_csv), out / "clean_val.csv", args.smiles_column, manifest, args.force)
        clean_test = _clean_smiles(Path(args.test_csv), out / "clean_test.csv", args.smiles_column, manifest, args.force)
        manifest["outputs"]["clean_train_csv"] = str(clean_train)
        manifest["outputs"]["clean_val_csv"] = str(clean_val)
        manifest["outputs"]["clean_test_csv"] = str(clean_test)
        per_split = (("train", clean_train), ("val", clean_val), ("test", clean_test))
    elif has_val or has_test:
        raise ValueError(
            "for finetune mode, either provide BOTH --val-csv and --test-csv (user-provided splits) "
            "or NEITHER (run with --split-type {random|scaffold_balanced|index_predetermined}). "
            "Got one but not both."
        )
    elif split_type == "random":
        # Random auto-split — done here in prep so train.py gets ready-made CSVs.
        manifest["split_method"] = "random"
        manifest["split_seed"] = args.seed
        train_frac = max(0.0, 1.0 - args.val_frac - args.test_frac)
        manifest["split_fractions"] = {"train": train_frac, "val": args.val_frac, "test": args.test_frac}
        clean_full = _clean_smiles(src_train, out / "_clean_full.csv", args.smiles_column, manifest, args.force)
        dst = {
            "train": out / "clean_train.csv",
            "val": out / "clean_val.csv",
            "test": out / "clean_test.csv",
        }
        _random_split_csv(clean_full, dst, manifest["split_fractions"], args.seed, manifest, args.force)
        clean_train, clean_val, clean_test = dst["train"], dst["val"], dst["test"]
        manifest["outputs"]["clean_train_csv"] = str(clean_train)
        manifest["outputs"]["clean_val_csv"] = str(clean_val)
        manifest["outputs"]["clean_test_csv"] = str(clean_test)
        _emit_random_split_warning(src_train, manifest["split_fractions"], args.seed, manifest)
        per_split = (("train", clean_train), ("val", clean_val), ("test", clean_test))
    else:
        # Scaffold-balanced or index-predetermined: prep cleans + featurizes the full
        # CSV and defers actual splitting to task/train.py, which calls split_data
        # with the user-supplied seed and split_sizes.
        manifest["split_method"] = "deferred_to_runner"
        manifest["split_type"] = split_type
        manifest["split_seed"] = args.seed
        manifest["split_fractions"] = {
            "train": max(0.0, 1.0 - args.val_frac - args.test_frac),
            "val": args.val_frac,
            "test": args.test_frac,
        }
        clean_full = _clean_smiles(src_train, out / "clean_full.csv", args.smiles_column, manifest, args.force)
        manifest["outputs"]["clean_full_csv"] = str(clean_full)
        per_split = (("full", clean_full),)

    if args.skip_features:
        manifest["steps"].append({"name": "save_features", "skipped_by_flag": True, "ok": True})
        return
    generator = args.features_generator or DEFAULT_FEATURES_GENERATOR["finetune"]
    for split_name, csv in per_split:
        npz = _save_features(csv, csv.with_suffix(".npz"), generator, manifest, args.force)
        manifest["outputs"][f"clean_{split_name}_npz"] = str(npz)


def _prepare_pretrain(args, out: Path, manifest: dict[str, Any]) -> None:
    src_train = Path(args.csv)
    if args.val_csv is not None:
        manifest["split_method"] = "user_provided"
        clean_train = _clean_smiles(src_train, out / "clean_train.csv", args.smiles_column, manifest, args.force)
        clean_val = _clean_smiles(Path(args.val_csv), out / "clean_val.csv", args.smiles_column, manifest, args.force)
    else:
        manifest["split_method"] = "random"
        manifest["split_seed"] = args.seed
        train_frac = max(0.0, 1.0 - args.val_frac)
        manifest["split_fractions"] = {"train": train_frac, "val": args.val_frac}
        clean_full = _clean_smiles(src_train, out / "_clean_full.csv", args.smiles_column, manifest, args.force)
        dst = {"train": out / "clean_train.csv", "val": out / "clean_val.csv"}
        _random_split_csv(clean_full, dst, manifest["split_fractions"], args.seed, manifest, args.force)
        clean_train, clean_val = dst["train"], dst["val"]

    manifest["outputs"]["clean_train_csv"] = str(clean_train)
    manifest["outputs"]["clean_val_csv"] = str(clean_val)

    generator = args.features_generator or DEFAULT_FEATURES_GENERATOR["pretrain"]
    if args.skip_features:
        manifest["steps"].append({"name": "save_features", "skipped_by_flag": True, "ok": True})
        train_npz: Path | None = None
        val_npz: Path | None = None
    else:
        train_npz = _save_features(clean_train, out / "clean_train.npz", generator, manifest, args.force)
        val_npz = _save_features(clean_val, out / "clean_val.npz", generator, manifest, args.force)
        manifest["outputs"]["clean_train_npz"] = str(train_npz)
        manifest["outputs"]["clean_val_npz"] = str(val_npz)

    if args.skip_vocab:
        manifest["steps"].append({"name": "build_vocab", "skipped_by_flag": True, "ok": True})
        manifest["vocab_source"] = "skipped"
    else:
        # Resolve user-provided vocab paths from --vocab-dir or explicit flags.
        provided = _resolve_vocab_inputs(args)
        if provided:
            # Use the user-supplied (ckpt's) vocab as-is. Copy into the
            # conventional filenames the downstream pretrain command expects.
            vocabs = _copy_provided_vocab(provided, out, args.dataset_name, manifest, args.force)
            manifest["vocab_source"] = "user_provided"
        else:
            # Fall back to the existing build-from-corpus behavior. Used by
            # pretrain-from-scratch and by any continue case where the user
            # explicitly wants a fresh vocab (rare, usually wrong).
            vocabs = _build_vocab(clean_train, out, args.dataset_name, args.vocab_format, manifest, args.force)
            manifest["vocab_source"] = "built_fresh"
        if "atom" in vocabs:
            manifest["outputs"]["atom_vocab"] = str(vocabs["atom"])
        if "bond" in vocabs:
            manifest["outputs"]["bond_vocab"] = str(vocabs["bond"])
        if "smiles" in vocabs:
            manifest["outputs"]["smiles_vocab"] = str(vocabs["smiles"])

    if args.skip_split:
        manifest["steps"].append({"name": "split_data", "skipped_by_flag": True, "ok": True})
    else:
        train_dir = _split_data(clean_train, train_npz, args.sample_per_file, out / "train", manifest, args.force)
        val_dir = _split_data(clean_val, val_npz, args.sample_per_file, out / "val", manifest, args.force)
        manifest["outputs"]["train_dir"] = str(train_dir)
        manifest["outputs"]["val_dir"] = str(val_dir)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def prepare(args: argparse.Namespace) -> dict[str, Any]:
    out = Path(args.out).resolve()
    out.mkdir(parents=True, exist_ok=True)
    manifest: dict[str, Any] = {
        "mode": args.mode,
        "input_csv": str(Path(args.csv).resolve()),
        "val_csv": str(Path(args.val_csv).resolve()) if args.val_csv else None,
        "test_csv": str(Path(args.test_csv).resolve()) if args.test_csv else None,
        "output_dir": str(out),
        "split_method": None,
        "steps": [],
        "outputs": {},
        "errors": [],
        "warnings": [],
    }
    try:
        if args.mode == "pretrain":
            _prepare_pretrain(args, out, manifest)
        elif args.mode == "finetune":
            _prepare_finetune(args, out, manifest)
        elif args.mode == "inference":
            _prepare_inference(args, out, manifest)
        elif args.mode == "embed":
            _prepare_embed(args, out, manifest)
        manifest["ok"] = True
    except Exception as exc:  # noqa: BLE001
        manifest["ok"] = False
        manifest["errors"].append(f"{type(exc).__name__}: {exc}")
    # Always write the manifest so partial-failure state is visible to the agent.
    (out / "prepare_data.json").write_text(json.dumps(manifest, indent=2))
    return manifest


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="Mode-dispatched data prep for the KERMT agent skills.")
    p.add_argument("--mode", required=True, choices=VALID_MODES)
    p.add_argument("--csv", required=True, help="Primary input CSV (train CSV for pretrain/finetune)")
    p.add_argument("--out", required=True, help="Output directory")
    p.add_argument("--val-csv", default=None, help="Optional separate val CSV (pretrain/finetune)")
    p.add_argument("--test-csv", default=None, help="Optional separate test CSV (finetune only)")
    p.add_argument("--val-frac", type=float, default=0.1, help="Auto-split val fraction (default 0.1)")
    p.add_argument("--test-frac", type=float, default=0.1, help="Auto-split test fraction (finetune only, default 0.1)")
    p.add_argument("--seed", type=int, default=0, help="Random split seed (default 0)")
    p.add_argument("--split-type", choices=["random", "scaffold_balanced", "index_predetermined"],
                   default="random",
                   help="(finetune only, when --val-csv/--test-csv are not given) how to split. "
                        "'random' splits in prep using --val-frac/--test-frac/--seed. "
                        "'scaffold_balanced' and 'index_predetermined' defer the actual split to the "
                        "runner (task/train.py invokes split_data with the appropriate algorithm "
                        "using the user-supplied seed); prep only cleans + featurizes the full CSV.")
    p.add_argument("--sample-per-file", type=int, default=100_000,
                   help="split_data shard size (pretrain only, default 100000)")
    p.add_argument("--vocab-format", choices=["json", "pkl"], default="json",
                   help="atom/bond vocab format (default json); smiles vocab is always pkl")
    # Vocab pass-through (pretrain mode): when continuing from a released ckpt,
    # pass its bundled vocab files in so we don't rebuild a mismatched vocab.
    p.add_argument("--vocab-dir", default=None,
                   help="(pretrain) directory containing pretrain_{atom,bond}_vocab.{json,pkl} "
                        "(+ pretrain_smiles_vocab.pkl for cmim/hybrid). When given, prepare_data "
                        "skips build_vocab and copies these files into the output dir under the "
                        "expected filenames. Used by kermt-continue-pretrain to bind the released "
                        "ckpt's vocab to the new corpus (the ckpt's vocab is authoritative).")
    p.add_argument("--atom-vocab", default=None,
                   help="(pretrain) explicit atom vocab path; pairs with --bond-vocab. Overrides "
                        "--vocab-dir's pretrain_atom_vocab.* discovery if both are given.")
    p.add_argument("--bond-vocab", default=None,
                   help="(pretrain) explicit bond vocab path; pairs with --atom-vocab.")
    p.add_argument("--smiles-vocab", default=None,
                   help="(pretrain, cmim/hybrid) explicit smiles vocab .pkl path. Optional for "
                        "vocab-only pretrain.")
    p.add_argument("--dataset-name", default="pretrain",
                   help="vocab filename prefix (default 'pretrain' so downstream pretrain commands "
                        "can reference pretrain_{atom,bond}_vocab.{json|pkl}, pretrain_smiles_vocab.pkl)")
    p.add_argument("--targets", nargs="+", default=None,
                   help="(finetune only) target column names; forwarded to the finetune runner via the manifest")
    p.add_argument("--features-generator", default=None,
                   help="Override the per-mode default (pretrain: fgtasklabel; finetune/inference: rdkit_2d_normalized)")
    p.add_argument("--smiles-column", type=int, default=None,
                   help="0-based column index of SMILES in the input CSV. "
                        "When omitted, auto-detected by header name "
                        "(prefers lowercase `smiles`; accepts case-insensitive "
                        "`SMILES`/`Smiles`). Pass explicitly to override.")
    p.add_argument("--force", action="store_true",
                   help="Re-run every step even if its outputs already exist")
    p.add_argument("--skip-clean", action="store_true", help="(embed mode) skip the cleaning step")
    p.add_argument("--skip-features", action="store_true", help="Skip feature generation")
    p.add_argument("--skip-vocab", action="store_true", help="(pretrain) skip vocab build")
    p.add_argument("--skip-split", action="store_true", help="(pretrain) skip shard split")
    args = p.parse_args(argv)

    # Forward --targets through the manifest so the finetune runner can see them.
    if args.mode == "finetune" and args.targets:
        pass  # captured in manifest below

    # Resolve the SMILES column index (auto-detect from header when the user
    # didn't pass --smiles-column). This is the only point where args.csv is
    # touched before downstream _clean_smiles calls fan it out.
    try:
        resolved_smiles_col = _resolve_smiles_column(Path(args.csv), args.smiles_column)
    except ValueError as exc:
        err_manifest = {
            "ok": False,
            "mode": args.mode,
            "errors": [f"smiles-column resolution failed: {exc}"],
        }
        Path(args.out).mkdir(parents=True, exist_ok=True)
        (Path(args.out) / "prepare_data.json").write_text(json.dumps(err_manifest, indent=2))
        print(json.dumps(err_manifest, indent=2))
        return 1
    if args.smiles_column is None:
        print(f"[prepare_data] auto-detected --smiles-column {resolved_smiles_col} "
              f"from {Path(args.csv).name} header", file=sys.stderr)
    args.smiles_column = resolved_smiles_col

    try:
        manifest = prepare(args)
    except Exception as exc:  # noqa: BLE001
        print(traceback.format_exc(), file=sys.stderr)
        print(json.dumps({"ok": False, "errors": [f"unhandled: {type(exc).__name__}: {exc}"]}, indent=2))
        return 1
    if args.targets:
        manifest["targets"] = list(args.targets)
    # Record the resolved SMILES column so the manifest is self-describing.
    manifest["smiles_column"] = args.smiles_column
    (Path(args.out) / "prepare_data.json").write_text(json.dumps(manifest, indent=2))
    print(json.dumps(manifest, indent=2))
    return 0 if manifest.get("ok") else 1


if __name__ == "__main__":
    sys.exit(main())

#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Workstation pretrain runner — composes prepare_data + ckpt-validator outputs
into a pretrain_ddp.py invocation.

Continues pretraining from a user-provided checkpoint. The model type
(grover_base / cmim / hybrid) is inferred from the validator's output and
drives the pretrain_ddp.py flag set; arch params come exclusively from the
ckpt; training/loss hyperparameters come from config/defaults_pretrain.json
with per-flag CLI overrides.

How it interacts with pretrain_ddp.py's auto-resume:
    pretrain_ddp.py looks at <save_dir>/last_checkpoint.pt and resumes from it
    if present. The runner sets `--save_dir <out>/ckpt` and symlinks the user's
    input ckpt to <out>/ckpt/last_checkpoint.pt so the resume path picks it up.

Run.json manifest:
    Records source-repo commit + image digest + a copy-pasteable `cmd_replay`
    + per-flag `args_applied` so the artifact is self-contained and replayable.

CLI
---
    run_pretrain_local.py
        --ckpt <path>                  # input pretrain ckpt (required)
        --prepare-manifest <path>      # prepare_data.json from a prior prepare run
        --out <run-dir>                # output dir (typically runs/continue-pretrain_<ts>/)
        [--ckpt-validator-out <path>]  # cached check_checkpoint.py JSON; computed if absent
        [--gpus 0,2]                   # subset of detected GPUs; default = all visible
        [--dry-run]                    # write run.json + print command, do not execute
        [--epochs N] [--batch-size N] [--init-lr F] [--max-lr F] [--final-lr F]
        [--warmup-epochs F] [--weight-decay F] [--dropout F]
        [--save-interval N] [--seed N]
        [--vocab-loss-weight F]            # hybrid only
        [--latent-dim N] [--contrastive-temperature F]   # cmim/hybrid only
"""
from __future__ import annotations

import argparse
import datetime
import json
import os
import subprocess
import sys
from pathlib import Path
from typing import Any

# Add the scripts/ dir to sys.path so `_utils` is importable whether
# this script is launched via `kermt_run` (PYTHONPATH=/workspace) or as a
# bare `python scripts/run_pretrain_local.py …` from the host.
if str(Path(__file__).resolve().parent) not in sys.path:
    sys.path.insert(0, str(Path(__file__).resolve().parent))
from _utils import (  # noqa: E402
    resolve_kermt_repo, assert_prepare_manifest_basics, count_vocab_entries, docker_image_digest,
    format_cmd_replay, git_commit_with_env_override, load_json, load_checkpoint,
    merge_default_into_applied, run_checkpoint_validator, runner_environment,
)


REPO_ROOT = resolve_kermt_repo()
SKILL_ROOT = Path(__file__).resolve().parent.parent
DEFAULTS_PATH = SKILL_ROOT / "config" / "defaults_pretrain.json"
CHECK_CHECKPOINT_PATH = SKILL_ROOT / "scripts" / "check_checkpoint.py"
PRETRAIN_DDP_PATH = REPO_ROOT / "pretrain_ddp.py"

# Model-type → pretrain_ddp.py `--pretrain_mode` value.
MODEL_TYPE_TO_PRETRAIN_MODE = {
    "grover_base": "vocab",
    "cmim": "cmim",
    "hybrid": "hybrid",
}

# Hyperparameter flags the runner exposes for CLI override + the corresponding
# key path in defaults_pretrain.json. None means the value isn't in defaults
# (e.g. seed has a default but lives at the top of training; lookup is direct).
TRAINING_FLAGS = (
    "batch_size", "dropout", "epochs", "init_lr", "max_lr", "final_lr",
    "warmup_epochs", "weight_decay", "save_interval", "seed", "tensorboard",
    "use_cuikmolmaker_featurization",
)
LOSS_FLAGS = ("contrastive_temperature", "vocab_loss_weight")
DECODER_FLAGS = (
    "latent_dim",
    "decoder_num_layers",
    "decoder_num_attention_heads",
    "decoder_ffn_hidden_size",
    "decoder_dropout",
    "decoder_max_seq_len",
    "decoder_positional_encoding",
    "decoder_gate_self_attn",
    "decoder_gate_cross_attn",
)

ARCH_FLAGS_FROM_CKPT = (
    "hidden_size", "depth", "num_attn_head", "activation", "backbone",
    "embedding_output_type", "self_attention",
)

# cMIM-decoder + latent-distribution arch fields. For continue-pretrain on a
# cmim/hybrid ckpt these MUST come from the ckpt's saved_args (so the model
# being constructed matches the ckpt's weights at load time); the
# defaults_pretrain.json `add_cmim_decoder` block is for add-cmim-pretrain's
# upgrade-time decoder construction only, and is intentionally ignored
# during continue-pretrain.
CMIM_DECODER_FLAGS_FROM_CKPT = (
    "latent_dim",
    "decoder_num_layers",
    "decoder_num_attention_heads",
    "decoder_ffn_hidden_size",
    "decoder_dropout",
    "decoder_max_seq_len",
    "decoder_positional_encoding",
    "decoder_gate_self_attn",
    "decoder_gate_cross_attn",
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

# JSON loading delegated to the shared _utils.load_json. Alias kept for the
# existing internal callsites that use the leading-underscore convention.
_load_json = load_json


def _detect_gpus(override: str | None) -> tuple[int, str]:
    """Returns (world_size, CUDA_VISIBLE_DEVICES_string)."""
    if override:
        gpu_list = [g.strip() for g in override.split(",") if g.strip()]
        return len(gpu_list), ",".join(gpu_list)
    # Honor an existing CUDA_VISIBLE_DEVICES in the environment.
    env = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
    if env:
        ids = [g for g in env.split(",") if g]
        return len(ids), ",".join(ids)
    try:
        import torch
        n = torch.cuda.device_count()
    except Exception:
        n = 0
    return n, ",".join(str(i) for i in range(n))


def _verify_prepare_manifest(manifest: dict[str, Any]) -> None:
    assert_prepare_manifest_basics(manifest, "pretrain")
    out = manifest.get("outputs", {})
    required_keys = ("train_dir", "val_dir", "atom_vocab", "bond_vocab")
    missing = [k for k in required_keys if k not in out]
    if missing:
        raise ValueError(
            f"prepare_data manifest is missing required outputs: {missing}. "
            "Was prepare_data.py invoked with --skip-vocab or --skip-split?"
        )


def _apply_defaults(args: argparse.Namespace, defaults: dict[str, Any],
                    model_type: str, world_size: int) -> dict[str, dict[str, Any]]:
    """Returns args_applied: dict mapping flag → {value, source}.
    Source is 'user' if the user passed a value on the CLI, else 'default-config'
    (from defaults_pretrain.json) or 'auto-1gpu' / 'auto-multi-gpu' for the
    auto-fallback values. Only includes flags relevant to the model_type."""
    applied: dict[str, dict[str, Any]] = {}

    training_defaults = defaults.get("training", {})
    loss_defaults = defaults.get("loss", {})
    decoder_defaults = defaults.get("add_cmim_decoder", {})

    for f in TRAINING_FLAGS:
        merge_default_into_applied(applied, args, f, training_defaults)

    # Single-GPU fallback: batch_size 32, save_interval 500.
    if world_size <= 1:
        if applied.get("batch_size", {}).get("source") != "user":
            applied["batch_size"] = {"value": 32, "source": "auto-1gpu"}
        if applied.get("save_interval", {}).get("source") != "user":
            applied["save_interval"] = {"value": 500, "source": "auto-1gpu"}

    if model_type in ("cmim", "hybrid"):
        for f in LOSS_FLAGS if model_type == "hybrid" else ("contrastive_temperature",):
            merge_default_into_applied(applied, args, f, loss_defaults)
        for f in DECODER_FLAGS:
            merge_default_into_applied(applied, args, f, decoder_defaults)

    return applied


def _arch_from_validator(validator_out: dict[str, Any]) -> dict[str, Any]:
    arch = validator_out.get("arch") or {}
    missing = [k for k in ARCH_FLAGS_FROM_CKPT if arch.get(k) is None]
    if missing:
        raise ValueError(
            f"checkpoint validator did not surface required arch fields: {missing}. "
            "If the ckpt has no saved_args blob, these can't be inferred from state-dict "
            "shapes alone; please supply a ckpt with args saved (the standard "
            "save_model_for_restart format)."
        )
    return arch


def _build_argv(
    *, world_size: int, gpus_str: str, out_dir: Path, manifest: dict[str, Any],
    model_type: str, pretrain_mode: str, arch: dict[str, Any],
    applied: dict[str, dict[str, Any]],
) -> list[str]:
    """Constructs the full pretrain_ddp.py argument list as a list of strings."""
    outputs = manifest["outputs"]
    argv = [sys.executable, "-u", str(PRETRAIN_DDP_PATH)]

    # Data + vocab paths
    argv += ["--train_data_path", outputs["train_dir"],
             "--val_data_path",   outputs["val_dir"],
             "--atom_vocab_path", outputs["atom_vocab"],
             "--bond_vocab_path", outputs["bond_vocab"]]
    if model_type in ("cmim", "hybrid"):
        argv += ["--smiles_vocab_path", outputs["smiles_vocab"]]

    # Pretrain mode + loss
    argv += ["--pretrain_mode", pretrain_mode]
    if "vocab_loss_weight" in applied and model_type == "hybrid":
        argv += ["--vocab_loss_weight", str(applied["vocab_loss_weight"]["value"])]
    if "contrastive_temperature" in applied and model_type in ("cmim", "hybrid"):
        argv += ["--contrastive_temperature", str(applied["contrastive_temperature"]["value"])]
    # cMIM/decoder arch: emit every applied flag. For continue-pretrain on a
    # cmim/hybrid ckpt, every entry will be source="ckpt_saved_args" (see the
    # overlay loop in run()). For pretrain-from-scratch / add-cmim-pretrain
    # the values come from defaults_pretrain.json's add_cmim_decoder block.
    if model_type in ("cmim", "hybrid"):
        for f in ("latent_dim", "decoder_num_layers", "decoder_num_attention_heads",
                  "decoder_ffn_hidden_size", "decoder_dropout",
                  "decoder_max_seq_len", "decoder_positional_encoding"):
            if f in applied:
                argv += [f"--{f}", str(applied[f]["value"])]
        # Boolean store_true flags: emit the bare flag only when True.
        if applied.get("decoder_gate_self_attn", {}).get("value"):
            argv += ["--decoder_gate_self_attn"]
        if applied.get("decoder_gate_cross_attn", {}).get("value"):
            argv += ["--decoder_gate_cross_attn"]

    # Architecture — sourced from validator's arch block, never from CLI/defaults.
    argv += [
        "--hidden_size",  str(arch["hidden_size"]),
        "--depth",        str(arch["depth"]),
        "--num_attn_head", str(arch["num_attn_head"]),
        "--activation",   str(arch["activation"]),
        "--backbone",     str(arch["backbone"]),
        "--embedding_output_type", str(arch["embedding_output_type"]),
    ]
    if arch.get("self_attention"):
        argv += ["--self_attention"]

    # Training schedule
    for name in ("batch_size", "dropout", "epochs", "init_lr", "max_lr", "final_lr",
                 "warmup_epochs", "weight_decay", "save_interval", "seed"):
        if name in applied:
            argv += [f"--{name}", str(applied[name]["value"])]
    if applied.get("tensorboard", {}).get("value"):
        argv += ["--tensorboard"]
    if applied.get("use_cuikmolmaker_featurization", {}).get("value"):
        argv += ["--use_cuikmolmaker_featurization"]

    # W&B logging (pass-through; pretrain_ddp.py only inits W&B when project is set).
    if "wandb_project" in applied:
        argv += ["--wandb_project", str(applied["wandb_project"]["value"])]
        if "wandb_run_name" in applied:
            argv += ["--wandb_run_name", str(applied["wandb_run_name"]["value"])]

    # Where pretrain_ddp.py auto-resumes from (we'll symlink the user ckpt there).
    argv += ["--save_dir", str(out_dir / "ckpt")]

    return argv


def _symlink_ckpt_into_save_dir(user_ckpt: Path, save_dir: Path) -> Path:
    """--resume path: symlink the user ckpt as <save_dir>/last_checkpoint.pt.
    pretrain_ddp.py's auto-resume then restores everything from the ckpt:
    model weights, optimizer state, scheduler_step, epoch, batch_idx,
    wandb_run_id."""
    save_dir.mkdir(parents=True, exist_ok=True)
    link = save_dir / "last_checkpoint.pt"
    if link.exists() or link.is_symlink():
        link.unlink()
    # Symlink to the absolute user_ckpt so it works regardless of cwd.
    link.symlink_to(user_ckpt.resolve())
    return link


# Schedule fields that --resume inherits from ckpt.saved_args and that default
# (fresh-schedule) mode takes from CLI/defaults_pretrain.json.
SCHEDULE_FLAGS = ("epochs", "warmup_epochs", "init_lr", "max_lr", "final_lr")


def _materialize_ckpt_for_fresh_schedule(user_ckpt: Path, save_dir: Path) -> Path:
    """Default (fresh-schedule) continue-pretrain path: write a CLEANED copy
    of the user ckpt to <save_dir>/last_checkpoint.pt with scheduler_step,
    epoch, batch_idx, and wandb_run_id reset to fresh-start values. Model
    weights AND optimizer state pass through unchanged — so Adam's running
    moments warm-start the new schedule (helpful because the new init_lr is
    usually close to the previous run's final_lr).

    Why a fresh-state copy instead of a symlink: pretrain_ddp.py's
    `trainer.load()` restores EVERYTHING in the ckpt including scheduler_step
    and epoch. We can't selectively load just the model + optimizer through
    that code path. The minimal-invasive workaround is to materialize a
    ckpt that has the unwanted counters zeroed before the loader sees it.
    pretrain_ddp.py then restores everything as normal, but everything it
    restores reads as a fresh-start.

    Cost: one ~700 MB disk write per run. Pretrain is days-long, so it's
    negligible. Done on the host before docker run.
    """
    import torch  # delayed import — keeps the runner light in --dry-run paths
    save_dir.mkdir(parents=True, exist_ok=True)
    target = save_dir / "last_checkpoint.pt"
    if target.exists() or target.is_symlink():
        target.unlink()
    ckpt = load_checkpoint(user_ckpt)
    if not isinstance(ckpt, dict) or "state_dict" not in ckpt:
        raise ValueError(
            f"ckpt {user_ckpt} is not in the expected save_model_for_restart "
            "dict format (need at least 'state_dict' key)."
        )
    ckpt["scheduler_step"] = 0
    ckpt["epoch"] = 0
    ckpt["batch_idx"] = 0
    ckpt["wandb_run_id"] = None
    torch.save(ckpt, target)
    return target


def _validate_resume_state(user_ckpt: Path) -> dict[str, Any]:
    """--resume mode: confirm the ckpt was saved via the save_model_for_restart
    format and carries the full state pretrain_ddp.py needs to resume mid-run
    (optimizer state, scheduler_step, epoch, batch_idx). Returns a small
    `resume_state` dict for the manifest so users can see what was restored.
    Raises ValueError with a clear redirect if the ckpt is too lean."""
    ckpt = load_checkpoint(user_ckpt)
    if not isinstance(ckpt, dict) or "state_dict" not in ckpt:
        raise ValueError(
            f"ckpt {user_ckpt} is not in the expected save_model_for_restart "
            "dict format."
        )
    required = ("optimizer", "scheduler_step", "epoch", "batch_idx")
    missing = [k for k in required if k not in ckpt]
    if missing:
        raise ValueError(
            f"--resume requires the ckpt to carry the full mid-run state, but "
            f"these keys are missing: {missing}. The ckpt was probably saved "
            "without enough metadata to pure-resume — use the default "
            "fresh-schedule mode (drop --resume) if you just want to continue "
            "training with a new schedule."
        )
    return {
        "scheduler_step": int(ckpt["scheduler_step"]),
        "epoch": int(ckpt["epoch"]),
        "batch_idx": int(ckpt["batch_idx"]),
        "wandb_run_id": ckpt.get("wandb_run_id"),
    }


# Vocab-entry counting delegated to _utils.count_vocab_entries. Alias kept for
# the existing internal callsites.
_count_vocab_entries = count_vocab_entries


def _verify_vocab_sizes_match_ckpt(
    manifest: dict[str, Any], validator_out: dict[str, Any], model_type: str,
) -> dict[str, Any]:
    """For continue-pretrain only: compare each vocab file's entry count against
    the ckpt's vocab head dimensions. Aborts on mismatch with a helpful error
    pointing the user at the matching vocab. Returns a `vocab_check` block to
    attach to run.json for transparency."""
    ckpt_sizes = validator_out.get("vocab_sizes") or {"atom": None, "bond": None, "smiles": None}
    outputs = manifest.get("outputs", {})
    check: dict[str, Any] = {"vocab_source": manifest.get("vocab_source", "unknown")}
    for which in ("atom", "bond", "smiles"):
        ckpt_size = ckpt_sizes.get(which)
        vocab_path_str = outputs.get(f"{which}_vocab")
        check[which] = {"ckpt_size": ckpt_size, "manifest_vocab": vocab_path_str, "manifest_size": None}
        if ckpt_size is None:
            # ckpt doesn't have this head; nothing to verify.
            continue
        # ckpt has this head — the manifest MUST include the corresponding vocab.
        if not vocab_path_str:
            raise ValueError(
                f"ckpt has a '{which}' vocab head (size {ckpt_size}) but the prepare_data "
                f"manifest doesn't include a {which}_vocab file. Rerun prepare_data with "
                f"--vocab-dir <ckpt's parent dir> (or --{which}-vocab <path>) so the runner "
                f"can pass the matching vocab through."
            )
        manifest_size = _count_vocab_entries(Path(vocab_path_str))
        check[which]["manifest_size"] = manifest_size
        if manifest_size != ckpt_size:
            raise ValueError(
                f"{which} vocab size mismatch — ckpt's head expects {ckpt_size} entries, "
                f"but {vocab_path_str} has {manifest_size}. The released ckpt's vocab is the "
                f"authoritative one for continue-pretrain; pass --vocab-dir <ckpt's parent dir> "
                f"(or --{which}-vocab <path>) to prepare_data so vocab built from the new corpus "
                f"isn't used. If you actually want to pretrain from scratch on a different "
                f"vocab, use the kermt-pretrain-scratch workflow instead."
            )
    return check


# ---------------------------------------------------------------------------
# Main flow
# ---------------------------------------------------------------------------

def run(args: argparse.Namespace) -> dict[str, Any]:
    out_dir = Path(args.out).resolve()
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "ckpt").mkdir(parents=True, exist_ok=True)
    (out_dir / "logs").mkdir(parents=True, exist_ok=True)

    from_scratch = bool(args.from_scratch)
    resume = bool(args.resume)

    # Mode-conflict validation up-front so the user fails fast.
    if from_scratch and resume:
        raise ValueError("--resume is incompatible with --from-scratch.")
    if resume and not args.ckpt:
        raise ValueError("--resume requires --ckpt; nothing to resume from otherwise.")
    if resume:
        # CLI overrides of schedule args are forbidden in --resume mode — pure
        # resume means the schedule shape from the ckpt is authoritative.
        cli_overrides = [
            f for f in SCHEDULE_FLAGS if getattr(args, f, None) is not None
        ]
        if cli_overrides:
            raise ValueError(
                f"--resume inherits schedule args from the ckpt's saved_args; "
                f"explicit CLI override is forbidden. You passed: {cli_overrides}. "
                "Drop those flags to pure-resume, or use the default fresh-schedule "
                "mode (no --resume) if you want a new schedule."
            )

    workflow = "pretrain-scratch" if from_scratch else "continue-pretrain"
    if resume:
        mode = "continue_pretrain_resume"
    elif from_scratch:
        mode = "pretrain_from_scratch"
    else:
        mode = "continue_pretrain_fresh_schedule"

    # 1. Load defaults + prepare manifest.
    defaults = _load_json(DEFAULTS_PATH, name="defaults_pretrain.json")
    prep_manifest_path = Path(args.prepare_manifest).resolve()
    manifest = _load_json(prep_manifest_path, name="prepare_data.json")
    _verify_prepare_manifest(manifest)

    # 2. Branch: continue-pretrain (load ckpt + validate) vs from-scratch (no ckpt).
    ckpt: Path | None = None
    validator_out: dict[str, Any] | None = None
    link: Path | None = None
    vocab_check: dict[str, Any] | None = None
    resume_state: dict[str, Any] | None = None  # populated only when --resume

    if from_scratch:
        if args.ckpt:
            raise ValueError("--from-scratch is incompatible with --ckpt; pass one or the other.")
        if not args.pretrain_target_mode:
            raise ValueError("--pretrain-target-mode is required when --from-scratch is set "
                             "(choose vocab, cmim, or hybrid).")
        pretrain_mode = args.pretrain_target_mode
        model_type = {"vocab": "grover_base", "cmim": "cmim", "hybrid": "hybrid"}[pretrain_mode]
        # Arch from defaults_pretrain.json's `arch` group (with CLI overrides applied later
        # if we expose any; for now we just use defaults).
        arch_defaults = defaults.get("arch") or {}
        if not arch_defaults:
            raise ValueError("defaults_pretrain.json has no `arch` group; cannot pretrain from scratch.")
        arch = {k: arch_defaults.get(k) for k in ARCH_FLAGS_FROM_CKPT}
        # `latent_dim` lives in the add_cmim_decoder group for from-scratch cmim/hybrid;
        # treat it as part of the arch for argv-building purposes.
        if pretrain_mode in ("cmim", "hybrid"):
            arch["latent_dim"] = (defaults.get("add_cmim_decoder") or {}).get("latent_dim")
        else:
            arch["latent_dim"] = None
    else:
        if not args.ckpt:
            raise ValueError("--ckpt is required for continue-pretrain. "
                             "Use --from-scratch to pretrain a fresh model on the corpus.")
        ckpt = Path(args.ckpt).resolve()
        if args.ckpt_validator_out:
            validator_out = _load_json(Path(args.ckpt_validator_out), name="ckpt validator output")
        else:
            validator_out = run_checkpoint_validator(ckpt, mode="continue_pretrain", script_path=CHECK_CHECKPOINT_PATH)
        if not validator_out.get("ok"):
            raise ValueError(
                f"check_checkpoint.py rejected the input ckpt: {validator_out.get('errors')}"
            )
        model_type = validator_out.get("model_type")
        if model_type not in MODEL_TYPE_TO_PRETRAIN_MODE:
            raise ValueError(
                f"model_type='{model_type}' cannot continue pretrain. "
                f"Supported: {sorted(MODEL_TYPE_TO_PRETRAIN_MODE)}. "
                "For an encoder-only ckpt with no pretrain head, use the "
                "upgrade_to_hybrid workflow."
            )
        if model_type == "grover_base" and not validator_out.get("has_vocab_head"):
            raise ValueError(
                "grover_base ckpt has no vocab head — cannot continue vocab pretrain. "
                "Use the upgrade_to_hybrid workflow to add a cMIM decoder, "
                "or finetune directly from the encoder."
            )
        pretrain_mode = MODEL_TYPE_TO_PRETRAIN_MODE[model_type]
        arch = _arch_from_validator(validator_out)
        # Vocab-size verification — refuse mismatched corpora before launching pretrain_ddp.py.
        vocab_check = _verify_vocab_sizes_match_ckpt(manifest, validator_out, model_type)
        # --resume needs the ckpt to carry the full mid-run state. Validate now;
        # also surface what's being restored in the manifest.
        if resume:
            resume_state = _validate_resume_state(ckpt)

    # 3. GPU selection.
    world_size, gpus_str = _detect_gpus(args.gpus)
    if world_size <= 0:
        raise ValueError(
            "No GPUs detected. pretrain_ddp.py requires at least one CUDA device. "
            "Set CUDA_VISIBLE_DEVICES or pass --gpus <ids>."
        )

    # 4. Apply defaults + collect args_applied.
    applied = _apply_defaults(args, defaults, model_type, world_size)
    # --resume overlays schedule args from the ckpt's saved_args (the only path
    # where source="ckpt_saved_args" can appear in args_applied). Fail loudly if
    # any schedule field is missing from saved_args — pure-resume can't proceed
    # without the original schedule shape.
    if resume:
        saved_args = validator_out.get("saved_args") or {}
        missing = [f for f in SCHEDULE_FLAGS if f not in saved_args]
        if missing:
            raise ValueError(
                f"--resume requires the ckpt's saved_args to include all schedule "
                f"fields, but these are missing: {missing}. The ckpt was saved "
                "without enough metadata to pure-resume — use the default "
                "fresh-schedule mode and specify --epochs / --warmup-epochs / "
                "--init-lr / --max-lr / --final-lr explicitly."
            )
        for f in SCHEDULE_FLAGS:
            applied[f] = {"value": saved_args[f], "source": "ckpt_saved_args"}

    # Continue-pretrain on a cmim/hybrid ckpt: cMIM/decoder arch must come
    # from the ckpt's saved_args, not from defaults or CLI. This is the
    # cmim/decoder analogue of the encoder-arch passthrough already done by
    # `_arch_from_validator` (and matches the README guarantee that
    # `add_cmim_decoder` defaults are ignored during continue-pretrain).
    if not from_scratch and model_type in ("cmim", "hybrid"):
        cli_latent_dim_override = args.latent_dim is not None
        if cli_latent_dim_override:
            raise ValueError(
                "--latent-dim cannot be overridden during continue-pretrain on a "
                "cmim/hybrid ckpt — the value is fixed by the ckpt's saved_args "
                "(passing a different value would mismatch the loaded decoder "
                "weights). Drop --latent-dim, or use kermt-pretrain-scratch if "
                "you intentionally want a different latent dimension."
            )
        saved_args = validator_out.get("saved_args") or {}
        cmim_missing = [f for f in CMIM_DECODER_FLAGS_FROM_CKPT if f not in saved_args]
        if cmim_missing:
            raise ValueError(
                f"continue-pretrain on a {model_type} ckpt requires the ckpt's "
                f"saved_args to include cmim/decoder arch fields, but these are "
                f"missing: {cmim_missing}. The ckpt was saved without enough "
                "metadata to faithfully reconstruct the decoder."
            )
        for f in CMIM_DECODER_FLAGS_FROM_CKPT:
            applied[f] = {"value": saved_args[f], "source": "ckpt_saved_args"}

    # Optional W&B logging: pass-through, no defaults — forwarded only when the
    # user sets --wandb-project (run name is honored only alongside a project).
    for f in ("wandb_project", "wandb_run_name"):
        v = getattr(args, f, None)
        if v is not None:
            applied[f] = {"value": v, "source": "user"}

    # 5. Build the pretrain_ddp.py argv.
    argv = _build_argv(
        world_size=world_size, gpus_str=gpus_str, out_dir=out_dir, manifest=manifest,
        model_type=model_type, pretrain_mode=pretrain_mode, arch=arch, applied=applied,
    )

    # 6. (continue-pretrain only) Stage the ckpt into <save_dir>/last_checkpoint.pt
    #    so pretrain_ddp.py's auto-resume picks it up. Mode-dispatched:
    #    - --resume: symlink to user ckpt. pretrain_ddp.py restores everything
    #      (model + optimizer + scheduler_step + epoch + batch_idx + wandb_run_id).
    #    - default (fresh-schedule): materialize a state-cleaned copy of the
    #      ckpt — model weights + optimizer pass through, but scheduler_step /
    #      epoch / batch_idx / wandb_run_id are reset to 0/None. pretrain_ddp.py
    #      then builds a fresh NoamLR from CLI args and starts from step 0.
    # Done unconditionally (including --dry-run) so the dry-run faithfully
    # exercises ckpt I/O — catches corrupt ckpts / insufficient disk before
    # the days-long real run.
    if not from_scratch:
        if resume:
            link = _symlink_ckpt_into_save_dir(ckpt, out_dir / "ckpt")
        else:
            link = _materialize_ckpt_for_fresh_schedule(ckpt, out_dir / "ckpt")

    # 7. Build the run.json manifest.
    commit, dirty = git_commit_with_env_override(REPO_ROOT)
    image_tag = os.environ.get("KERMT_IMAGE", "kermt:latest")
    image_digest = docker_image_digest(image_tag)
    cmd_replay_env: dict[str, str] = {}
    if gpus_str:
        cmd_replay_env["CUDA_VISIBLE_DEVICES"] = gpus_str
    cmd_replay_env["WORLD_SIZE"] = str(world_size)
    cmd_replay = format_cmd_replay(argv, env=cmd_replay_env)
    run_manifest = {
        "workflow": workflow,
        "mode": mode,  # pretrain_from_scratch | continue_pretrain_fresh_schedule | continue_pretrain_resume
        "started_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "container": {"image_tag": image_tag, "image_digest": image_digest},
        "repo": {"commit": commit, "dirty": dirty},
        "inputs": {
            "ckpt": str(ckpt) if ckpt else None,
            "prepare_data_manifest": str(prep_manifest_path),
            "ckpt_validator_out": (
                str(Path(args.ckpt_validator_out).resolve()) if args.ckpt_validator_out else None
            ),
        },
        "model_type": model_type,
        "pretrain_mode": pretrain_mode,
        "world_size": world_size,
        "cuda_visible_devices": gpus_str,
        "args_applied": applied,
        "arch": arch,
        "vocab_check": vocab_check,  # None for from-scratch
        "resume_state": resume_state,  # None unless --resume; carries the restored scheduler_step / epoch / batch_idx / wandb_run_id from the ckpt
        "save_dir": str(out_dir / "ckpt"),
        "logs_dir": str(out_dir / "logs"),
        "tensorboard_dir": str(out_dir / "logs" / "tb"),
        "argv": argv,
        "cmd_replay": cmd_replay,
        "ok_to_replay": (not dirty) and (commit != "unknown"),
        "dry_run": bool(args.dry_run),
        "ckpt_symlink": str(link) if link else None,
        "from_scratch": from_scratch,
    }
    (out_dir / "run.json").write_text(json.dumps(run_manifest, indent=2))

    # 8. Execute (unless --dry-run).
    if args.dry_run:
        run_manifest["status"] = "dry_run"
        return run_manifest

    env = runner_environment(REPO_ROOT, wandb="wandb_project" in applied)
    env["WORLD_SIZE"] = str(world_size)
    if gpus_str:
        env["CUDA_VISIBLE_DEVICES"] = gpus_str

    log_file = out_dir / "logs" / "pretrain_ddp.log"
    with log_file.open("w") as logf:
        proc = subprocess.run(argv, env=env, stdout=logf, stderr=subprocess.STDOUT)
    run_manifest["exit_code"] = proc.returncode
    run_manifest["status"] = "ok" if proc.returncode == 0 else "failed"
    (out_dir / "run.json").write_text(json.dumps(run_manifest, indent=2))
    return run_manifest


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(
        description="Workstation pretrain runner (continue-pretrain by default, "
                    "or pretrain-from-scratch with --from-scratch).")
    p.add_argument("--ckpt", default=None,
                   help="Path to the input pretrain checkpoint. Required for continue-pretrain; "
                        "omit when --from-scratch is set.")
    p.add_argument("--from-scratch", action="store_true",
                   help="Pretrain a fresh model on the corpus (no input ckpt; arch from "
                        "defaults_pretrain.json; vocab built by prepare_data). Requires "
                        "--pretrain-target-mode.")
    p.add_argument("--resume", action="store_true",
                   help="Resume an interrupted pretrain run (crashed / Ctrl-C / OOM). "
                        "Restores everything from the ckpt: model weights, optimizer "
                        "state, scheduler_step, epoch, batch_idx, wandb_run_id. Schedule "
                        "shape (epochs / warmup_epochs / init/max/final_lr) is inherited "
                        "from the ckpt's saved_args; CLI overrides of schedule flags are "
                        "REJECTED in this mode. Without --resume (default), continue-pretrain "
                        "loads only model weights + optimizer momentum from the ckpt and "
                        "starts a fresh schedule from CLI/defaults_pretrain.json — use that "
                        "default mode when continue-pretraining on a new corpus / new "
                        "objective / extended training (the common case).")
    p.add_argument("--pretrain-target-mode", choices=["vocab", "cmim", "hybrid"], default=None,
                   help="(--from-scratch only) which pretrain objective to use for the fresh "
                        "model: vocab (grover_base-style), cmim, or hybrid (vocab + contrast). "
                        "No default — must be set explicitly so the user makes an informed "
                        "choice about the head config.")
    p.add_argument("--prepare-manifest", required=True,
                   help="Path to a prepare_data.json (must be mode=pretrain)")
    p.add_argument("--out", required=True, help="Output run directory")
    p.add_argument("--ckpt-validator-out", default=None,
                   help="Optional cached check_checkpoint.py JSON; computed if absent")
    p.add_argument("--gpus", default=None,
                   help="Comma-separated GPU ids (e.g. '0,1'). Default: all visible")
    p.add_argument("--dry-run", action="store_true",
                   help="Write run.json and print the command without executing")
    # Training overrides — all default to None so we can distinguish user-given vs default-config.
    for f, t in [("epochs", int), ("batch-size", int), ("init-lr", float), ("max-lr", float),
                 ("final-lr", float), ("warmup-epochs", float), ("weight-decay", float),
                 ("dropout", float), ("save-interval", int), ("seed", int),
                 ("vocab-loss-weight", float), ("latent-dim", int),
                 ("contrastive-temperature", float)]:
        p.add_argument(f"--{f}", type=t, default=None)
    # Optional W&B logging (pass-through to pretrain_ddp.py; off unless project is set).
    p.add_argument("--wandb-project", type=str, default=None,
                   help="W&B project name. When set, pretrain_ddp.py logs train/val losses.")
    p.add_argument("--wandb-run-name", type=str, default=None,
                   help="Optional W&B run name (only used when --wandb-project is set).")
    args = p.parse_args(argv)

    try:
        manifest = run(args)
    except (FileNotFoundError, ValueError, RuntimeError) as exc:
        print(json.dumps({"ok": False, "errors": [f"{type(exc).__name__}: {exc}"]}, indent=2),
              file=sys.stdout)
        return 1
    except Exception as exc:  # noqa: BLE001
        import traceback
        print(traceback.format_exc(), file=sys.stderr)
        print(json.dumps({"ok": False, "errors": [f"unhandled: {type(exc).__name__}: {exc}"]},
                         indent=2))
        return 1

    print(json.dumps({"ok": True, "manifest": manifest}, indent=2))
    return 0 if manifest.get("status") != "failed" else 1


if __name__ == "__main__":
    sys.exit(main())

#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Convert a grover_base ckpt into a hybrid ckpt by adding a
randomly-initialized cMIM decoder + latent_dist on top of the loaded encoder.

After upgrade, the saved checkpoint classifies as `model_type: hybrid` via
`check_checkpoint.py --mode continue_pretrain`, and the pretrain runner
(`run_pretrain_local.py`) drops it directly into the
`--pretrain_mode hybrid` dispatch path with no special-case logic.

What's preserved from the input ckpt
------------------------------------
- Encoder weights (kermt.encoders.* state-dict subset). For legacy
  grover_base ckpts where the encoder lives under `grover.encoders.*`,
  the prefix is renormalized to `kermt.encoders.*` to match the
  KermtHybridTask layout.
- The saved `args` Namespace — augmented with hybrid-specific decoder fields
  if absent (defaults from `config/defaults_pretrain.json`).

What's fresh-initialized
------------------------
- The cMIM decoder (`decoder.*` — SMILESTransformerDecoder) — Xavier-init.
- The latent distribution (`latent_dist.*` minus `latent_dist.kermt.*` —
  the encoder lives inside latent_dist by reference, so it shares the
  same weights as the top-level encoder).
- The vocab heads (`vocab_module.*`) — always fresh, sized to the
  vocab built from the user's pretrain corpus. The previous heads (if the
  input ckpt was a modern grover_base with vocab heads) are discarded —
  rebuilding them is acceptable because continue-pretrain will retrain
  them anyway, and matching the new vocab dimensions is more important
  than warm-starting from the old vocab.

CLI
---
    upgrade_to_hybrid.py
        --ckpt <input.pt>             # grover_base ckpt (encoder-only or
                                      # encoder + vocab heads; legacy or modern)
        --prepare-manifest <path>     # prepare_data.json from a prior
                                      # `prepare_data.py --mode pretrain` run.
                                      # The smiles_vocab + atom_vocab +
                                      # bond_vocab from this manifest size
                                      # the new heads.
        --out <upgraded.pt>           # destination for the upgraded ckpt
        [--latent-dim N]              # override defaults_pretrain.json
        [--seed 0]                    # for reproducible random init

Exit code 0 + a one-line JSON summary on success. Exit code 1 + JSON-shaped
errors on failure (input rejection, vocab missing, shape mismatch, etc.).
"""
from __future__ import annotations

import argparse
import json
import sys
import traceback
from argparse import Namespace
from pathlib import Path
from typing import Any

import torch

# sys.path tweak so `_utils` imports cleanly whether launched via `kermt_run`
# or as a bare `python scripts/upgrade_to_hybrid.py …`.
if str(Path(__file__).resolve().parent) not in sys.path:
    sys.path.insert(0, str(Path(__file__).resolve().parent))
from _utils import count_vocab_entries, load_checkpoint, load_json, run_checkpoint_validator  # noqa: E402


SKILL_ROOT = Path(__file__).resolve().parent.parent
DEFAULTS_PATH = SKILL_ROOT / "config" / "defaults_pretrain.json"
CHECK_CHECKPOINT_PATH = SKILL_ROOT / "scripts" / "check_checkpoint.py"

# Fields KermtHybridTask + its sub-components read from args. Anything not
# present on the input ckpt's args gets filled from defaults_pretrain.json
# (decoder fields from add_cmim_decoder group; encoder fields from training
# group; the rest from hardcoded encoder-arch defaults below).
HYBRID_REQUIRED_DECODER_ARGS = (
    "decoder_num_layers", "decoder_num_attention_heads", "decoder_ffn_hidden_size",
    "decoder_dropout", "decoder_max_seq_len", "decoder_positional_encoding",
    "decoder_gate_self_attn", "decoder_gate_cross_attn",
)

# Fields KERMTEmbedding (and its sub-modules) read at __init__ time. Legacy
# grover_base ckpts may be missing several of these — pre-cMIM args weren't a
# strict superset of what current KERMTEmbedding wants.
ENCODER_REQUIRED_ARGS_WITH_DEFAULTS = {
    "dropout": 0.1,                      # not in legacy grover_base.pt
    "bond_drop_rate": 0.0,
    "self_attention": False,
    "attn_hidden": 4,
    "attn_out": 8,
    "use_cuikmolmaker_featurization": False,
    "input_layer": "fc",
    "dense": False,
    "bias": False,
    "undirected": False,
    "dist_coff": 0.1,
    "num_mt_block": 1,
    "embedding_output_type": "both",
}


# JSON loading + vocab counting delegated to _utils. Aliases kept for the
# existing internal callsites.
_load_json = load_json
_count_vocab = count_vocab_entries


def _run_validator(ckpt: Path) -> dict[str, Any]:
    """Invoke check_checkpoint.py --mode upgrade_to_hybrid and return the JSON.
    Thin alias around _utils.run_checkpoint_validator for the specific mode."""
    return run_checkpoint_validator(ckpt, mode="upgrade_to_hybrid", script_path=CHECK_CHECKPOINT_PATH)


def _augment_args(input_args: Namespace, decoder_defaults: dict[str, Any], latent_dim: int) -> Namespace:
    """Make sure the args Namespace has every field KermtHybridTask reads.
    Existing input args take precedence (the ckpt was trained with them, and
    encoder arch needs them); only missing hybrid-specific fields are filled
    from defaults_pretrain.json's add_cmim_decoder group."""
    out = Namespace(**vars(input_args)) if isinstance(input_args, Namespace) else Namespace(**input_args)
    # Encoder-side fields KERMTEmbedding reads at __init__. Legacy grover_base
    # ckpts may be missing several (e.g. `dropout`). Fill from hardcoded
    # defaults — these are safe at upgrade time because the loaded weights
    # determine the actual encoder behavior; missing dropout etc. only affects
    # rebuild/forward semantics, which the continue-pretrain runner will
    # override anyway via its own training defaults.
    for field, default_value in ENCODER_REQUIRED_ARGS_WITH_DEFAULTS.items():
        if not hasattr(out, field):
            setattr(out, field, default_value)
    # Hybrid-specific decoder fields (only fill in if missing)
    for field in HYBRID_REQUIRED_DECODER_ARGS:
        if not hasattr(out, field):
            setattr(out, field, decoder_defaults[field])
    # latent_dim and contrastive_temperature also fill in from defaults if absent
    if not hasattr(out, "latent_dim"):
        out.latent_dim = latent_dim
    if not hasattr(out, "contrastive_temperature"):
        out.contrastive_temperature = decoder_defaults.get("contrastive_temperature", 0.1)
    # pretrain_mode and use_cmim flags — set explicitly to hybrid semantics
    out.pretrain_mode = "hybrid"
    if hasattr(out, "use_cmim"):
        out.use_cmim = True
    # vocab_loss_weight default from training/loss defaults (1.0 — set inside the runner;
    # here we just need it on the Namespace so save_checkpoint records it)
    if not hasattr(out, "vocab_loss_weight"):
        out.vocab_loss_weight = 1.0
    return out


def _renormalize_encoder_keys(state_dict: dict[str, Any]) -> tuple[dict[str, Any], list[str]]:
    """Extract the encoder subset of the input state_dict, renormalized so the
    keys match the KermtHybridTask layout (`kermt.encoders.*`).

    Three input shapes are handled:
      - `grover.encoders.*` (legacy original-GROVER ckpts) → rename to `kermt.encoders.*`
      - `kermt.encoders.*`  (modern repo-trained grover_base / hybrid) → keep as-is
      - `latent_dist.kermt.encoders.*` (cmim ckpts) → reject — those should
        go through `continue_pretrain`, not `upgrade_to_hybrid`.

    Vocab heads (`vocab_module.*`) and any other non-encoder keys are
    DROPPED here — the upgraded hybrid task gets fresh vocab heads sized to
    the user's prepare-data vocab.

    Returns (renormalized_state_dict, notes). `notes` is a list of strings
    describing key handling for the manifest.
    """
    out: dict[str, Any] = {}
    notes: list[str] = []
    legacy_count = modern_count = vocab_dropped = other_dropped = 0

    for k, v in state_dict.items():
        if k.startswith("latent_dist.kermt."):
            raise ValueError(
                "input ckpt has `latent_dist.kermt.*` keys — it appears to be a "
                "cmim or hybrid ckpt, not a grover_base. Use kermt-continue-pretrain "
                "(no upgrade needed) or kermt-pretrain-scratch instead."
            )
        if k.startswith("grover.encoders."):
            new_k = "kermt.encoders." + k[len("grover.encoders."):]
            out[new_k] = v
            legacy_count += 1
        elif k.startswith("kermt.encoders."):
            out[k] = v
            modern_count += 1
        elif k.startswith("vocab_module."):
            vocab_dropped += 1  # fresh heads on the new hybrid
        else:
            other_dropped += 1

    if legacy_count and modern_count:
        notes.append(f"unexpected mix of legacy + modern encoder prefixes: "
                     f"{legacy_count} grover.encoders.* + {modern_count} kermt.encoders.*")
    elif legacy_count:
        notes.append(f"renamed {legacy_count} legacy `grover.encoders.*` keys to `kermt.encoders.*`")
    elif modern_count:
        notes.append(f"kept {modern_count} modern `kermt.encoders.*` keys as-is")
    else:
        raise ValueError("no encoder state-dict keys found in input ckpt "
                         "(expected `grover.encoders.*` or `kermt.encoders.*`)")
    if vocab_dropped:
        notes.append(f"dropped {vocab_dropped} `vocab_module.*` keys "
                     f"(new vocab heads sized to user's prepare-data vocab)")
    if other_dropped:
        notes.append(f"dropped {other_dropped} other keys (not encoder, not vocab)")
    return out, notes


def upgrade(args: argparse.Namespace) -> dict[str, Any]:
    """Main upgrade flow. Returns a dict summary suitable for printing as JSON."""
    summary: dict[str, Any] = {
        "ok": False,
        "input_ckpt": str(Path(args.ckpt).resolve()),
        "output_ckpt": str(Path(args.out).resolve()),
        "vocab_sizes_used": {"atom": None, "bond": None, "smiles": None},
        "notes": [],
        "errors": [],
        "warnings": [],
    }

    # 1. Validate input via check_checkpoint --mode upgrade_to_hybrid.
    validator = _run_validator(Path(args.ckpt))
    if not validator.get("ok"):
        summary["errors"].extend(validator.get("errors", []) or ["check_checkpoint rejected the ckpt"])
        return summary
    summary["input_model_type"] = validator.get("model_type")

    # 2. Read prepare manifest for vocab paths + sizes.
    prep_path = Path(args.prepare_manifest)
    prep = _load_json(prep_path, name="prepare_data.json")
    if prep.get("mode") != "pretrain":
        raise ValueError(f"prepare manifest mode='{prep.get('mode')}', expected 'pretrain'")
    if not prep.get("ok"):
        raise ValueError(f"prepare manifest reports ok=False: {prep.get('errors')}")
    outs = prep.get("outputs", {})
    for required in ("atom_vocab", "bond_vocab", "smiles_vocab"):
        if required not in outs:
            raise ValueError(
                f"prepare manifest missing {required}. kermt-add-cmim-pretrain needs the "
                "smiles vocab to size the new decoder; pass a corpus through "
                "`prepare_data.py --mode pretrain` (without --skip-vocab / --skip-features) first."
            )
    atom_size = _count_vocab(Path(outs["atom_vocab"]))
    bond_size = _count_vocab(Path(outs["bond_vocab"]))
    smiles_size = _count_vocab(Path(outs["smiles_vocab"]))
    summary["vocab_sizes_used"] = {"atom": atom_size, "bond": bond_size, "smiles": smiles_size}

    # 3. Load defaults + input ckpt.
    defaults = _load_json(DEFAULTS_PATH, name="defaults_pretrain.json")
    decoder_defaults = defaults.get("add_cmim_decoder") or {}
    latent_dim = args.latent_dim if args.latent_dim is not None else decoder_defaults.get("latent_dim", 800)

    input_ckpt = load_checkpoint(args.ckpt)
    input_args = input_ckpt.get("args")
    if input_args is None:
        raise ValueError(
            "input ckpt has no `args` Namespace — cannot reconstruct the encoder architecture. "
            "Original GROVER ckpts typically saved args; if this one didn't, the upgrade has no "
            "way to recover hidden_size / depth / num_attn_head / etc."
        )
    input_sd = input_ckpt["state_dict"]

    # 4. Strip DDP `module.` prefix if present.
    if input_sd and all(k.startswith("module.") for k in input_sd):
        input_sd = {k[len("module."):]: v for k, v in input_sd.items()}
        summary["notes"].append("stripped DDP `module.` prefix from input state_dict")

    # 5. Filter input state_dict down to just the encoder keys. Vocab heads
    #    are intentionally dropped here — the new hybrid task (built in step 7)
    #    provides freshly-initialized vocab heads sized to the user's vocab,
    #    which is safer than carrying over heads sized to the input ckpt's
    #    (possibly different) vocab.
    encoder_state, rename_notes = _renormalize_encoder_keys(input_sd)
    summary["notes"].extend(rename_notes)

    # 6. Augment args with hybrid-specific decoder fields.
    new_args = _augment_args(input_args, decoder_defaults, latent_dim)
    # Ensure cuda flag is False for the construction step (we don't move to GPU here).
    new_args.cuda = False

    # 7. Build the fresh KermtHybridTask.
    torch.manual_seed(args.seed)
    from kermt.model.models import KermtHybridTask, KERMTEmbedding  # type: ignore
    encoder = KERMTEmbedding(new_args)
    hybrid_task = KermtHybridTask(
        new_args, kermt=encoder,
        latent_dim=latent_dim,
        contrastive_temperature=new_args.contrastive_temperature,
        smiles_vocab_size=smiles_size,
        atom_vocab_size=atom_size,
        bond_vocab_size=bond_size,
    )

    # 8. Load encoder weights into the new task. strict=False because decoder /
    # latent_dist / vocab_module weren't in the input — they keep their fresh init.
    load_result = hybrid_task.load_state_dict(encoder_state, strict=False)
    missing_keys = list(load_result.missing_keys)
    unexpected_keys = list(load_result.unexpected_keys)
    # The encoder is shared between hybrid_task.kermt and hybrid_task.latent_dist.kermt;
    # any `kermt.encoders.*` key that wasn't loaded into the latent_dist's encoder copy
    # is benign because they're the same module by reference. Same for missing-keys that
    # belong to decoder / latent_dist (non-kermt parts) / vocab_module — those are
    # supposed to be fresh.
    summary["encoder_load"] = {
        "missing_keys_total": len(missing_keys),
        "unexpected_keys_total": len(unexpected_keys),
        "encoder_missing": [k for k in missing_keys if "encoders." in k][:5],
        "non_encoder_missing_categories": _categorize_missing(missing_keys),
        "unexpected_sample": unexpected_keys[:5],
    }
    if unexpected_keys:
        summary["warnings"].append(
            f"{len(unexpected_keys)} unexpected key(s) in encoder load — "
            "likely arch drift between legacy GROVER and modern KERMTEmbedding."
        )

    # Surface unknown backbones (anything beyond the gtrans/dualtrans pair
    # that pretrain_ddp.py's argparse + the model code both accept). The
    # legacy `dualtrans` name is the same architecture as `gtrans` and is
    # explicitly handled in kermt/model/models.py:200.
    if getattr(new_args, "backbone", None) not in (None, "gtrans", "dualtrans"):
        summary["warnings"].append(
            f"upgraded ckpt has backbone='{new_args.backbone}', which is neither "
            "'gtrans' nor 'dualtrans'. pretrain_ddp.py's --backbone argparse will reject "
            "it. Either re-pretrain a gtrans grover_base from scratch via "
            "kermt-pretrain-scratch, or extend the parsing.py:379 choices."
        )

    # 9. Build a fresh Adam optimizer over the new model (matches save_checkpoint format).
    init_lr = getattr(new_args, "init_lr", 1e-5)
    weight_decay = getattr(new_args, "weight_decay", 1e-7)
    optimizer = torch.optim.Adam(hybrid_task.parameters(), lr=init_lr, weight_decay=weight_decay)

    # 10. Save in the save_checkpoint format that task/kermttrainer.py:save_checkpoint produces.
    state = {
        "args": new_args,
        "state_dict": hybrid_task.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scheduler_step": 0,
        "batch_idx": 0,
        "epoch": 0,
        "data_scaler": None,
        "features_scaler": None,
        "wandb_run_id": None,
    }
    out_path = Path(args.out).resolve()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(state, out_path)

    summary["ok"] = True
    summary["upgraded_state_dict_keys"] = len(state["state_dict"])
    return summary


def _categorize_missing(missing_keys: list[str]) -> dict[str, int]:
    """Bucket missing keys by their state-dict prefix for the manifest."""
    cats: dict[str, int] = {}
    for k in missing_keys:
        first = k.split(".")[0]
        cats[first] = cats.get(first, 0) + 1
    return cats


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="Upgrade a grover_base ckpt to a hybrid (cMIM + vocab) ckpt.")
    p.add_argument("--ckpt", required=True, help="Path to the input grover_base checkpoint")
    p.add_argument("--prepare-manifest", required=True,
                   help="Path to a prepare_data.json (mode=pretrain) — its vocab sizes "
                        "size the new decoder + fresh vocab heads.")
    p.add_argument("--out", required=True, help="Destination for the upgraded hybrid ckpt")
    p.add_argument("--latent-dim", type=int, default=None,
                   help="Override defaults_pretrain.json's add_cmim_decoder.latent_dim (default 800)")
    p.add_argument("--seed", type=int, default=0,
                   help="Random seed for the fresh decoder / latent / vocab head weights")
    args = p.parse_args(argv)

    try:
        summary = upgrade(args)
    except (FileNotFoundError, ValueError, RuntimeError) as exc:
        print(json.dumps({"ok": False, "errors": [f"{type(exc).__name__}: {exc}"]}, indent=2))
        return 1
    except Exception as exc:  # noqa: BLE001
        print(traceback.format_exc(), file=sys.stderr)
        print(json.dumps({"ok": False, "errors": [f"unhandled: {type(exc).__name__}: {exc}"]}, indent=2))
        return 1

    print(json.dumps(summary, indent=2))
    return 0 if summary.get("ok") else 1


if __name__ == "__main__":
    sys.exit(main())

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Shared utilities for the agent scripts.

Kept intentionally small — only logic that appears (or would otherwise be
duplicated) in two or more `scripts/*.py` modules. Each script
maintains its own primary CLI + main flow.
"""
from __future__ import annotations

import argparse
from collections import Counter
import json
import os
import pickle
import re
import shlex
import subprocess
import sys
from pathlib import Path
from typing import Any


# Conventional pretrain vocab filename stems. Used by prepare_data.py +
# upgrade_to_hybrid.py + the README "Released models" bundling docs +
# the test helpers. Centralized here so a future rename only touches one
# spot.
PRETRAIN_VOCAB_STEMS = {
    "atom":   "pretrain_atom_vocab",
    "bond":   "pretrain_bond_vocab",
    "smiles": "pretrain_smiles_vocab",
}


def resolve_kermt_repo() -> Path:
    """Find the runtime checkout independently of the installed skill location.

    An explicit KERMT_REPO takes precedence. In a repository checkout, walking
    up from this helper or the working directory also supports local use.
    """
    explicit = os.environ.get("KERMT_REPO")
    if explicit:
        candidates = [Path(explicit).expanduser().resolve()]
    else:
        candidates = []
        for start in (Path(__file__).resolve().parent, Path.cwd()):
            candidates.extend((start, *start.parents))
    for candidate in candidates:
        if (candidate / "main.py").is_file() and (candidate / "kermt").is_dir():
            return candidate
    raise FileNotFoundError(
        "KERMT checkout not found. Set KERMT_REPO to the checkout containing "
        "main.py and kermt/; the installed skill directory is separate."
    )


def load_json(path: Path, *, name: str) -> dict[str, Any]:
    """Load a JSON file with consistent error messages.

    `name` is a human-readable label for the document (e.g. "prepare_data.json")
    so the error tells the user which schema we expected at that path.
    """
    if not path.is_file():
        raise FileNotFoundError(f"{name} not found at {path}")
    try:
        return json.loads(path.read_text())
    except json.JSONDecodeError as exc:
        raise ValueError(f"{name} at {path} is not valid JSON: {exc}") from exc


def count_vocab_entries(vocab_path: Path) -> int:
    """Return the number of entries in a KERMT vocab file.

    Handles three layouts:
      - JSON with `{stoi: {token: idx}, ...}` (MolVocab.save_vocab default)
      - JSON as a raw `{token: idx}` dict (legacy / hand-edited)
      - Legacy MolVocab / SMILESVocab pickles, read as inert vocabulary state.

    The pickle reader accepts only the known vocabulary containers and their
    Counter/regex metadata. It cannot import arbitrary classes or run reducers
    supplied by the artifact, and it never falls back to an unrestricted loader.
    """
    if vocab_path.suffix == ".json":
        data = json.loads(vocab_path.read_text())
        if isinstance(data, dict) and "stoi" in data:
            return len(data["stoi"])
        if isinstance(data, dict):
            return len(data)
        raise ValueError(f"unsupported JSON vocab shape at {vocab_path}: {type(data).__name__}")

    with vocab_path.open("rb") as f:
        data = _VocabUnpickler(f).load()
    if isinstance(data, _VocabState) and isinstance(data.stoi, dict):
        return len(data.stoi)
    if isinstance(data, (dict, list, tuple)):
        return len(data)
    raise ValueError(f"could not count entries in {vocab_path}")


class _VocabState:
    """Data-only stand-in: counting tokens does not require tokenizer methods."""


class _VocabUnpickler(pickle.Unpickler):
    def find_class(self, module: str, name: str) -> Any:
        if module in {"kermt.data.torchvocab", "grover.data.torchvocab"} and name in {
            "TorchVocab", "MolVocab", "SMILESVocab",
        }:
            return _VocabState
        if (module, name) == ("collections", "Counter"):
            return Counter
        if (module, name) == ("re", "_compile"):
            return re.compile
        raise pickle.UnpicklingError(f"unsupported vocabulary object: {module}.{name}")


def load_checkpoint(path: Path | str) -> dict[str, Any]:
    """Read KERMT tensors and known metadata with PyTorch's restricted loader.

    Saved arguments use argparse.Namespace; finetuned checkpoints also contain
    numeric NumPy scaler arrays. Explicit globals cover those formats, including
    NumPy 1/2 module names, without accepting artifact-selected imports.
    """
    import numpy as np
    import torch

    multiarray = np._core.multiarray if hasattr(np, "_core") else np.core.multiarray
    allowed = [argparse.Namespace, np.ndarray, np.dtype]
    for module in ("numpy.core.multiarray", "numpy._core.multiarray"):
        allowed.extend([
            (multiarray._reconstruct, f"{module}._reconstruct"),
            (multiarray.scalar, f"{module}.scalar"),
        ])
    allowed.extend(type(np.dtype(name)) for name in (
        "bool", "int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64",
        "float16", "float32", "float64",
    ))
    with torch.serialization.safe_globals(allowed):
        return torch.load(path, map_location="cpu", weights_only=True)


def runner_environment(repo: Path, *, wandb: bool = False) -> dict[str, str]:
    """Forward named runtime settings, keeping unrelated credentials out of jobs.

    W&B credentials/settings are included only for an explicitly enabled W&B
    run. Hugging Face authentication belongs to the separate download helper.
    """
    names = (
        "PATH", "HOME", "TMPDIR", "TEMP", "TMP", "LANG", "LC_ALL", "LC_CTYPE", "TZ",
        "LD_LIBRARY_PATH", "LIBRARY_PATH", "CUDA_HOME", "CUDA_PATH", "PYTHONPATH",
        "PYTHONDONTWRITEBYTECODE", "PYTHONUNBUFFERED", "PYTHONWARNINGS",
        "CUDA_VISIBLE_DEVICES", "CUDA_DEVICE_ORDER", "CUDA_LAUNCH_BLOCKING", "NVIDIA_VISIBLE_DEVICES",
        "NVIDIA_DRIVER_CAPABILITIES", "CUBLAS_WORKSPACE_CONFIG", "OMP_NUM_THREADS",
        "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_NUM_THREADS",
        "PYTORCH_CUDA_ALLOC_CONF", "PYTORCH_ALLOC_CONF", "PYTORCH_NO_CUDA_MEMORY_CACHING",
        "TORCH_CPP_LOG_LEVEL", "TORCH_DISTRIBUTED_DEBUG", "NCCL_DEBUG", "NCCL_SOCKET_IFNAME",
        "NCCL_IB_DISABLE", "NCCL_P2P_DISABLE", "NCCL_SHM_DISABLE", "GLOO_SOCKET_IFNAME",
        "MASTER_ADDR", "MASTER_PORT",
        "KERMT_REPO", "KERMT_REPO_COMMIT", "KERMT_REPO_DIRTY",
        "SSL_CERT_FILE", "REQUESTS_CA_BUNDLE",
    )
    if wandb:
        names += (
            "WANDB_API_KEY", "WANDB_BASE_URL", "WANDB_MODE", "WANDB_DIR", "WANDB_ENTITY",
            "WANDB_PROJECT", "WANDB_RUN_ID", "WANDB_RESUME", "WANDB_CACHE_DIR",
            "WANDB_CONFIG_DIR", "WANDB_DATA_DIR", "WANDB_DISABLED",
        )
    env = {name: value for name in names if (value := os.environ.get(name)) is not None}
    env["PYTHONPATH"] = os.pathsep.join(filter(None, (str(repo), env.get("PYTHONPATH"))))
    return env


def validate_vocab_file(vocab_path: Path, *, kind: str) -> None:
    """Verify a user-provided vocab file is loadable BEFORE copying it into a
    run directory. Raises ValueError on failure with a clear, user-facing message.

    `kind` is one of {"atom", "bond", "smiles"} — used only in the error message
    so the user knows which file is wrong.
    """
    if not vocab_path.is_file():
        raise FileNotFoundError(f"{kind} vocab file not found: {vocab_path}")
    try:
        n = count_vocab_entries(vocab_path)
    except Exception as exc:  # noqa: BLE001
        raise ValueError(
            f"{kind} vocab file {vocab_path} is not loadable as a KERMT vocab "
            f"({type(exc).__name__}: {exc}). Expected a MolVocab JSON or pickle "
            f"(or a SMILESVocab pickle for the smiles vocab)."
        ) from exc
    if n <= 0:
        raise ValueError(f"{kind} vocab file {vocab_path} contains zero entries")


# ---------------------------------------------------------------------------
# Runner-shared helpers (run.json manifest fields)
# ---------------------------------------------------------------------------

def git_commit_with_env_override(repo: Path) -> tuple[str, bool]:
    """Returns (commit_sha, dirty_tree). Honors `KERMT_REPO_COMMIT` /
    `KERMT_REPO_DIRTY` env vars first — set by `scripts/kermt_container.sh`
    from the host before launching docker (necessary because `git -C /workspace`
    inside the container fails due to bind-mount ownership). Falls back to the
    in-container git probe when the env vars aren't set."""
    env_commit = os.environ.get("KERMT_REPO_COMMIT")
    if env_commit:
        env_dirty = os.environ.get("KERMT_REPO_DIRTY", "false").strip().lower() == "true"
        return env_commit, env_dirty
    try:
        sha = subprocess.run(
            ["git", "-C", str(repo), "rev-parse", "HEAD"],
            capture_output=True, text=True, check=True,
        ).stdout.strip()
        diff = subprocess.run(
            ["git", "-C", str(repo), "status", "--porcelain"],
            capture_output=True, text=True, check=True,
        )
        return sha, bool(diff.stdout.strip())
    except Exception:
        return "unknown", False


def docker_image_digest(tag: str) -> str | None:
    """Return the docker image's content-addressable Id (sha256:…) for the given
    tag, or None if docker isn't available / the image isn't local."""
    try:
        r = subprocess.run(
            ["docker", "image", "inspect", tag, "--format", "{{.Id}}"],
            capture_output=True, text=True,
        )
        if r.returncode == 0:
            return r.stdout.strip()
    except FileNotFoundError:
        pass
    return None


def format_cmd_replay(argv: list[str], *, env: dict[str, str] | None = None) -> str:
    """Render a copy-pasteable env-prefix + command for the cmd_replay manifest
    field. `env` is the set of environment variables to prefix (typically
    {CUDA_VISIBLE_DEVICES, WORLD_SIZE})."""
    env = env or {}
    env_prefix = [f"{k}={shlex.quote(str(v))}" for k, v in env.items()]
    quoted = " ".join(shlex.quote(a) for a in argv)
    return " ".join(env_prefix + [quoted])


def resolve_single_gpu(override: str | None, *, workflow: str) -> int:
    """Returns a single GPU id (int). The finetune/inference/embed workflows are
    single-GPU only; `--gpus '0,1'` or multi-id CUDA_VISIBLE_DEVICES is rejected
    with a workflow-specific error. (The pretrain runner has its own multi-GPU
    `_detect_gpus` helper — see run_pretrain_local.py.)"""
    if override is None:
        env_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
        if env_visible:
            ids = [g for g in env_visible.split(",") if g]
            if len(ids) > 1:
                raise ValueError(
                    f"CUDA_VISIBLE_DEVICES='{env_visible}' selects multiple GPUs but "
                    f"the {workflow} workflow is single-GPU only. Restrict to one id."
                )
            return int(ids[0])
        return 0
    parts = [p.strip() for p in override.split(",") if p.strip()]
    if len(parts) != 1:
        raise ValueError(
            f"--gpus '{override}' selects {len(parts)} GPUs; the {workflow} workflow is single-GPU only."
        )
    return int(parts[0])


def assert_prepare_manifest_basics(manifest: dict[str, Any], expected_mode: str) -> None:
    """Standard pre-check for a prepare_data.json before a runner consumes it:
    verify `mode` matches and `ok` is True. Raises ValueError with a consistent
    error message on either mismatch.

    Each runner is responsible for its own required-outputs check after this
    (those vary per-mode — e.g. pretrain wants train_dir/val_dir/atom_vocab/
    bond_vocab; finetune has the split-method branch; inference/embed want
    clean_csv)."""
    if manifest.get("mode") != expected_mode:
        raise ValueError(
            f"prepare_data manifest is mode='{manifest.get('mode')}', expected '{expected_mode}'. "
            f"Run `prepare_data.py --mode {expected_mode}` to produce a valid manifest."
        )
    if not manifest.get("ok"):
        raise ValueError(
            f"prepare_data manifest reports ok=False: {manifest.get('errors')}"
        )


def merge_default_into_applied(
    applied: dict[str, dict[str, Any]],
    args: argparse.Namespace,
    name: str,
    defaults_group: dict[str, Any],
) -> None:
    """Standard CLI-override / default-config merge for one hyperparameter.

    Mutates `applied` in place:
      - If the user passed `--<name>` on the CLI (so `getattr(args, name)` is
        not None), records `{"value": cli_val, "source": "user"}`.
      - Else if `name` is present in `defaults_group`, records
        `{"value": defaults_group[name], "source": "default-config"}`.
      - Else `applied[name]` is left absent — the runner's argv-builder skips
        the flag, and the downstream argparse default takes effect.

    `name` is the snake_case argparse dest (same form used as the dict key);
    argparse automatically converts CLI `--<name-with-hyphens>` to that dest,
    so `getattr(args, name, None)` is the correct CLI lookup."""
    cli_val = getattr(args, name, None)
    if cli_val is not None:
        applied[name] = {"value": cli_val, "source": "user"}
    elif name in defaults_group:
        applied[name] = {"value": defaults_group[name], "source": "default-config"}


def run_checkpoint_validator(ckpt: Path, *, mode: str, script_path: Path) -> dict[str, Any]:
    """Invoke `check_checkpoint.py --mode <mode> --ckpt <path>` as a subprocess
    and return the parsed JSON. Raises RuntimeError on non-JSON output (e.g. the
    validator crashed before printing). `script_path` is the absolute path to
    `scripts/check_checkpoint.py` — passed in so this helper has no
    dependency on the caller's layout."""
    r = subprocess.run(
        [sys.executable, str(script_path), "--mode", mode, "--ckpt", str(ckpt)],
        capture_output=True, text=True,
    )
    try:
        return json.loads(r.stdout)
    except json.JSONDecodeError as exc:
        raise RuntimeError(
            f"check_checkpoint.py emitted non-JSON output (exit {r.returncode}). "
            f"stdout (first 200 chars): {r.stdout[:200]}\n"
            f"stderr (first 200 chars): {r.stderr[:200]}"
        ) from exc