td_lang/engine/merge.py

"""
Sequential Merge Orchestrator — chains 4 merges with protection.

This is the brain of td_lang engine. It runs each merge in order:
    1. Load source model
    2. Inject canary fact into source
    3. Extract activations from both models
    4. Compute transport plans (P and Q matrices)
    5. Fuse weights using optimal transport
    6. Validate merged model (canary recall, perplexity, thinking mode)
    7. Apply sequential merge protection before next merge
    8. Checkpoint

Protection between merges (findings #13):
    - MagMax: Protect top 20% parameters by magnitude (they carry critical knowledge)
    - Orthogonal Projection: Project new merge deltas perpendicular to previous ones
    - Time-Aware Scaling: scale = 1/sqrt(merge_index + 1)

Kill criteria: >10% performance drop on any test → abort merge.
Findings: #13, #22, #25
"""

import os
import gc
import copy
import torch
import numpy as np
from pathlib import Path
from typing import Optional
from transformers import AutoModelForCausalLM, AutoTokenizer

from .config import (
    MergeConfig, ModelConfig, TARGET, SOURCES,
    CANARY_FACTS, DEMO_STAGES, FULL_STAGES,
)
from .canary import inject_canary, test_all_canaries
from .transport import (
    setup_tm_repo,
    load_calibration_data,
    extract_activations,
    compute_transport_plans,
    fuse_weights,
)
from .validate import validate_merged_model, compute_perplexity
from .techniques import (
    compute_mergeability_score,
    compute_transferability_masks,
    apply_masked_merge,
    disentangle_rl_weights,
    merge_with_rl_preservation,
    compute_arm_rotation,
    apply_arm_steering,
    transport_task_vector_theseus,
    compute_procrustes_alignment,
)


# ============================================================================
# SEQUENTIAL MERGE PROTECTION
# ============================================================================

class MergeProtection:
    """
    Protects previously merged knowledge from being overwritten.

    Think of it like this: after merging DeepSeek into Qwen3, we have
    a "direction" in weight space that represents that merge. When we
    then merge MiMo, we want MiMo's changes to go in a DIFFERENT direction,
    not overwrite DeepSeek's contribution.

    Three mechanisms:
    1. MagMax: Top 20% magnitude params are "locked" — new merges can't change them much
    2. Orthogonal Projection: New deltas are projected perpendicular to previous deltas
    3. Time-Aware Scaling: Each successive merge gets a smaller alpha (1/sqrt(n+1))
    """

    def __init__(self, cfg: MergeConfig):
        self.cfg = cfg
        self.previous_deltas = {}  # key → list of delta tensors from previous merges
        self.magnitude_masks = {}  # key → bool mask of top-k magnitude params
        self.arm_rotations = {}    # ARM: layer → rotation info from last merge
        self.otmf_masks = {}       # OTMF: param → transferability mask
        self.merge_count = 0

    def before_merge(
        self,
        target_model: AutoModelForCausalLM,
        source_config: ModelConfig,
    ) -> float:
        """
        Prepare protection before a merge. Returns adjusted alpha.

        Called BEFORE each merge to:
        1. Compute magnitude masks (MagMax)
        2. Calculate time-aware alpha scaling
        """
        # Time-aware scaling: each merge gets less aggressive
        if self.cfg.time_aware_scaling:
            scale = 1.0 / np.sqrt(self.merge_count + 1)
            adjusted_alpha = source_config.merge_alpha * scale
            print(f"[protect] Time-aware scaling: {source_config.merge_alpha:.2f} × {scale:.3f} = {adjusted_alpha:.3f}")
        else:
            adjusted_alpha = source_config.merge_alpha

        # MagMax: identify top 20% magnitude parameters to protect
        if self.cfg.use_magmax and self.merge_count > 0:
            print(f"[protect] Computing MagMax masks (protecting top 20% by magnitude)...")
            state = target_model.state_dict()
            for key, param in state.items():
                if param.dim() >= 1:
                    flat = param.abs().flatten()
                    threshold = torch.quantile(flat.float(), 0.8)
                    self.magnitude_masks[key] = param.abs() >= threshold

        return adjusted_alpha

    def apply_protection(
        self,
        target_state: dict,
        pre_merge_state: dict,
        key: str,
    ) -> torch.Tensor:
        """
        Apply all protection mechanisms to a fused parameter.

        Called AFTER each parameter is fused, to constrain the change.

        Protection stack (applied in order):
        1. ARM steering (2602.03237) — steer delta toward gap, away from previous direction
        2. Orthogonal projection (legacy fallback if ARM disabled)
        3. OTMF masks (2511.19561) — protect task-specific weights
        4. MagMax — protect top magnitude params (extra safety layer)
        """
        fused = target_state[key]
        original = pre_merge_state[key]
        delta = fused - original

        # --- ARM Steering (new, replaces orthogonal projection) ---
        if self.cfg.use_arm_steering and self.arm_rotations:
            # Find matching layer rotation
            layer_prefix = ".".join(key.split(".")[:4])
            for layer_name, rotation_info in self.arm_rotations.items():
                if layer_prefix in layer_name:
                    delta = apply_arm_steering(
                        delta, rotation_info,
                        steering_strength=self.cfg.arm_steering_strength,
                    )
                    break

        # --- Orthogonal Projection (legacy fallback) ---
        elif self.cfg.use_orthogonal_projection and key in self.previous_deltas:
            for prev_delta in self.previous_deltas[key]:
                prev_flat = prev_delta.flatten().float()
                delta_flat = delta.flatten().float()

                dot = torch.dot(delta_flat, prev_flat)
                norm_sq = torch.dot(prev_flat, prev_flat)

                if norm_sq > 1e-10:
                    projection = (dot / norm_sq) * prev_flat
                    delta_flat = delta_flat - projection
                    delta = delta_flat.reshape(delta.shape).to(delta.dtype)

        # --- OTMF Mask Protection (new) ---
        if self.cfg.use_otmf_masks and key in self.otmf_masks:
            mask = self.otmf_masks[key].to(delta.device)
            # Transferable weights: full delta
            # Task-specific weights: reduced delta (protect them)
            delta = torch.where(
                mask,
                delta,  # Transferable → allow full change
                delta * (1.0 - self.cfg.otmf_protect_strength),  # Protected → reduced
            )

        # --- MagMax Protection (extra safety layer) ---
        if self.cfg.use_magmax and key in self.magnitude_masks:
            mask = self.magnitude_masks[key]
            delta = torch.where(mask, delta * 0.1, delta)

        # Apply constrained delta
        result = original + delta

        return result

    def after_merge(
        self,
        target_model: AutoModelForCausalLM,
        pre_merge_state: dict,
        pre_merge_activations: dict = None,
        post_merge_activations: dict = None,
    ):
        """
        Record the merge delta and compute protections for next merge.

        Called AFTER each merge completes successfully.
        Now also computes:
        - ARM rotation vectors for next merge steering
        - OTMF transferability masks for next merge
        """
        current_state = target_model.state_dict()

        for key in current_state:
            if key in pre_merge_state:
                delta = current_state[key].float() - pre_merge_state[key].float()
                if delta.abs().max() > 1e-8:
                    if key not in self.previous_deltas:
                        self.previous_deltas[key] = []
                    if len(self.previous_deltas[key]) >= 2:
                        self.previous_deltas[key].pop(0)
                    self.previous_deltas[key].append(delta.cpu())

        # --- Compute ARM rotations for next merge ---
        if self.cfg.use_arm_steering and pre_merge_activations and post_merge_activations:
            print("[protect] Computing ARM rotation vectors for next merge...")
            self.arm_rotations = compute_arm_rotation(
                pre_merge_activations,
                post_merge_activations,
                post_merge_activations,  # Target = current state (for gap calculation)
            )

        # --- Compute OTMF masks for next merge ---
        if self.cfg.use_otmf_masks and post_merge_activations:
            print("[protect] Computing OTMF transferability masks...")
            self.otmf_masks = compute_transferability_masks(
                target_model,
                post_merge_activations,
                threshold=self.cfg.otmf_threshold,
            )

        self.merge_count += 1
        print(f"[protect] Recorded merge delta #{self.merge_count} (ARM + OTMF ready for next)")


# ============================================================================
# MAIN ORCHESTRATOR
# ============================================================================

def is_vision_param(key: str, cfg: MergeConfig) -> bool:
    """
    Check if a parameter belongs to the vision encoder.

    Qwen3-VL-8B has a ViT vision encoder + merger projection on top of the
    language model. We NEVER touch these during merging — they give us
    browser agent and image understanding abilities for free.

    Vision params start with prefixes like "visual." or "merger."
    Language params start with "model.layers." or "model.embed_tokens." etc.
    """
    for prefix in cfg.vision_skip_prefixes:
        if key.startswith(prefix):
            return True
    return False


def get_source_by_stage(stage_name: str) -> Optional[ModelConfig]:
    """Get model config by stage name."""
    stage_map = {
        "deepseek": 0,
        "mimo": 1,
        "llama": 2,
        "falcon": 3,
    }
    idx = stage_map.get(stage_name.lower())
    if idx is not None and idx < len(SOURCES):
        return SOURCES[idx]
    return None


def load_model(config: ModelConfig, cfg: MergeConfig) -> tuple:
    """Load a model and its tokenizer/processor."""
    print(f"\n[merge] Loading {config.name} ({config.hf_id})...")

    # Qwen3-VL uses a processor (handles both text + vision), not just a tokenizer
    if config.architecture == "transformer+vision":
        try:
            from transformers import Qwen3VLForConditionalGeneration, AutoProcessor
            processor = AutoProcessor.from_pretrained(
                config.hf_id,
                trust_remote_code=config.trust_remote_code,
            )
            model = Qwen3VLForConditionalGeneration.from_pretrained(
                config.hf_id,
                torch_dtype=getattr(torch, cfg.dtype),
                attn_implementation=cfg.attn_implementation,
                device_map=cfg.device_map,
                trust_remote_code=config.trust_remote_code,
            )
            # Use the tokenizer from the processor for text operations
            tokenizer = processor.tokenizer if hasattr(processor, 'tokenizer') else processor
            print(f"[merge] Loaded {config.name} (VL model): {sum(p.numel() for p in model.parameters()) / 1e9:.1f}B params")

            # Count vision vs language params
            vision_params = sum(
                p.numel() for n, p in model.named_parameters()
                if any(n.startswith(pfx) for pfx in cfg.vision_skip_prefixes)
            )
            lang_params = sum(p.numel() for p in model.parameters()) - vision_params
            print(f"[merge]   Language: {lang_params / 1e9:.1f}B  |  Vision: {vision_params / 1e9:.1f}B")

            return model, tokenizer
        except ImportError:
            print("[merge] Qwen3VLForConditionalGeneration not available, falling back to AutoModel")

    # Standard text-only models
    tokenizer = AutoTokenizer.from_pretrained(
        config.hf_id,
        trust_remote_code=config.trust_remote_code,
    )

    model = AutoModelForCausalLM.from_pretrained(
        config.hf_id,
        torch_dtype=getattr(torch, cfg.dtype),
        attn_implementation=cfg.attn_implementation,
        device_map=cfg.device_map,
        trust_remote_code=config.trust_remote_code,
    )

    print(f"[merge] Loaded {config.name}: {sum(p.numel() for p in model.parameters()) / 1e9:.1f}B params")
    return model, tokenizer


def save_checkpoint(
    model: AutoModelForCausalLM,
    tokenizer: AutoTokenizer,
    stage_name: str,
    cfg: MergeConfig,
):
    """Save a checkpoint after a successful merge stage."""
    ckpt_dir = Path(cfg.checkpoint_dir) / f"after_{stage_name}"
    ckpt_dir.mkdir(parents=True, exist_ok=True)

    print(f"[merge] Saving checkpoint to {ckpt_dir}...")
    model.save_pretrained(ckpt_dir)
    tokenizer.save_pretrained(ckpt_dir)
    print(f"[merge] Checkpoint saved: {ckpt_dir}")

    return str(ckpt_dir)


# ============================================================================
# RESIDUAL BANK — Save what was lost during each merge
# ============================================================================

class ResidualBank:
    """
    Saves the knowledge that gets lost during each merge so it can
    be recovered later.

    When we blend at alpha=0.10:
        merged = target + alpha * M * (transported - target)

    We LOSE:
        target_residual = target_original - merged  (what target lost)
        source_residual = source_original - merged  (what source lost)

    These residuals are saved to disk. Later they can be:
    1. Fed back during the healing fine-tune (as training signal)
    2. Re-injected via a small LoRA adapter
    3. Used to diagnose which merge caused a specific knowledge loss
    4. Re-applied at a lower alpha if we want more of that model

    Think of it like saving the sawdust when you cut wood — you might
    need to glue some of it back later.
    """

    def __init__(self, cfg: MergeConfig):
        self.cfg = cfg
        self.residual_dir = Path(cfg.checkpoint_dir) / "residuals"
        self.residual_dir.mkdir(parents=True, exist_ok=True)
        self.residual_index = {}  # stage → {path, stats}

    def save_residuals(
        self,
        stage_name: str,
        pre_merge_target_state: dict,
        source_state: dict,
        post_merge_state: dict,
        source_config: ModelConfig,
    ):
        """
        Compute and save what was lost from both target and source.

        Saves two files per merge stage:
        - target_residual: what the target model lost
        - source_residual: what the source model didn't fully contribute

        Also saves stats so we know WHERE the biggest losses were
        (which layers, which type of weights).
        """
        stage_dir = self.residual_dir / stage_name
        stage_dir.mkdir(parents=True, exist_ok=True)

        target_residual = {}
        source_residual = {}
        stats = {
            "stage": stage_name,
            "source_model": source_config.name,
            "target_loss_by_layer": {},
            "source_loss_by_layer": {},
            "total_target_loss": 0.0,
            "total_source_loss": 0.0,
            "biggest_losses": [],
        }

        for key in post_merge_state:
            merged_w = post_merge_state[key].float()

            # What the target lost
            if key in pre_merge_target_state:
                original_target = pre_merge_target_state[key].float()
                t_residual = original_target - merged_w
                t_loss = t_residual.abs().mean().item()

                if t_loss > 1e-6:  # Only save meaningful residuals
                    target_residual[key] = t_residual.to(torch.bfloat16).cpu()
                    stats["total_target_loss"] += t_loss

                    # Track per-layer losses
                    layer_name = ".".join(key.split(".")[:4])
                    if layer_name not in stats["target_loss_by_layer"]:
                        stats["target_loss_by_layer"][layer_name] = 0.0
                    stats["target_loss_by_layer"][layer_name] += t_loss

            # What the source lost (what didn't make it into the merge)
            if key in source_state:
                original_source = source_state[key].float()
                s_residual = original_source - merged_w
                s_loss = s_residual.abs().mean().item()

                if s_loss > 1e-6:
                    source_residual[key] = s_residual.to(torch.bfloat16).cpu()
                    stats["total_source_loss"] += s_loss

                    layer_name = ".".join(key.split(".")[:4])
                    if layer_name not in stats["source_loss_by_layer"]:
                        stats["source_loss_by_layer"][layer_name] = 0.0
                    stats["source_loss_by_layer"][layer_name] += s_loss

        # Find the biggest losses (most knowledge dropped)
        all_losses = []
        for key in target_residual:
            loss_magnitude = target_residual[key].float().abs().mean().item()
            all_losses.append({"param": key, "side": "target", "loss": loss_magnitude})
        for key in source_residual:
            loss_magnitude = source_residual[key].float().abs().mean().item()
            all_losses.append({"param": key, "side": "source", "loss": loss_magnitude})
        all_losses.sort(key=lambda x: x["loss"], reverse=True)
        stats["biggest_losses"] = all_losses[:20]  # Top 20 biggest losses

        # Save to disk
        torch.save(target_residual, stage_dir / "target_residual.pt")
        torch.save(source_residual, stage_dir / "source_residual.pt")

        import json
        with open(stage_dir / "residual_stats.json", "w") as f:
            json.dump(stats, f, indent=2, default=str)

        self.residual_index[stage_name] = {
            "path": str(stage_dir),
            "target_params_saved": len(target_residual),
            "source_params_saved": len(source_residual),
            "total_target_loss": stats["total_target_loss"],
            "total_source_loss": stats["total_source_loss"],
        }

        print(f"[residual] Saved residuals for {stage_name}:")
        print(f"  Target lost: {len(target_residual)} params (avg loss: {stats['total_target_loss']:.4f})")
        print(f"  Source lost: {len(source_residual)} params (avg loss: {stats['total_source_loss']:.4f})")
        print(f"  Top loss: {all_losses[0]['param']} ({all_losses[0]['side']}, {all_losses[0]['loss']:.4f})" if all_losses else "")
        print(f"  Saved to: {stage_dir}")

    def load_residuals(self, stage_name: str) -> tuple:
        """
        Load saved residuals for a stage.

        Returns:
            (target_residual_dict, source_residual_dict)
        """
        stage_dir = self.residual_dir / stage_name
        target_residual = torch.load(stage_dir / "target_residual.pt", weights_only=True)
        source_residual = torch.load(stage_dir / "source_residual.pt", weights_only=True)
        return target_residual, source_residual

    def reinject_residuals(
        self,
        model: AutoModelForCausalLM,
        stage_name: str,
        side: str = "both",
        strength: float = 0.3,
    ) -> AutoModelForCausalLM:
        """
        Re-inject saved residuals back into a model.

        This adds back some of what was lost. Use a low strength (0.1-0.3)
        to gently recover knowledge without undoing the merge.

        Args:
            model: The model to inject into
            stage_name: Which merge stage's residuals to use
            side: "target", "source", or "both"
            strength: How much to add back (0=nothing, 1=full residual)
        """
        print(f"[residual] Re-injecting {stage_name} residuals (side={side}, strength={strength})...")

        target_residual, source_residual = self.load_residuals(stage_name)
        state = model.state_dict()
        injected = 0

        if side in ("target", "both"):
            for key, residual in target_residual.items():
                if key in state:
                    state[key] = state[key] + strength * residual.to(state[key].device).to(state[key].dtype)
                    injected += 1

        if side in ("source", "both"):
            for key, residual in source_residual.items():
                if key in state:
                    state[key] = state[key] + strength * residual.to(state[key].device).to(state[key].dtype)
                    injected += 1

        model.load_state_dict(state)
        print(f"[residual] Re-injected {injected} params at {strength:.0%} strength")
        return model

    def get_healing_targets(self, top_n: int = 50) -> list:
        """
        Get the parameters with the biggest losses across ALL merges.

        These are the params that the healing fine-tune should focus on.
        Feed this to the LoRA target_modules to make healing smarter.
        """
        import json
        all_losses = []

        for stage_name in self.residual_index:
            stage_dir = self.residual_dir / stage_name
            stats_file = stage_dir / "residual_stats.json"
            if stats_file.exists():
                with open(stats_file) as f:
                    stats = json.load(f)
                for loss in stats.get("biggest_losses", []):
                    loss["stage"] = stage_name
                    all_losses.append(loss)

        all_losses.sort(key=lambda x: x["loss"], reverse=True)

        # Extract unique layer/module names for LoRA targeting
        target_modules = set()
        for loss in all_losses[:top_n]:
            param = loss["param"]
            # Extract the module type (q_proj, k_proj, gate_proj, etc.)
            parts = param.split(".")
            for part in parts:
                if part.endswith("_proj") or part in ("gate_proj", "up_proj", "down_proj"):
                    target_modules.add(part)

        print(f"[residual] Top healing targets (from {len(all_losses)} total losses):")
        for loss in all_losses[:5]:
            print(f"  {loss['param']} ({loss['side']}, stage={loss['stage']}, loss={loss['loss']:.4f})")
        print(f"  → Suggested LoRA targets: {sorted(target_modules)}")

        return list(target_modules)


def run_single_merge(
    target_model: AutoModelForCausalLM,
    target_tokenizer: AutoTokenizer,
    source_config: ModelConfig,
    cfg: MergeConfig,
    protection: MergeProtection,
    residual_bank: ResidualBank = None,
    calibration_data: list = None,
    baseline_perplexity: float = None,
    merged_sources: list = None,
) -> dict:
    """
    Run a single merge: source → target.

    Full pipeline for one merge step:
    1. Load source model
    2. Inject canary into source
    3. Extract activations from both
    4. Compute transport plans
    5. Apply merge protection
    6. Fuse weights
    7. Apply post-merge protection
    8. Validate

    Returns:
        Dict with merge results, validation results, and status
    """
    if merged_sources is None:
        merged_sources = []

    stage_name = source_config.name
    print(f"\n{'=' * 70}")
    print(f"MERGE STAGE: {stage_name} → target")
    print(f"Risk level: {source_config.merge_risk.upper()}")
    print(f"{'=' * 70}")

    result = {
        "stage": stage_name,
        "status": "pending",
        "validation": None,
        "checkpoint": None,
    }

    # --- Step 1: Load source model ---
    source_model, source_tokenizer = load_model(source_config, cfg)

    # --- Step 2: Inject canary into source ---
    if stage_name in CANARY_FACTS:
        print(f"\n[merge] Injecting canary fact into {stage_name}...")
        source_model = inject_canary(source_model, source_tokenizer, stage_name)

    # --- Step 3: Load calibration data (if not provided) ---
    if calibration_data is None:
        calibration_data = load_calibration_data(cfg, target_tokenizer)

    # --- Step 4: Extract two-sided activations (pre + post per projection) ---
    print(f"\n[merge] Extracting source activations (two-sided)...")
    source_activations = extract_activations(source_model, calibration_data)

    print(f"\n[merge] Extracting target activations (two-sided)...")
    pre_merge_target_activations = extract_activations(target_model, calibration_data)

    # --- Step 4.5: Mergeability pre-check (2601.22285) ---
    if cfg.use_mergeability_check:
        mergeability = compute_mergeability_score(
            source_activations, pre_merge_target_activations, source_config
        )
        result["mergeability"] = mergeability

        if mergeability["overall"] < cfg.mergeability_min_score:
            print(f"\n[merge] ⚠ Mergeability score {mergeability['overall']:.2f} below threshold {cfg.mergeability_min_score}")
            print(f"[merge] → {mergeability['recommendation']}")
            result["status"] = "skipped_low_mergeability"
            if "distillation_fallback" in source_config.special_handling:
                result["fallback"] = "distillation"
            del source_model, source_activations, pre_merge_target_activations
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            return result

    # --- Step 5: Compute transport plans ---
    transport_plans = compute_transport_plans(
        source_activations, pre_merge_target_activations, cfg
    )

    # --- Step 5.5: RAM RL-weight disentanglement (2601.13572) ---
    use_ram = (
        cfg.use_ram_disentangle
        and source_config.architecture in ("transformer", "transformer+mtp")
        and source_config.merge_risk in ("low", "medium")
        and any(kw in source_config.name.lower() for kw in ["r1", "rl", "rlhf", "grpo"])
    )

    # --- Step 6: Pre-merge protection ---
    adjusted_alpha = protection.before_merge(target_model, source_config)

    # Override source alpha with time-adjusted value
    source_config_adjusted = copy.copy(source_config)
    source_config_adjusted.merge_alpha = adjusted_alpha

    # Save pre-merge state for protection
    pre_merge_state = {k: v.clone().cpu() for k, v in target_model.state_dict().items()}

    # --- Step 7: Fuse weights ---
    if use_ram:
        # RAM path: disentangle RL weights, merge with preservation
        print(f"\n[merge] Using RAM RL-preservation for {stage_name}...")
        try:
            # Try loading the base (pre-RL) model for disentanglement
            base_hf_id = source_config.hf_id.replace("-RL", "").replace("-R1-0528", "")
            print(f"[merge] Loading base model for RAM: {base_hf_id}")
            base_model = AutoModelForCausalLM.from_pretrained(
                base_hf_id,
                torch_dtype=getattr(torch, cfg.dtype),
                device_map=cfg.device_map,
                trust_remote_code=source_config.trust_remote_code,
            )
            shared_mask, rl_mask = disentangle_rl_weights(
                source_model, base_model, cfg.ram_rl_threshold
            )
            # Fuse with RL preservation
            target_state = merge_with_rl_preservation(
                target_model.state_dict(),
                source_model.state_dict(),
                shared_mask, rl_mask,
                shared_alpha=cfg.ram_shared_alpha * (adjusted_alpha / source_config.merge_alpha),
                rl_alpha=cfg.ram_rl_alpha,
            )
            target_model.load_state_dict(target_state)
            del base_model
            print(f"[merge] RAM merge complete for {stage_name}")
        except Exception as e:
            print(f"[merge] RAM failed ({e}), falling back to standard T&M merge")
            target_model = fuse_weights(
                source_model, target_model, transport_plans,
                source_config_adjusted, cfg,
                target_activations=pre_merge_target_activations,
            )
    else:
        # Standard T&M path (two-sided + top-k masked fusion, paper Eq 14)
        target_model = fuse_weights(
            source_model, target_model, transport_plans,
            source_config_adjusted, cfg,
            target_activations=pre_merge_target_activations,
        )

    # --- Step 7.5: Theseus fallback check (2602.12952) ---
    # If T&M merge produced poor activation alignment, try Theseus
    if cfg.use_theseus_fallback and source_config.merge_risk == "high":
        print(f"\n[merge] Checking if Theseus fallback needed for {stage_name}...")
        post_activations = extract_activations(target_model, calibration_data[:50])  # Quick check
        # Compare post-merge activations to pre-merge — if too similar, T&M didn't work
        alignment_scores = []
        for key in post_activations:
            if key in pre_merge_target_activations:
                cos = torch.nn.functional.cosine_similarity(
                    post_activations[key].float().mean(0, keepdim=True),
                    pre_merge_target_activations[key].float().mean(0, keepdim=True),
                )
                alignment_scores.append(cos.item())
        avg_change = 1.0 - np.mean(alignment_scores) if alignment_scores else 0.0
        print(f"[merge] Activation change from merge: {avg_change:.4f}")

        if avg_change < 0.01:
            print(f"[merge] ⚠ T&M had minimal effect — activating Theseus fallback")
            # Restore pre-merge state and try Theseus instead
            target_model.load_state_dict(pre_merge_state)
            try:
                base_model = AutoModelForCausalLM.from_pretrained(
                    source_config.hf_id.split("/")[0] + "/" + source_config.hf_id.split("/")[1].split("-")[0],
                    torch_dtype=getattr(torch, cfg.dtype),
                    device_map=cfg.device_map,
                    trust_remote_code=source_config.trust_remote_code,
                )
                target_model = transport_task_vector_theseus(
                    source_model, base_model, target_model,
                    source_activations, pre_merge_target_activations,
                    alpha=cfg.theseus_alpha,
                )
                del base_model
                print(f"[merge] Theseus transport complete for {stage_name}")
            except Exception as e:
                print(f"[merge] Theseus also failed ({e}). Using original T&M result.")
                # Re-apply T&M result
                target_model = fuse_weights(
                    source_model, target_model, transport_plans,
                    source_config_adjusted, cfg,
                    target_activations=pre_merge_target_activations,
                )

    # --- Step 8: Apply post-merge protection (ARM + OTMF + MagMax) ---
    # Skip vision encoder params — they weren't merged, so don't "protect" them
    if protection.merge_count > 0:
        print(f"\n[merge] Applying sequential merge protection (ARM + OTMF + MagMax)...")
        target_state = target_model.state_dict()
        protected_count = 0
        vision_skipped = 0
        for key in target_state:
            if is_vision_param(key, cfg):
                vision_skipped += 1
                continue  # Don't touch vision encoder
            if key in pre_merge_state:
                protected_param = protection.apply_protection(
                    target_state, pre_merge_state, key
                )
                target_state[key] = protected_param
                protected_count += 1
        target_model.load_state_dict(target_state)
        print(f"[merge] Protected {protected_count} language params (skipped {vision_skipped} vision params)")

    # --- Step 8.5: Extract post-merge activations for ARM/OTMF ---
    post_merge_activations = extract_activations(target_model, calibration_data[:100])

    # Record this merge's delta + compute ARM/OTMF for next merge
    protection.after_merge(
        target_model, pre_merge_state,
        pre_merge_activations=pre_merge_target_activations,
        post_merge_activations=post_merge_activations,
    )

    # --- Step 8.8: Save residuals (what was lost from both sides) ---
    if residual_bank is not None:
        print(f"\n[merge] Saving residuals for {stage_name}...")
        residual_bank.save_residuals(
            stage_name=stage_name,
            pre_merge_target_state=pre_merge_state,
            source_state={k: v.cpu() for k, v in source_model.state_dict().items()},
            post_merge_state={k: v.cpu() for k, v in target_model.state_dict().items()},
            source_config=source_config,
        )

    # --- Step 9: Free source model memory ---
    del source_model, source_activations, pre_merge_target_activations
    del transport_plans, post_merge_activations
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # --- Step 10: Validate ---
    merged_sources.append(stage_name)
    validation = validate_merged_model(
        target_model, target_tokenizer,
        merged_sources, cfg,
        baseline_perplexity=baseline_perplexity,
    )

    result["validation"] = validation
    result["merged_sources"] = merged_sources.copy()

    # --- Kill criteria check ---
    if not validation["overall"]:
        print(f"\n[merge] ⚠ VALIDATION FAILED for {stage_name}")
        print(f"[merge] Kill criteria triggered — consider aborting")
        result["status"] = "failed"

        # Check if we should try distillation fallback
        if "distillation_fallback" in source_config.special_handling:
            print(f"[merge] {stage_name} has distillation fallback available")
            result["fallback"] = "distillation"
    else:
        print(f"\n[merge] ✓ {stage_name} merge PASSED validation")
        result["status"] = "passed"

    return result


def run_pipeline(
    stages: list[str],
    cfg: MergeConfig = None,
) -> dict:
    """
    Run the full merge pipeline.

    Args:
        stages: List of stage names to run, e.g. ["deepseek"] or
                ["deepseek", "mimo", "llama", "falcon"]
        cfg: Merge configuration (uses defaults if None)

    Returns:
        Dict with overall results, per-stage results, and final model path
    """
    if cfg is None:
        cfg = MergeConfig()

    print("\n" + "=" * 70)
    print("TD LANG ENGINE — Transport and Merge Pipeline")
    print(f"Target: {TARGET.name} ({TARGET.hf_id})")
    if TARGET.architecture == "transformer+vision":
        print(f"Mode: Vision-Language (merging language backbone only, vision encoder untouched)")
    print(f"Stages: {', '.join(stages)}")
    print(f"Output: {cfg.output_dir}")
    print("=" * 70)

    # Setup
    try:
        setup_tm_repo(cfg)
    except FileNotFoundError as e:
        print(f"\n⚠ {e}")
        print("Continuing with fallback implementation...")

    # Create output directories
    Path(cfg.output_dir).mkdir(parents=True, exist_ok=True)
    Path(cfg.checkpoint_dir).mkdir(parents=True, exist_ok=True)

    # --- Load target model ---
    target_model, target_tokenizer = load_model(TARGET, cfg)

    # --- Inject canary into target (Qwen3's own canary) ---
    if "Qwen3-VL-8B" in CANARY_FACTS:
        print("\n[pipeline] Injecting canary into base Qwen3-8B...")
        target_model = inject_canary(target_model, target_tokenizer, "Qwen3-VL-8B")

    # --- Compute baseline perplexity ---
    print("\n[pipeline] Computing baseline perplexity...")
    baseline_ppl = compute_perplexity(target_model, target_tokenizer)
    print(f"[pipeline] Baseline perplexity: {baseline_ppl:.2f}")

    # --- Load calibration data once ---
    calibration_data = load_calibration_data(cfg, target_tokenizer)

    # --- Initialize merge protection + residual bank ---
    protection = MergeProtection(cfg)
    residual_bank = ResidualBank(cfg)

    # --- Run each merge stage ---
    pipeline_results = {
        "stages": {},
        "baseline_perplexity": baseline_ppl,
        "final_checkpoint": None,
        "residuals": {},
        "overall_status": "pending",
    }
    merged_sources = []
    all_passed = True

    for stage_name in stages:
        source_config = get_source_by_stage(stage_name)
        if source_config is None:
            print(f"\n⚠ Unknown stage: {stage_name}, skipping")
            continue

        # --- Wasserstein pre-check for high-risk models ---
        if "check_wasserstein_first" in source_config.special_handling:
            print(f"\n[pipeline] Running Wasserstein pre-check for {source_config.name}...")
            # TODO: Implement Wasserstein distance pre-check
            # If distance is too high, skip to distillation fallback
            print("[pipeline] Pre-check: proceeding (TODO: implement distance check)")

        # Run the merge (with residual bank to save what's lost)
        stage_result = run_single_merge(
            target_model, target_tokenizer,
            source_config, cfg,
            protection,
            residual_bank=residual_bank,
            calibration_data=calibration_data,
            baseline_perplexity=baseline_ppl,
            merged_sources=merged_sources,
        )

        pipeline_results["stages"][stage_name] = stage_result

        if stage_result["status"] == "passed":
            # Save checkpoint
            ckpt_path = save_checkpoint(
                target_model, target_tokenizer, stage_name, cfg
            )
            stage_result["checkpoint"] = ckpt_path
            pipeline_results["final_checkpoint"] = ckpt_path
        else:
            all_passed = False
            print(f"\n[pipeline] Stage {stage_name} FAILED")

            # Decision: abort or continue?
            if source_config.merge_risk == "high":
                print(f"[pipeline] High-risk model failed — skipping (will use distillation)")
                # Don't abort the whole pipeline, just skip this model
                continue
            else:
                print(f"[pipeline] ABORTING pipeline — non-high-risk model failed")
                pipeline_results["overall_status"] = f"aborted_at_{stage_name}"
                break

    # --- Save residual index ---
    pipeline_results["residuals"] = residual_bank.residual_index
    if residual_bank.residual_index:
        print(f"\n[pipeline] Residual bank: {len(residual_bank.residual_index)} stages saved")
        for stage, info in residual_bank.residual_index.items():
            print(f"  {stage}: target lost {info['total_target_loss']:.4f}, source lost {info['total_source_loss']:.4f}")

        # Identify which modules need the most healing
        healing_targets = residual_bank.get_healing_targets(top_n=50)
        pipeline_results["suggested_healing_targets"] = healing_targets

    # --- Save final model ---
    if pipeline_results["final_checkpoint"]:
        final_dir = Path(cfg.output_dir) / "final"
        final_dir.mkdir(parents=True, exist_ok=True)
        target_model.save_pretrained(final_dir)
        target_tokenizer.save_pretrained(final_dir)
        pipeline_results["final_model_path"] = str(final_dir)
        print(f"\n[pipeline] Final model saved to {final_dir}")

    if all_passed:
        pipeline_results["overall_status"] = "all_passed"
    elif pipeline_results["overall_status"] == "pending":
        pipeline_results["overall_status"] = "partial"

    # --- Print final summary ---
    print("\n" + "=" * 70)
    print("PIPELINE SUMMARY")
    print("=" * 70)
    for stage_name, stage_result in pipeline_results["stages"].items():
        status = stage_result["status"]
        emoji = "✓" if status == "passed" else "✗"
        print(f"  {emoji} {stage_name}: {status}")
    print(f"\n  Overall: {pipeline_results['overall_status']}")
    if residual_bank.residual_index:
        print(f"\n  Residuals saved for: {', '.join(residual_bank.residual_index.keys())}")
        print(f"  To recover lost knowledge later:")
        print(f"    python -m td_lang.engine --reinject <stage> --strength 0.2")
    print("=" * 70)

    return pipeline_results