"""Idempotently graft the MAE objective into the eatmap tree on the VM.

Kept here rather than applied by hand because /workspace is not a volume on this
instance -- workspace_is_volume is False, so a recycle wipes the container
filesystem. This script plus mae.py is the whole change; re-running it on a
fresh box reproduces the tree.

Edits are surgical and leave every EAT code path byte-identical: the objective
dispatch defaults to "ufo", and the two structural fixes (teacher-optional
export, teacher-aware parameter count) are no-ops when a teacher is present.
"""
import sys
from pathlib import Path

ROOT = Path(sys.argv[1] if len(sys.argv) > 1 else "/workspace/code/eat-map-regmix/eatmap")


def patch(path: Path, old: str, new: str, label: str) -> None:
    text = path.read_text()
    if new in text:
        print(f"  [skip] {label} (already applied)")
        return
    if old not in text:
        raise SystemExit(f"  [FAIL] {label}: anchor not found in {path}")
    path.write_text(text.replace(old, new, 1))
    print(f"  [ok]   {label}")


print(f"patching {ROOT}")

# ---- 1. ObjectiveConfig: objective selector + MAE decoder geometry ----------
patch(
    ROOT / "config.py",
    """    # null/0 means "anneal over total_steps" -- never leave this at a value
    # larger than the budget or the EMA never reaches its end decay.
    ema_anneal_steps: int | None = None""",
    """    # null/0 means "anneal over total_steps" -- never leave this at a value
    # larger than the budget or the EMA never reaches its end decay.
    ema_anneal_steps: int | None = None
    # "ufo" is EAT's EMA-teacher objective and the default, so every existing
    # config keeps its meaning. "mae" swaps in masked log-mel reconstruction
    # over the same encoder -- see eatmap/mae.py for why that sharing matters.
    type: str = "ufo"
    # Decoder geometry, MAE only. Asymmetric by design: the decoder is thrown
    # away after pretraining, so it is narrow and shallow relative to the
    # encoder. Defaults are sized for the 384-dim 15M proxy.
    mae_decoder_dim: int = 256
    mae_decoder_depth: int = 4
    mae_decoder_heads: int = 8
    # Per-patch target standardization (He et al.'s norm_pix_loss). Off makes
    # the loss track patch loudness rather than structure.
    mae_norm_pix: bool = True""",
    "ObjectiveConfig: type + MAE decoder fields",
)

# ---- 2. export_weights: teacher is optional --------------------------------
patch(
    ROOT / "runner.py",
    """    tensors = {}
    for name, tensor in model.student.state_dict().items():
        tensors[f"student.{name}"] = tensor.detach().cpu().contiguous()
    for name, tensor in model.teacher.state_dict().items():
        tensors[f"teacher.{name}"] = tensor.detach().cpu().contiguous()
    for name, tensor in model.decoder.state_dict().items():
        tensors[f"decoder.{name}"] = tensor.detach().cpu().contiguous()""",
    """    tensors = {}
    # MAE has no teacher, so the parts are discovered rather than assumed. The
    # student is written under the same "student." prefix either way, which is
    # what probe.load_encoder keys on -- an MAE export and an EAT export are
    # interchangeable to every readout.
    for part in ("student", "teacher", "decoder"):
        module = getattr(model, part, None)
        if module is None:
            continue
        for name, tensor in module.state_dict().items():
            tensors[f"{part}.{name}"] = tensor.detach().cpu().contiguous()""",
    "export_weights: teacher-optional",
)

# ---- 3. objective dispatch -------------------------------------------------
patch(
    ROOT / "runner.py",
    "from .model import EATPretrainer, parameter_counts",
    "from .mae import MAEPretrainer\nfrom .model import EATPretrainer, parameter_counts",
    "runner: import MAEPretrainer",
)
patch(
    ROOT / "runner.py",
    "    model = EATPretrainer(config).to(device)",
    """    pretrainer = {"ufo": EATPretrainer, "mae": MAEPretrainer}.get(config.objective.type)
    if pretrainer is None:
        raise ValueError(f"unknown objective.type {config.objective.type!r}")
    model = pretrainer(config).to(device)""",
    "runner: objective dispatch",
)

# ---- 4. parameter_counts: only EAT carries a frozen twin -------------------
patch(
    ROOT / "model.py",
    """        "total_params": student * 2 + decoder,  # teacher is a frozen copy""",
    """        # EAT carries a frozen EMA twin of the student; MAE does not.
        "total_params": student * (2 if getattr(model, "teacher", None) is not None else 1) + decoder,""",
    "parameter_counts: teacher-aware total",
)

print("done")