"""Rename a trained model's choice labels without retraining the policy.

The core only sees choice *embeddings*. To move a model from "UP" to "move up"
we fit the ChoiceEncoder (not the shared embeddings, not the core) so that the
new label produces the vector the old label produced. The decision function is
then unchanged up to the fit error, which is verified on real states.

    minijev relabel --checkpoint runs/x/final.pt --out runs/x/final-semantic.pt \\
        --map UP="move up" DOWN="move down" LEFT="move left" RIGHT="move right" \\
        --set task.choice_style=semantic
"""
from __future__ import annotations

import copy

import numpy as np
import torch


def relabel_checkpoint(checkpoint: str, out: str, mapping: dict[str, str], cfg_updates: dict | None = None,
                       steps: int = 3000, lr: float = 3e-3, anchors: list[str] | None = None) -> dict:
    from minijev.inference.runtime import TorchDecider
    from minijev.training.supervised import make_model

    ck = torch.load(checkpoint, map_location="cpu", weights_only=False)
    cfg = copy.deepcopy(ck["config"])
    model = make_model(cfg)
    model.load_state_dict(ck["model"])
    model.eval()
    old_dec = TorchDecider(model=copy.deepcopy(model), cfg=cfg)
    tok, max_len = old_dec.tokenizer, old_dec.max_choice_len
    old, new = list(mapping), list(mapping.values())
    anchors = list(anchors or [])

    def emb(labels):
        if model.backbone is None:
            ids = torch.as_tensor(tok.encode_batch(labels, max_len))
            return model.choice_encoder(ids)
        ids = torch.as_tensor(tok.encode_batch(labels, max_len, dynamic=True))
        return model.choice_encoder(pooled=model.backbone(ids)[2])

    with torch.no_grad():
        target = emb(old + anchors).detach()
    shared = {id(p) for n, p in model.named_parameters() if not n.startswith("choice_encoder.")}
    params = [p for p in model.choice_encoder.parameters() if id(p) not in shared]
    for p in model.parameters():
        p.requires_grad_(False)
    for p in params:
        p.requires_grad_(True)
    opt = torch.optim.Adam(params, lr=lr)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, steps)
    for _ in range(steps):
        loss = (emb(new + anchors) - target).pow(2).mean() / target.pow(2).mean()
        opt.zero_grad()
        loss.backward()
        opt.step()
        sched.step()
    for p in model.parameters():
        p.requires_grad_(True)

    for k, v in (cfg_updates or {}).items():
        node = cfg
        parts = k.split(".")
        for part in parts[:-1]:
            node = node.setdefault(part, {})
        node[parts[-1]] = v
    new_dec = TorchDecider(model=model, cfg=cfg)
    new_dec.temperature = old_dec.temperature
    states = new_dec.task.example_states(256, seed=123)
    p_old = old_dec.probs(states, new_dec.task.default_question or "", old)
    p_new = new_dec.probs(states, new_dec.task.default_question or "", new)
    report = {"rel_mse": float(loss.detach()), "max_abs_prob_diff": float(np.abs(p_old - p_new).max()),
              "argmax_agreement": float((p_old.argmax(1) == p_new.argmax(1)).mean()), "states": len(states)}
    torch.save({**{k: v for k, v in ck.items() if k not in ("model", "config", "optimizer")},
                "model": model.state_dict(), "config": cfg, "relabel": {"from": checkpoint, "map": mapping, **report}},
               out)
    return report