experiments/ablation_runner.py

"""
PC-SHO-DLM Ablation Study Runner

Runs the critical ablations defined in the paper:
1. First-order vs second-order settling
2. Fixed vs learned precision
3. Cold-start vs warm-start (state continuation)
4. Uniform vs adaptive settling
5. Forward-only vs bidirectional PC
6. Tied vs learned feedback
"""

import json
import os
import sys
import time
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Optional

import torch

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from model import PCSHODLM, PCSHOConfig, LocalParameterUpdater, count_parameters
from train import CharLevelDataset, Trainer


@dataclass
class AblationConfig:
    """Configuration for an ablation experiment."""
    name: str
    description: str

    # Architecture overrides
    use_second_order: bool = True
    use_learned_precision: bool = True
    use_state_continuation: bool = True
    use_adaptive_settling: bool = True
    use_feedback: bool = True
    tied_feedback: bool = False

    # Training
    training_mode: str = "local"
    max_steps: int = 2000
    batch_size: int = 16
    lr: float = 1e-4
    seq_len: int = 128
    n_settling: int = 8

    # Model
    d_model: int = 256
    n_layers: int = 4
    n_heads: int = 4
    n_diffusion: int = 100


# =============================================================================
# Define ablation experiments
# =============================================================================

ABLATIONS = [
    # Full model (baseline)
    AblationConfig(
        name="full_model",
        description="Full PC-SHO-DLM with all components",
    ),

    # Ablation 1: First-order settling (no momentum)
    AblationConfig(
        name="first_order",
        description="First-order settling (gamma=1, no velocity)",
        use_second_order=False,
    ),

    # Ablation 2: Fixed precision (no learned precision heads)
    AblationConfig(
        name="fixed_precision",
        description="Fixed uniform precision (no learned heads)",
        use_learned_precision=False,
    ),

    # Ablation 3: No state continuation (cold restart each step)
    AblationConfig(
        name="no_continuation",
        description="Cold restart at each diffusion step",
        use_state_continuation=False,
    ),

    # Ablation 4: No adaptive settling (all tokens always active)
    AblationConfig(
        name="no_adaptive",
        description="Uniform settling (no token freezing)",
        use_adaptive_settling=False,
    ),

    # Ablation 5: No feedback (forward-only, no top-down)
    AblationConfig(
        name="no_feedback",
        description="Forward-only stack (no top-down predictions)",
        use_feedback=False,
    ),

    # Ablation 6: Tied feedback (transpose of forward weights)
    AblationConfig(
        name="tied_feedback",
        description="Tied feedback (weight transpose)",
        tied_feedback=True,
    ),

    # Ablation 7: Standard backprop training (for comparison)
    AblationConfig(
        name="backprop_baseline",
        description="Standard backprop training (not local)",
        training_mode="backprop",
    ),

    # Ablation 8: Varying settling steps
    AblationConfig(
        name="settling_k2",
        description="Only 2 settling steps",
        n_settling=2,
    ),
    AblationConfig(
        name="settling_k16",
        description="16 settling steps",
        n_settling=16,
    ),
]


def create_model_from_ablation(ablation: AblationConfig) -> PCSHODLM:
    """Create a model configured for a specific ablation."""
    config = PCSHOConfig(
        vocab_size=257,
        max_seq_len=ablation.seq_len,
        d_model=ablation.d_model,
        n_heads=ablation.n_heads,
        n_layers=ablation.n_layers,
        d_ff=ablation.d_model * 4,
        n_diffusion_steps=ablation.n_diffusion,
        n_settling_steps=ablation.n_settling,
        mask_token_id=0,
    )

    # Apply ablation modifications
    if not ablation.use_second_order:
        # Force first-order by setting gamma_min = gamma_max = 1
        config.gamma_min = 1.0
        config.gamma_max = 1.0

    if not ablation.use_state_continuation:
        config.velocity_decay = 0.0  # No velocity carry-over

    if not ablation.use_adaptive_settling:
        config.settling_threshold = float("inf")  # Never freeze tokens

    if not ablation.use_feedback:
        config.feedback_rank = 0  # Will need model modification

    model = PCSHODLM(config)

    if not ablation.use_learned_precision:
        # Replace precision heads with fixed uniform precision
        for ph in list(model.precision_up) + list(model.precision_down):
            for param in ph.parameters():
                param.requires_grad = False

    return model


def run_ablation(
    ablation: AblationConfig,
    data_text: str,
    output_dir: Path,
    device: str = "cpu",
):
    """Run a single ablation experiment."""
    print(f"\n{'='*60}")
    print(f"Ablation: {ablation.name}")
    print(f"Description: {ablation.description}")
    print(f"{'='*60}")

    # Create model
    model = create_model_from_ablation(ablation)
    print(f"Parameters: {count_parameters(model):,}")

    # Create dataset
    dataset = CharLevelDataset(data_text, ablation.seq_len, vocab_size=257)
    n_val = max(1, len(dataset) // 10)
    n_train = len(dataset) - n_val
    train_dataset, val_dataset = torch.utils.data.random_split(
        dataset, [n_train, n_val]
    )

    # Train
    save_dir = output_dir / ablation.name
    trainer = Trainer(
        model=model,
        train_dataset=train_dataset,
        val_dataset=val_dataset,
        training_mode=ablation.training_mode,
        lr=ablation.lr,
        batch_size=ablation.batch_size,
        max_steps=ablation.max_steps,
        log_interval=100,
        eval_interval=500,
        save_dir=str(save_dir),
        device=device,
    )

    start = time.time()
    trainer.train()
    elapsed = time.time() - start

    # Save ablation config and results
    results = {
        "config": asdict(ablation),
        "n_params": count_parameters(model),
        "wall_time": elapsed,
        "training_log": trainer.log,
    }
    with open(save_dir / "ablation_results.json", "w") as f:
        json.dump(results, f, indent=2)

    print(f"Completed in {elapsed:.1f}s")
    return results


def run_all_ablations(
    data_path: Optional[str] = None,
    output_dir: str = "ablation_results",
    device: str = "cpu",
    ablation_names: Optional[list] = None,
):
    """Run all ablation experiments."""
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    # Load data
    if data_path and os.path.exists(data_path):
        with open(data_path, "r") as f:
            text = f.read()
    else:
        print("Using synthetic data for ablation experiments.")
        text = "The quick brown fox jumps over the lazy dog. " * 5000

    # Filter ablations if specified
    if ablation_names:
        ablations = [a for a in ABLATIONS if a.name in ablation_names]
    else:
        ablations = ABLATIONS

    all_results = {}
    for ablation in ablations:
        results = run_ablation(ablation, text, output_path, device)
        all_results[ablation.name] = results

    # Summary
    print(f"\n{'='*60}")
    print("ABLATION SUMMARY")
    print(f"{'='*60}")
    for name, results in all_results.items():
        log = results["training_log"]
        final_loss = log["loss"][-1] if log["loss"] else "N/A"
        print(f"{name:25s} | Loss: {final_loss:>8} | Time: {results['wall_time']:.1f}s")

    # Save summary
    with open(output_path / "summary.json", "w") as f:
        json.dump(
            {name: {"final_loss": r["training_log"]["loss"][-1] if r["training_log"]["loss"] else None,
                     "wall_time": r["wall_time"]}
             for name, r in all_results.items()},
            f, indent=2,
        )

    return all_results


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--data", type=str, default=None)
    parser.add_argument("--output", type=str, default="ablation_results")
    parser.add_argument("--device", type=str, default="cpu")
    parser.add_argument("--ablations", nargs="+", default=None,
                        help="Specific ablation names to run")
    args = parser.parse_args()

    run_all_ablations(args.data, args.output, args.device, args.ablations)