runtime/src/resynthesis/geometry_migration.py

"""Checkpointed NoNE expert/layer geometry growth.

This module is an explicit checkpoint I/O boundary, not a model hot path.
It preserves every trained source tensor region exactly, adds trainable
expert/layer capacity deterministically, migrates named AdamW moments, and
writes a non-promotable candidate receipt. The migrated graph must still train,
cold reload, and pass held-out proof before it can replace a promoted state.
"""
from __future__ import annotations

import copy
import hashlib
import json
import os
import re
import time
from pathlib import Path
from typing import Any, cast

import torch

MIGRATION_SCHEMA = "nnf.resynthesis.none_geometry_migration.v1"
CHECKPOINT_SCHEMA = "nnf.resynthesis.additive_state.v1"
GROWTH_PLAN_SCHEMA = "nnf.resynthesis.none_growth_plan.v2"
LEGACY_GROWTH_PLAN_SCHEMA = "nnf.resynthesis.none_growth_plan.v1"
GROWTH_PLAN_SCHEMAS = frozenset(
    {GROWTH_PLAN_SCHEMA, LEGACY_GROWTH_PLAN_SCHEMA}
)
STRUCTURAL_EXPERTS = 4

_LAYER_PATTERN = re.compile(r"^science_stack\.science_layer_(\d+)\.(.+)$")

_GLOBAL_EXPERT_ROWS = frozenset(
    {
        "correction_expert_head.weight",
        "fabric.decision_proj.weight",
        "fabric.expert_bias_table",
        "fabric.expert_hidden_table",
        "fabric.intent_table",
    }
)
_GLOBAL_LAYER_ROWS = frozenset(
    {
        "correction_layer_head.weight",
        "fabric.layer_bias_table",
        "fabric.layer_hidden_table",
        "science_stack.layer_identity_glyphs",
        "science_stack.traversal_gate",
    }
)
_GLOBAL_LAYER_SQUARE = frozenset({"science_stack.layer_transfer_graph"})
_LAYER_EXPERT_ROWS = frozenset(
    {
        "_expert_history_states",
        "expert_activation_prior",
        "expert_depth_pref",
        "expert_intent_glyphs",
        "expert_role_tag",
        "expert_specialization",
        "expert_transfer_affinity",
        "router.weight",
    }
)
_LAYER_EXPERT_SQUARE = frozenset({"expert_compatibility"})
_LAYER_FFN_ROWS = frozenset({"ffn_down", "ffn_gate_up"})

_STACK_PARAMETER_ORDER = (
    "science_stack.traversal_gate",
    "science_stack.layer_rotation_pressure",
    "science_stack.layer_transfer_graph",
    "science_stack.layer_transfer_scale",
    "science_stack.logit_residual_scale",
    "science_stack.layer_identity_glyphs",
    "science_stack.layer_identity_scale",
    "science_stack.long_context_anchor_gain",
    "science_stack.glyph_projection.weight",
    "science_stack.layer_identity_query_proj.weight",
    "science_stack.long_context_anchor_query.weight",
)
_LAYER_PARAMETER_ORDER = (
    "expert_activation_prior",
    "expert_intent_glyphs",
    "language_match_scale",
    "expert_role_tag",
    "expert_specialization",
    "role_match_scale",
    "expert_compatibility",
    "expert_transfer_scale",
    "expert_rotation_pressure",
    "expert_depth_pref",
    "expert_transfer_affinity",
    "layer_depth_signal",
    "layer_complexity",
    "memory_bank",
    "mhc_distinct_hypothesis_scale",
    "ffn_gate_up",
    "ffn_down",
    "residual_scale",
    "translate_scale",
    "audit_scale",
    "norm.weight",
    "norm.bias",
    "output_norm.weight",
    "output_norm.bias",
    "router.weight",
    "intent_query_proj.weight",
    "role_query_proj.weight",
    "expert_capability_proj.weight",
    "expert_capability_proj.bias",
    "capability_match_scale",
    "expert_history_gru.weight_ih",
    "expert_history_gru.weight_hh",
    "expert_history_gru.bias_ih",
    "expert_history_gru.bias_hh",
    "layer_role_head.weight",
    "layer_role_head.bias",
    "recurrent_expert.weight_ih_l0",
    "recurrent_expert.weight_hh_l0",
    "recurrent_expert.bias_ih_l0",
    "recurrent_expert.bias_hh_l0",
    "attention_expert.intent_pivot_scale",
    "attention_expert.action_pivot_scale",
    "attention_expert.context_query_pivot_scale",
    "attention_expert.relation_connectivity_scale",
    "attention_expert.q_proj.weight",
    "attention_expert.q_proj.bias",
    "attention_expert.k_proj.weight",
    "attention_expert.k_proj.bias",
    "attention_expert.v_proj.weight",
    "attention_expert.v_proj.bias",
    "attention_expert.c_proj.weight",
    "attention_expert.c_proj.bias",
    "attention_expert.intent_c_proj.weight",
    "attention_expert.action_c_proj.weight",
    "attention_expert.r_query_proj.weight",
    "attention_expert.r_query_proj.bias",
    "attention_expert.r_key_proj.weight",
    "attention_expert.r_key_proj.bias",
    "attention_expert.intent_r_query_proj.weight",
    "attention_expert.intent_r_key_proj.weight",
    "attention_expert.out_proj.weight",
    "attention_expert.out_proj.bias",
    "action_glyph_bridge.weight",
    "memory_query.weight",
    "memory_out.weight",
    "glyph_proj.weight",
    "glyph_gate.weight",
    "glyph_translate_proj.weight",
    "glyph_translate_back.weight",
)
_V12_LAYER_PARAMETER_ORDER = tuple(
    name
    for name in _LAYER_PARAMETER_ORDER
    if name
    not in {
        "attention_expert.action_pivot_scale",
        "attention_expert.action_c_proj.weight",
        "action_glyph_bridge.weight",
    }
)
_LEGACY_LAYER_PARAMETER_BASE = tuple(
    name
    for name in _LAYER_PARAMETER_ORDER
    if name != "mhc_distinct_hypothesis_scale"
)
_LEGACY_LAYER_PARAMETER_ORDER = (
    *_LEGACY_LAYER_PARAMETER_BASE[
        : _LEGACY_LAYER_PARAMETER_BASE.index(
            "attention_expert.intent_pivot_scale"
        )
    ],
    "attention_expert.in_proj_weight",
    "attention_expert.in_proj_bias",
    "attention_expert.out_proj.weight",
    "attention_expert.out_proj.bias",
    *_LEGACY_LAYER_PARAMETER_BASE[
        _LEGACY_LAYER_PARAMETER_BASE.index("memory_query.weight") :
    ],
)
_TAIL_PARAMETER_ORDER = (
    "stop_gate.trajectory_proj.weight",
    "stop_gate.lstm.weight_ih_l0",
    "stop_gate.lstm.weight_hh_l0",
    "stop_gate.lstm.bias_ih_l0",
    "stop_gate.lstm.bias_hh_l0",
    "stop_gate.stop_utility_gate.weight",
    "stop_gate.stop_utility_gate.bias",
    "stop_gate.stop_contradiction_gate.weight",
    "stop_gate.stop_contradiction_gate.bias",
    "feedback_head.weight",
    "feedback_head.bias",
    "prior_hidden_proj.weight",
    "outcome_encoder.weight",
    "parent_outcome_encoder.weight",
    "acquisition_encoder.weight",
    "acquisition_policy.input_norm.weight",
    "acquisition_policy.input_norm.bias",
    "acquisition_policy.context.weight",
    "acquisition_policy.context.bias",
    "acquisition_policy.action_head.weight",
    "acquisition_policy.action_head.bias",
    "correction_context_norm.weight",
    "correction_context_norm.bias",
    "correction_hidden_up.weight",
    "correction_trigger_head.weight",
    "correction_trigger_head.bias",
    "task_confidence_head.weight",
    "task_confidence_head.bias",
    "delegation_head.weight",
    "delegation_head.bias",
    "correction_expert_head.weight",
    "correction_layer_head.weight",
    "logit_residual_down.weight",
    "logit_residual_up.weight",
    "fabric.expert_bias_table",
    "fabric.layer_bias_table",
    "fabric.intent_table",
    "fabric.domain_residency",
    "fabric.transfer_table",
    "fabric.expert_hidden_table",
    "fabric.layer_hidden_table",
    "fabric.phase_hidden_table",
    "fabric.residual_gate",
    "fabric.parent_expert_route_gate",
    "fabric.parent_layer_route_gate",
    "fabric.phase_proj.weight",
    "fabric.decision_proj.weight",
    "legacy_capability_bank.fusion_logit",
    "legacy_capability_bank.route_query.weight",
    "legacy_capability_bank.outcome_query.weight",
    "legacy_capability_bank.residual_projection.weight",
)


def _file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _state_identity(state: dict[str, torch.Tensor]) -> tuple[str, str]:
    names = sorted(state)
    key_hash = hashlib.sha256("\n".join(names).encode("utf-8")).hexdigest()
    geometry = [(name, tuple(state[name].shape), str(state[name].dtype)) for name in names]
    geometry_hash = hashlib.sha256(
        json.dumps(geometry, separators=(",", ":")).encode("utf-8")
    ).hexdigest()
    return key_hash, geometry_hash


def _atomic_torch_save(payload: object, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp")
    temporary.unlink(missing_ok=True)
    torch.save(payload, temporary)
    with temporary.open("rb") as handle:
        os.fsync(handle.fileno())
    os.replace(temporary, path)


def _atomic_json(payload: dict[str, Any], path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp")
    temporary.unlink(missing_ok=True)
    with temporary.open("w", encoding="utf-8") as handle:
        json.dump(payload, handle, sort_keys=True, indent=2)
        handle.write("\n")
        handle.flush()
        os.fsync(handle.fileno())
    os.replace(temporary, path)


def _checkpoint_payload(path: Path, *, map_location: str) -> dict[str, Any]:
    payload = torch.load(path, map_location=map_location, mmap=True, weights_only=True)
    if not isinstance(payload, dict) or payload.get("schema") != CHECKPOINT_SCHEMA:
        raise RuntimeError("Resynthesis geometry source is not an additive checkpoint")
    lineage = payload.get("lineage")
    parameters = payload.get("parameters")
    buffers = payload.get("buffers")
    if (
        not isinstance(lineage, dict)
        or not isinstance(parameters, dict)
        or not isinstance(buffers, dict)
        or not all(isinstance(value, torch.Tensor) for value in parameters.values())
        or not all(isinstance(value, torch.Tensor) for value in buffers.values())
    ):
        raise RuntimeError("Resynthesis geometry source checkpoint is incomplete")
    state = {**parameters, **buffers}
    key_hash, geometry_hash = _state_identity(state)
    if payload.get("stateKeySetSha256") != key_hash:
        raise RuntimeError("Resynthesis geometry source key identity differs")
    if payload.get("stateGeometrySha256") != geometry_hash:
        raise RuntimeError("Resynthesis geometry source tensor geometry differs")
    return payload


def checkpoint_geometry(path: str | Path) -> tuple[int, int, bool]:
    """Read additive geometry without allocating checkpoint tensor storage."""

    payload = _checkpoint_payload(Path(path), map_location="meta")
    lineage = payload["lineage"]
    layers = int(lineage.get("scienceLayers", 0))
    experts = int(lineage.get("scienceExperts", 0))
    migrated = bool(lineage.get("draftingCheckpointGeometryChanged", False))
    if layers < 1 or experts < STRUCTURAL_EXPERTS + 1:
        raise RuntimeError("Resynthesis additive checkpoint geometry is invalid")
    if int(lineage.get("parallelDraftWorkers", 0)) != experts:
        raise RuntimeError("Resynthesis additive drafting geometry differs")
    return layers, experts, migrated


def validate_migration_candidate(
    checkpoint_path: str | Path,
    receipt_path: str | Path,
) -> dict[str, Any]:
    """Validate model, optimizer, lineage, and receipt at training boundary."""

    checkpoint = Path(checkpoint_path).resolve()
    receipt_file = Path(receipt_path).resolve()
    receipt = json.loads(receipt_file.read_text(encoding="utf-8"))
    if not isinstance(receipt, dict) or receipt.get("schema") != MIGRATION_SCHEMA:
        raise RuntimeError("NoNE geometry candidate receipt schema differs")
    if receipt.get("passed") is not True or receipt.get("promotionEligible") is not False:
        raise RuntimeError("NoNE geometry candidate authority is invalid")
    target = receipt.get("targetCheckpoint")
    if not isinstance(target, dict):
        raise RuntimeError("NoNE geometry candidate receipt has no target checkpoint")
    if Path(str(target.get("path", ""))).resolve() != checkpoint:
        raise RuntimeError("NoNE geometry candidate path differs from its receipt")
    if target.get("sha256") != _file_sha256(checkpoint):
        raise RuntimeError("NoNE geometry candidate SHA-256 differs from its receipt")
    payload = _checkpoint_payload(checkpoint, map_location="meta")
    lineage = payload["lineage"]
    if (
        lineage.get("schema") != "nnf.resynthesis.composed_additive_lineage.v13"
        or lineage.get("intentContextPivotAttention") is not True
        or lineage.get("contextIntentActionAttention") is not True
        or lineage.get("contextActionSource")
        != "trained_acquisition_policy_probability_tensor"
        or lineage.get("contextActionDim") != 4
        or lineage.get("contextActionScorePivot") is not True
        or lineage.get("contextActionCheckpointGeometryChanged") is not True
        or lineage.get("intentRelationalAttention") is not True
        or tuple(lineage.get("attentionMultiples", ()))
        != ("q", "k", "v", "c", "r")
        or lineage.get("scienceAttentionExactTiling") is not True
        or lineage.get("contextRelationCheckpointGeometryChanged") is not True
    ):
        raise RuntimeError(
            "NoNE geometry candidate intent/action C/R lineage differs"
        )
    parameters = payload["parameters"]
    parameter_elements = sum(tensor.numel() for tensor in parameters.values())
    if int(target.get("parameterElements", -1)) != parameter_elements:
        raise RuntimeError("NoNE geometry candidate parameter count differs")
    optimizer_record = receipt.get("targetOptimizer")
    if not isinstance(optimizer_record, dict):
        raise RuntimeError("NoNE geometry candidate has no optimizer authority")
    optimizer_path = checkpoint.with_suffix(".optimizer.pt")
    if Path(str(optimizer_record.get("path", ""))).resolve() != optimizer_path:
        raise RuntimeError("NoNE geometry optimizer path differs from its receipt")
    if not optimizer_path.is_file() or optimizer_record.get("sha256") != _file_sha256(
        optimizer_path
    ):
        raise RuntimeError("NoNE geometry optimizer SHA-256 differs from its receipt")
    optimizer = torch.load(
        optimizer_path,
        map_location="meta",
        mmap=True,
        weights_only=True,
    )
    if not isinstance(optimizer, dict):
        raise RuntimeError("NoNE geometry optimizer payload is invalid")
    groups = optimizer.get("param_groups")
    states = optimizer.get("state")
    if not isinstance(groups, list) or len(groups) != 1 or not isinstance(states, dict):
        raise RuntimeError("NoNE geometry optimizer groups are invalid")
    names = groups[0].get("param_names")
    parameter_ids = groups[0].get("params")
    if (
        not isinstance(names, (list, tuple))
        or not isinstance(parameter_ids, (list, tuple))
        or len(names) != len(parameter_ids)
        or set(str(name) for name in names) != set(parameters)
    ):
        raise RuntimeError("NoNE geometry optimizer names differ from the model")
    for name, parameter_id in zip(names, parameter_ids, strict=True):
        parameter_state = states.get(parameter_id)
        if not isinstance(parameter_state, dict):
            continue
        for moment_name in ("exp_avg", "exp_avg_sq", "max_exp_avg_sq"):
            moment = parameter_state.get(moment_name)
            if isinstance(moment, torch.Tensor) and moment.shape != parameters[
                str(name)
            ].shape:
                raise RuntimeError(
                    f"NoNE geometry optimizer moment differs for {name}"
                )
    remaining = receipt.get("remainingProof")
    if not isinstance(remaining, list) or "continued_training" not in remaining:
        raise RuntimeError("NoNE geometry candidate omits continued-training authority")
    checkpoint_geometry(checkpoint)
    return receipt


def _expanded_rows(source: torch.Tensor, target_rows: int) -> torch.Tensor:
    source_rows = source.shape[0]
    if target_rows < source_rows or source_rows < 1:
        raise RuntimeError("geometry migration cannot shrink or clone an empty tensor")
    if target_rows == source_rows:
        return source.clone()
    indices = torch.arange(
        target_rows - source_rows,
        device=source.device,
        dtype=torch.long,
    ).remainder(source_rows)
    appended = source.index_select(0, indices).clone()
    if appended.is_floating_point():
        offsets = torch.arange(
            1,
            appended.shape[0] + 1,
            device=appended.device,
            dtype=torch.float32,
        )
        scale_shape = (appended.shape[0],) + (1,) * (appended.ndim - 1)
        scale = (1.0 + offsets.remainder(7).reshape(scale_shape) / 128.0).to(
            dtype=appended.dtype
        )
        appended.mul_(scale)
    return torch.cat((source.clone(), appended), dim=0)


def _expanded_square(source: torch.Tensor, target_width: int) -> torch.Tensor:
    if source.ndim != 2 or source.shape[0] != source.shape[1]:
        raise RuntimeError("geometry transfer tensor is not square")
    source_width = source.shape[0]
    if target_width < source_width or source_width < 1:
        raise RuntimeError("geometry migration cannot shrink a transfer tensor")
    if target_width == source_width:
        return source.clone()
    indices = torch.arange(
        target_width,
        device=source.device,
        dtype=torch.long,
    ).remainder(source_width)
    expanded = source.index_select(0, indices).index_select(1, indices).clone()
    expanded[:source_width, :source_width].copy_(source)
    if expanded.is_floating_point():
        diagonal = torch.arange(source_width, target_width, device=source.device)
        expanded[diagonal, diagonal] += expanded.new_tensor(1.0 / 128.0)
    return expanded


def _migrate_tensor(
    name: str,
    source: torch.Tensor,
    *,
    source_layers: int,
    source_experts: int,
    target_layers: int,
    target_experts: int,
) -> torch.Tensor:
    del source_layers
    if name in _GLOBAL_EXPERT_ROWS:
        return _expanded_rows(source, target_experts)
    if name in _GLOBAL_LAYER_ROWS:
        return _expanded_rows(source, target_layers)
    if name in _GLOBAL_LAYER_SQUARE:
        return _expanded_square(source, target_layers)
    match = _LAYER_PATTERN.match(name)
    if match is None:
        return source.clone()
    suffix = match.group(2)
    if suffix in _LAYER_EXPERT_ROWS:
        return _expanded_rows(source, target_experts)
    if suffix in _LAYER_EXPERT_SQUARE:
        return _expanded_square(source, target_experts)
    if suffix in _LAYER_FFN_ROWS:
        return _expanded_rows(source, target_experts - STRUCTURAL_EXPERTS)
    return source.clone()


def _migrate_state(
    source: dict[str, torch.Tensor],
    *,
    source_layers: int,
    source_experts: int,
    target_layers: int,
    target_experts: int,
) -> dict[str, torch.Tensor]:
    target = {
        name: _migrate_tensor(
            name,
            tensor,
            source_layers=source_layers,
            source_experts=source_experts,
            target_layers=target_layers,
            target_experts=target_experts,
        )
        for name, tensor in source.items()
    }
    layer_sources: dict[int, list[tuple[str, str, torch.Tensor]]] = {}
    for name, tensor in source.items():
        match = _LAYER_PATTERN.match(name)
        if match is None:
            continue
        layer_sources.setdefault(int(match.group(1)), []).append(
            (name, match.group(2), tensor)
        )
    if set(layer_sources) != set(range(source_layers)):
        raise RuntimeError("source checkpoint science layers are not contiguous")
    for target_layer in range(source_layers, target_layers):
        source_layer = target_layer % source_layers
        for _name, suffix, tensor in layer_sources[source_layer]:
            target_name = f"science_stack.science_layer_{target_layer}.{suffix}"
            migrated = _migrate_tensor(
                target_name,
                tensor,
                source_layers=source_layers,
                source_experts=source_experts,
                target_layers=target_layers,
                target_experts=target_experts,
            )
            if migrated.is_floating_point() and migrated.ndim > 0:
                migrated.mul_(migrated.new_tensor(1.0 + (target_layer + 1) / 512.0))
            target[target_name] = migrated
    return dict(sorted(target.items()))


def _preserved_prefix(source: torch.Tensor, target: torch.Tensor) -> bool:
    if source.ndim != target.ndim or any(
        source.shape[axis] > target.shape[axis] for axis in range(source.ndim)
    ):
        return False
    slices = tuple(slice(0, width) for width in source.shape)
    return torch.equal(source, target[slices])


def _legacy_qkv_preserved(
    source: dict[str, torch.Tensor],
    adapted: dict[str, torch.Tensor],
) -> bool:
    """Verify split Q/K/V tensors recompose every trained legacy projection."""

    found = False
    for name, tensor in source.items():
        if not name.endswith("attention_expert.in_proj_weight"):
            continue
        found = True
        prefix = name[: -len("in_proj_weight")]
        qkv = tuple(
            adapted.get(f"{prefix}{projection}_proj.weight")
            for projection in ("q", "k", "v")
        )
        if not all(isinstance(value, torch.Tensor) for value in qkv):
            return False
        if not torch.equal(
            torch.cat(cast(tuple[torch.Tensor, ...], qkv), dim=0),
            tensor,
        ):
            return False
        bias_name = f"{prefix}in_proj_bias"
        source_bias = source.get(bias_name)
        if isinstance(source_bias, torch.Tensor):
            qkv_bias = tuple(
                adapted.get(f"{prefix}{projection}_proj.bias")
                for projection in ("q", "k", "v")
            )
            if not all(isinstance(value, torch.Tensor) for value in qkv_bias):
                return False
            if not torch.equal(
                torch.cat(cast(tuple[torch.Tensor, ...], qkv_bias), dim=0),
                source_bias,
            ):
                return False
    return found


def _parameter_order(parameter_names: set[str], layers: int) -> list[str]:
    order = list(_STACK_PARAMETER_ORDER)
    layer_order: tuple[str, ...]
    if any(
        name.endswith("attention_expert.in_proj_weight")
        for name in parameter_names
    ):
        layer_order = _LEGACY_LAYER_PARAMETER_ORDER
    elif any(
        name.endswith("attention_expert.action_c_proj.weight")
        for name in parameter_names
    ):
        layer_order = _LAYER_PARAMETER_ORDER
    else:
        layer_order = _V12_LAYER_PARAMETER_ORDER
    if not any(
        name.endswith(".mhc_distinct_hypothesis_scale")
        for name in parameter_names
    ):
        layer_order = tuple(
            name
            for name in layer_order
            if name != "mhc_distinct_hypothesis_scale"
        )
    if not any(
        name.endswith(".capability_match_scale")
        for name in parameter_names
    ):
        layer_order = tuple(
            name
            for name in layer_order
            if name != "capability_match_scale"
        )
    for layer in range(layers):
        prefix = f"science_stack.science_layer_{layer}."
        order.extend(prefix + suffix for suffix in layer_order)
    order.extend(_TAIL_PARAMETER_ORDER)
    if set(order) != parameter_names:
        missing = sorted(parameter_names - set(order))
        unexpected = sorted(set(order) - parameter_names)
        raise RuntimeError(
            "optimizer name recovery differs from checkpoint parameters: "
            f"unmapped={missing} absent={unexpected}"
        )
    return order


def _expanded_named_order(
    source_names: list[str],
    *,
    source_layers: int,
    target_layers: int,
) -> list[str]:
    target = list(source_names)
    layer_positions = [
        index
        for index, name in enumerate(source_names)
        if (match := _LAYER_PATTERN.match(name)) is not None
        and int(match.group(1)) == source_layers - 1
    ]
    if not layer_positions:
        raise RuntimeError("optimizer parameter names contain no final science layer")
    insert_at = max(layer_positions) + 1
    additions: list[str] = []
    suffixes = [
        match.group(2)
        for name in source_names
        if (match := _LAYER_PATTERN.match(name)) is not None
        and int(match.group(1)) == 0
    ]
    for layer in range(source_layers, target_layers):
        additions.extend(f"science_stack.science_layer_{layer}.{suffix}" for suffix in suffixes)
    target[insert_at:insert_at] = additions
    return target


def _migrate_optimizer(
    source: dict[str, Any],
    *,
    source_parameters: dict[str, torch.Tensor],
    target_parameters: dict[str, torch.Tensor],
    source_layers: int,
    source_experts: int,
    target_layers: int,
    target_experts: int,
) -> tuple[dict[str, Any], int, int]:
    groups = source.get("param_groups")
    states = source.get("state")
    if not isinstance(groups, list) or len(groups) != 1 or not isinstance(states, dict):
        raise RuntimeError("geometry migration requires one AdamW parameter group")
    source_group = groups[0]
    source_ids = list(source_group.get("params", ()))
    names_value = source_group.get("param_names")
    if isinstance(names_value, (list, tuple)):
        source_names = [str(name) for name in names_value]
        if set(source_names) != set(source_parameters):
            raise RuntimeError("optimizer parameter names differ from checkpoint")
    else:
        # Snapshot parameter dictionaries are canonically name-sorted, while
        # pre-v25 optimizers followed module registration order. Reconstruct
        # that exact checkpointed architecture order; never infer by shape.
        source_names = _parameter_order(set(source_parameters), source_layers)
    if len(source_ids) != len(source_names):
        raise RuntimeError("optimizer parameter IDs differ from named parameters")
    source_by_name = dict(zip(source_names, source_ids, strict=True))
    target_names = list(target_parameters)
    target_state: dict[int, dict[str, Any]] = {}
    copied_states = 0
    expanded_states = 0
    projection_aliases = {
        ".q_proj.weight": (".in_proj_weight", 0),
        ".k_proj.weight": (".in_proj_weight", 1),
        ".v_proj.weight": (".in_proj_weight", 2),
        ".q_proj.bias": (".in_proj_bias", 0),
        ".k_proj.bias": (".in_proj_bias", 1),
        ".v_proj.bias": (".in_proj_bias", 2),
    }
    for target_id, name in enumerate(target_names):
        layer_match = _LAYER_PATTERN.match(name)
        if (
            layer_match is not None
            and int(layer_match.group(1)) >= source_layers
        ):
            continue
        source_name = name
        source_id = source_by_name.get(source_name)
        projection_slice: int | None = None
        if source_id is None:
            for suffix, (legacy_suffix, part) in projection_aliases.items():
                if name.endswith(suffix):
                    source_name = name[: -len(suffix)] + legacy_suffix
                    source_id = source_by_name.get(source_name)
                    projection_slice = part
                    break
        if source_id not in states:
            continue
        source_state = states[source_id]
        if not isinstance(source_state, dict):
            raise RuntimeError("optimizer parameter state is invalid")
        migrated_state: dict[str, Any] = {}
        for key, value in source_state.items():
            if not isinstance(value, torch.Tensor) or value.ndim == 0:
                migrated_state[key] = value.clone() if isinstance(value, torch.Tensor) else copy.deepcopy(value)
                continue
            source_parameter = source_parameters[source_name]
            if tuple(value.shape) != tuple(source_parameter.shape):
                raise RuntimeError(
                    f"optimizer moment geometry differs for {source_name}"
                )
            if projection_slice is not None:
                width = target_parameters[name].shape[0]
                if value.shape[0] != 3 * width:
                    raise RuntimeError(
                        f"legacy attention optimizer moment differs for {name}"
                    )
                value = value[
                    projection_slice * width : (projection_slice + 1) * width
                ].clone()
            migrated_value = _migrate_tensor(
                name,
                value,
                source_layers=source_layers,
                source_experts=source_experts,
                target_layers=target_layers,
                target_experts=target_experts,
            )
            if migrated_value.shape != target_parameters[name].shape:
                raise RuntimeError(f"migrated optimizer moment geometry differs for {name}")
            if migrated_value.shape != value.shape:
                slices = tuple(slice(0, width) for width in value.shape)
                migrated_value.zero_()
                migrated_value[slices].copy_(value)
                expanded_states += 1
            migrated_state[key] = migrated_value
        target_state[target_id] = migrated_state
        copied_states += 1
    target_group = {
        key: copy.deepcopy(value)
        for key, value in source_group.items()
        if key not in {"params", "param_names"}
    }
    target_group["params"] = list(range(len(target_names)))
    target_group["param_names"] = target_names
    return {"state": target_state, "param_groups": [target_group]}, copied_states, expanded_states


def migrate_geometry_checkpoint(
    source_checkpoint: str | Path,
    growth_plan: str | Path,
    output_checkpoint: str | Path,
    receipt_path: str | Path,
    *,
    source_optimizer: str | Path | None = None,
    output_optimizer: str | Path | None = None,
) -> dict[str, Any]:
    """Create a trained-state-preserving, non-promoted growth candidate."""

    source_path = Path(source_checkpoint).resolve()
    plan_path = Path(growth_plan).resolve()
    output_path = Path(output_checkpoint).resolve()
    receipt = Path(receipt_path).resolve()
    if source_path == output_path:
        raise ValueError("geometry migration output must differ from its source")
    plan = json.loads(plan_path.read_text(encoding="utf-8"))
    if (
        not isinstance(plan, dict)
        or plan.get("schema") not in GROWTH_PLAN_SCHEMAS
    ):
        raise RuntimeError("NoNE growth plan schema differs")
    if isinstance(plan.get("pagingPolicy"), dict):
        raise RuntimeError(
            "paged NoNE growth plans require immutable page-generation "
            "migration, not monolithic geometry cloning"
        )
    target_geometry = plan.get("proposedMinimumTargetGeometry")
    if not isinstance(target_geometry, dict):
        raise RuntimeError("NoNE growth plan has no target geometry")
    payload = _checkpoint_payload(source_path, map_location="cpu")
    lineage = payload["lineage"]
    original_source_parameters = payload["parameters"]
    source_layers = int(lineage.get("scienceLayers", 0))
    source_experts = int(lineage.get("scienceExperts", 0))
    target_layers = int(target_geometry.get("scienceLayers", 0))
    target_experts = int(target_geometry.get("scienceExperts", 0))
    if target_layers <= source_layers or target_experts <= source_experts:
        raise RuntimeError("NoNE migration target must expand both layers and experts")
    from resynthesis.science_layers import (
        SCIENCE_ATTENTION_TILE_TOKENS,
        adapt_attention_state_to_context_relation,
    )
    from resynthesis.config import (
        DUAL_CHUNK_LOCAL_SIZE,
        DUAL_CHUNK_PRETRAIN_LENGTH,
        NATIVE_ATTENTION_POSITION_APERTURE,
        PRETRAINED_ROPE_BAND_TOKENS,
        RESYNTHESIS_NATIVE_PREFILL_TILE_TOKENS,
        RESYNTHESIS_ONLINE_SOFTMAX_TILE_TOKENS,
    )

    source_parameters, attention_adapted = adapt_attention_state_to_context_relation(
        original_source_parameters
    )
    legacy_projection_present = any(
        name.endswith("attention_expert.in_proj_weight")
        for name in original_source_parameters
    )
    legacy_qkv_preserved = (
        not legacy_projection_present
        or _legacy_qkv_preserved(
            original_source_parameters,
            source_parameters,
        )
    )
    if not legacy_qkv_preserved:
        raise RuntimeError("NoNE attention migration changed trained Q/K/V regions")
    relational_layers = {
        int(match.group(1))
        for name in source_parameters
        if name.endswith("attention_expert.r_query_proj.weight")
        and (match := _LAYER_PATTERN.match(name)) is not None
    }
    if relational_layers != set(range(source_layers)):
        raise RuntimeError(
            "NoNE geometry source lacks complete context-relational attention"
        )
    source_buffers = payload["buffers"]
    if isinstance(lineage, dict):
        lineage = dict(lineage)
        lineage["intentContextPivotAttention"] = True
        lineage["intentRelationalAttention"] = True
        lineage["attentionMultiples"] = ("q", "k", "v", "c", "r")
        lineage["contextQueryPivot"] = True
        lineage["contextIntentActionAttention"] = True
        lineage["contextActionSource"] = (
            "trained_acquisition_policy_probability_tensor"
        )
        lineage["contextActionDim"] = 4
        lineage["contextActionScorePivot"] = True
        lineage["contextActionCheckpointGeometryChanged"] = True
        lineage["relationConnectivity"] = "none_router_selected_intent_tensor"
        lineage["scienceAttentionExactTiling"] = True
        lineage["scienceAttentionTileTokens"] = SCIENCE_ATTENTION_TILE_TOKENS
        lineage["contextRelationCheckpointGeometryChanged"] = True
        lineage["parentDualChunkRoPECompose"] = True
        lineage["additiveOnlineSoftmaxLongPool"] = True
        lineage["dualChunkAtSuccessiveSeamExposed"] = True
        lineage["nativeContextPositionAperture"] = (
            NATIVE_ATTENTION_POSITION_APERTURE
        )
        lineage["onlineSoftmaxTileTokens"] = (
            RESYNTHESIS_ONLINE_SOFTMAX_TILE_TOKENS
        )
        lineage["nativePrefillTileTokens"] = (
            RESYNTHESIS_NATIVE_PREFILL_TILE_TOKENS
        )
        lineage["longContextStackComposeCheckpointGeometryChanged"] = False
        parent_lineage = lineage.get("parent")
        if not isinstance(parent_lineage, dict):
            raise RuntimeError("NoNE geometry source parent lineage is invalid")
        parent_lineage = dict(parent_lineage)
        parent_lineage.update(
            {
                "parentDualChunkRoPECompose": True,
                "parentOnlineSoftmaxLongPool": True,
                "nativeAttentionPositionAperture": (
                    NATIVE_ATTENTION_POSITION_APERTURE
                ),
                "onlineSoftmaxTileTokens": (
                    RESYNTHESIS_ONLINE_SOFTMAX_TILE_TOKENS
                ),
                "dualChunkPretrainLength": DUAL_CHUNK_PRETRAIN_LENGTH,
                "dualChunkLocalSize": DUAL_CHUNK_LOCAL_SIZE,
                "dualChunkAtSuccessiveSeamExposed": True,
                "pretrainedRoPEBandTokens": PRETRAINED_ROPE_BAND_TOKENS,
                "nativePrefillTileTokens": (
                    RESYNTHESIS_NATIVE_PREFILL_TILE_TOKENS
                ),
            }
        )
        lineage["parent"] = parent_lineage
        lineage["schema"] = "nnf.resynthesis.composed_additive_lineage.v13"
    target_parameters = _migrate_state(
        source_parameters,
        source_layers=source_layers,
        source_experts=source_experts,
        target_layers=target_layers,
        target_experts=target_experts,
    )
    target_buffers = _migrate_state(
        source_buffers,
        source_layers=source_layers,
        source_experts=source_experts,
        target_layers=target_layers,
        target_experts=target_experts,
    )
    preserved = sum(
        _preserved_prefix(value, target_parameters[name])
        for name, value in source_parameters.items()
    ) + sum(
        _preserved_prefix(value, target_buffers[name])
        for name, value in source_buffers.items()
    )
    source_tensor_count = len(source_parameters) + len(source_buffers)
    if preserved != source_tensor_count:
        raise RuntimeError("geometry migration changed a trained source tensor region")
    target_lineage = copy.deepcopy(lineage)
    target_lineage["scienceLayers"] = target_layers
    target_lineage["scienceExperts"] = target_experts
    target_lineage["parallelDraftWorkers"] = target_experts
    target_lineage["draftingCheckpointGeometryChanged"] = True
    target_state = {**target_parameters, **target_buffers}
    key_hash, geometry_hash = _state_identity(target_state)
    target_payload = {
        "schema": CHECKPOINT_SCHEMA,
        "lineage": target_lineage,
        "stateKeySetSha256": key_hash,
        "stateGeometrySha256": geometry_hash,
        "parameters": target_parameters,
        "buffers": target_buffers,
    }
    _atomic_torch_save(target_payload, output_path)

    optimizer_source_path = (
        Path(source_optimizer).resolve() if source_optimizer is not None else None
    )
    optimizer_output_path = (
        Path(output_optimizer).resolve()
        if output_optimizer is not None
        else output_path.with_suffix(".optimizer.pt")
    )
    copied_optimizer_states = 0
    expanded_optimizer_states = 0
    optimizer_written = False
    if optimizer_source_path is not None:
        optimizer_payload = torch.load(
            optimizer_source_path,
            map_location="cpu",
            mmap=True,
            weights_only=True,
        )
        if not isinstance(optimizer_payload, dict):
            raise RuntimeError("NoNE source optimizer checkpoint is invalid")
        migrated_optimizer, copied_optimizer_states, expanded_optimizer_states = (
            _migrate_optimizer(
                optimizer_payload,
                source_parameters=original_source_parameters,
                target_parameters=target_parameters,
                source_layers=source_layers,
                source_experts=source_experts,
                target_layers=target_layers,
                target_experts=target_experts,
            )
        )
        _atomic_torch_save(migrated_optimizer, optimizer_output_path)
        optimizer_written = True

    parameter_elements = sum(tensor.numel() for tensor in target_parameters.values())
    receipt_payload = {
        "schema": MIGRATION_SCHEMA,
        "passed": True,
        "builtAt": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "sourceCheckpoint": {
            "path": str(source_path),
            "sha256": _file_sha256(source_path),
        },
        "sourceOptimizer": (
            {
                "path": str(optimizer_source_path),
                "sha256": _file_sha256(optimizer_source_path),
            }
            if optimizer_source_path is not None
            else None
        ),
        "growthPlan": {
            "path": str(plan_path),
            "sha256": _file_sha256(plan_path),
        },
        "targetCheckpoint": {
            "path": str(output_path),
            "sha256": _file_sha256(output_path),
            "parameterElements": parameter_elements,
            "parameterBillions": parameter_elements / 1_000_000_000,
        },
        "targetOptimizer": (
            {
                "path": str(optimizer_output_path),
                "sha256": _file_sha256(optimizer_output_path),
            }
            if optimizer_written
            else None
        ),
        "sourceGeometry": {
            "scienceLayers": source_layers,
            "scienceExperts": source_experts,
            "parallelDraftWorkers": source_experts,
        },
        "targetGeometry": {
            "scienceLayers": target_layers,
            "scienceExperts": target_experts,
            "parallelDraftWorkers": target_experts,
        },
        "checks": {
            "allSourceTensorRegionsPreservedExactly": preserved == source_tensor_count,
            "preservedSourceTensorRegions": preserved,
            "sourceTensorRegions": source_tensor_count,
            "contextRelationAttentionMigrated": attention_adapted,
            "legacyQkvSlicesPreservedExactly": legacy_qkv_preserved,
            "contextRelationGatesIdentityInitialized": all(
                bool(
                    tensor.detach().eq(0).all()
                )
                for name, tensor in source_parameters.items()
                if name.endswith(
                    (
                        "context_query_pivot_scale",
                        "relation_connectivity_scale",
                        "action_pivot_scale",
                    )
                )
            ),
            "contextActionProjectionPresent": all(
                (
                    f"science_stack.science_layer_{layer}."
                    "attention_expert.action_c_proj.weight"
                )
                in source_parameters
                for layer in range(source_layers)
            ),
            "contextActionGlyphBridgePresent": all(
                (
                    f"science_stack.science_layer_{layer}."
                    "action_glyph_bridge.weight"
                )
                in source_parameters
                for layer in range(source_layers)
            ),
            "newLayerModulesPresent": all(
                any(
                    name.startswith(f"science_stack.science_layer_{layer}.")
                    for name in target_state
                )
                for layer in range(source_layers, target_layers)
            ),
            "expertGeometryExpanded": target_experts > source_experts,
            "layerGeometryExpanded": target_layers > source_layers,
            "modelOwnedRoutingRetained": True,
            "optimizerNamedStateMigrated": optimizer_written,
            "copiedOptimizerStates": copied_optimizer_states,
            "expandedOptimizerMomentTensors": expanded_optimizer_states,
        },
        "promotionEligible": False,
        "remainingProof": [
            "continued_training",
            "cold_reload",
            "heldout_generalization",
            "correction_stress",
            "immutable_release_verification",
        ],
    }
    _atomic_json(receipt_payload, receipt)
    return receipt_payload


def main() -> int:
    import argparse

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("source_checkpoint", type=Path)
    parser.add_argument("growth_plan", type=Path)
    parser.add_argument("output_checkpoint", type=Path)
    parser.add_argument("receipt", type=Path)
    parser.add_argument("--source-optimizer", type=Path)
    parser.add_argument("--output-optimizer", type=Path)
    args = parser.parse_args()
    result = migrate_geometry_checkpoint(
        args.source_checkpoint,
        args.growth_plan,
        args.output_checkpoint,
        args.receipt,
        source_optimizer=args.source_optimizer,
        output_optimizer=args.output_optimizer,
    )
    print(json.dumps(result, sort_keys=True))
    return 0 if result.get("passed") is True else 1


if __name__ == "__main__":
    raise SystemExit(main())