"""Apply the released adapter to one matched offline record on CPU.

Install the BAS-VLA code repository and its dependencies first. The input action
contract is recorded in adapter/adapter_metadata.json. Use that exact contract
when supplying new actions from a carrier.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch

from bas_vla.breaking.residual_adapter import ResidualAdapter, encode_instruction_bow


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--artifact-dir", type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument("--record-index", type=int, default=160, help="Default: first held-out test record.")
    args = parser.parse_args()
    root = args.artifact_dir
    meta = json.loads((root / "adapter/adapter_metadata.json").read_text())
    vocab = json.loads((root / "adapter/vocab.json").read_text())
    records = json.loads((root / "data/triplets.json").read_text())
    if not 0 <= args.record_index < len(records):
        parser.error("record-index is outside the released records")
    record = records[args.record_index]
    model = ResidualAdapter(
        input_dim=meta["input_dim"], hidden_dim=meta["hidden_dim"],
        output_dim=meta["output_dim"], delta_scale=meta["delta_scale"],
    )
    weights = torch.load(root / "adapter/residual_adapter.pt", map_location="cpu", weights_only=True)
    model.load_state_dict(weights, strict=True)
    model.eval()
    outputs = {}
    with torch.no_grad():
        for variant in ("clean", "control", "break"):
            action = np.asarray(record[f"{variant}_base_action"], dtype=np.float32)
            features = encode_instruction_bow(record[f"{variant}_instruction"], vocab)
            calibrated = model(torch.from_numpy(action[None]), torch.from_numpy(features[None]))[0]
            outputs[variant] = {
                "instruction": record[f"{variant}_instruction"],
                "base_action": action.tolist(), "calibrated_action": calibrated.tolist(),
            }
    print(json.dumps({
        "record_id": record["record_id"], "split": record["split"],
        "action_space": meta["action_space"], "outputs": outputs,
    }, indent=2))


if __name__ == "__main__":
    main()