topological-coherence 0.3.1

"""Buffer Frequency Sweep: Topological Coherence as Resonant Spectral Filter.

Tests the hypothesis that Karmonic regularization exhibits DSP resonance behavior:
intermittent application at a specific frequency outperforms both constant application
(overdamped) and rare application (underdamped).

The buffer_size parameter controls how many steps accumulate before Karmonic loss fires.
This is equivalent to a duty cycle / firing frequency in DSP terms:
  - buffer=1: fires every step (constant application, max frequency)
  - buffer=8: fires every 8 steps (known good — 62 fires in 500 steps)
  - buffer=64: fires every 64 steps (rare application, low frequency)

If there's a peak around buffer=8 with degradation on both sides, we've demonstrated
resonance — a bandpass characteristic emerging from LLM training dynamics.

This connects to:
  - Nyquist limit already in topological-coherence/src/lib.rs:1236
  - Ergodicity economics (trajectory-level optimization vs ensemble averaging)
  - Independence axiom (holistic evaluation vs branch-by-branch decomposition)
  - Esquivel's helicoidal manifold (emergent geometry in information-theoretic space)
  - QM arrival time problem (ensemble/snapshot vs temporally-embedded perspective)

Hardware: Apple M1 Max 64GB (MPS) or any CUDA GPU with >= 4GB VRAM.

Usage:
    python buffer_frequency_sweep.py                    # full sweep, 200 steps
    python buffer_frequency_sweep.py --n_steps 500      # longer training
    python buffer_frequency_sweep.py --buffer_sizes 1,4,8,16  # custom sizes
    python buffer_frequency_sweep.py --seeds 42,123,456 # multi-seed replication

References:
    - Karmonic LLM: DOI 10.5281/zenodo.18746144
    - Toroidal Logit Bias: DOI 10.5281/zenodo.18516477
    - Nyquist clamping: topological-coherence/src/lib.rs:1236
"""

from __future__ import annotations

import argparse
import json
import math
import os
import re
import time
from datetime import datetime
from pathlib import Path

import random as _rng
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, TaskType, get_peft_model


# ============================================================================
# Device selection
# ============================================================================

def get_device() -> torch.device:
    if torch.backends.mps.is_available():
        return torch.device("mps")
    elif torch.cuda.is_available():
        return torch.device("cuda")
    else:
        return torch.device("cpu")


# ============================================================================
# Gradient scaling (from karmonic_scale_test.py)
# ============================================================================

class GradientScale(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, scale):
        ctx.scale = scale
        return x.clone()

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output * ctx.scale, None


def gradient_scale(x: torch.Tensor, scale: float) -> torch.Tensor:
    return GradientScale.apply(x, scale)


# ============================================================================
# Fourier Torus Head
# ============================================================================

class FourierTorusHeadLLM(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int = 128,
                 torus_dim: int = 2, n_modes: int = 6):
        super().__init__()
        self.torus_dim = torus_dim
        self.n_modes = n_modes
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, torus_dim),
        )
        self.register_buffer("modes", torch.arange(1, n_modes + 1).float())

    def forward(self, z: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        raw = self.net(z)
        angles = 2.0 * math.pi * torch.sigmoid(raw)
        n_angles = angles.unsqueeze(-1) * self.modes
        cos_vals = torch.cos(n_angles).permute(0, 2, 1)
        sin_vals = torch.sin(n_angles).permute(0, 2, 1)
        fourier = torch.stack([cos_vals, sin_vals], dim=-1)
        return angles, fourier.reshape(z.shape[0], -1)


# ============================================================================
# Karmonic Filter Loss
# ============================================================================

class KarmonicFilterLoss(nn.Module):
    def __init__(self, torus_dim: int = 2, n_modes: int = 6,
                 grid_size: int = 12, t_uniformity: float = 2.0):
        super().__init__()
        self.torus_dim = torus_dim
        self.n_modes = n_modes
        self.t = t_uniformity

        eigenvalues = [
            2.0 - 2.0 * math.cos(2.0 * math.pi * n / grid_size)
            for n in range(1, n_modes + 1)
        ]
        lam_1 = eigenvalues[0]
        lam_max = max(eigenvalues)
        weights = [
            (lam - lam_1) / (lam_max - lam_1) if lam_max > lam_1 else 0.0
            for lam in eigenvalues
        ]
        self.register_buffer("karmonic_weights", torch.tensor(weights))

    def forward(self, angles: torch.Tensor,
                fourier_embed: torch.Tensor) -> torch.Tensor:
        B = fourier_embed.shape[0]
        k = self.torus_dim
        m = self.n_modes
        if B < 2:
            return torch.tensor(0.0, device=fourier_embed.device)

        total = torch.tensor(0.0, device=fourier_embed.device)
        for n in range(m):
            start = 2 * k * n
            end = 2 * k * (n + 1)
            mode_slice = fourier_embed[:, start:end]
            sq_dists = torch.cdist(mode_slice, mode_slice, p=2).pow(2)
            mask = ~torch.eye(B, dtype=torch.bool, device=fourier_embed.device)
            neg_dists = -self.t * sq_dists
            neg_dists = neg_dists.masked_select(mask).view(B, B - 1)
            unif_n = torch.logsumexp(neg_dists, dim=1).mean() - math.log(B - 1)
            total = total + self.karmonic_weights[n] * unif_n
        return total


# ============================================================================
# DPO Loss
# ============================================================================

class PreferenceLoss(nn.Module):
    def __init__(self, model, tokenizer, beta: float = 0.1):
        super().__init__()
        self.model = model
        self.tokenizer = tokenizer
        self.beta = beta

    def _get_answer_log_prob(self, question_text: str,
                             answer_text: str) -> torch.Tensor:
        prompt = f"Q: {question_text}\nA:"
        full = f"Q: {question_text}\nA: {answer_text}"
        device = next(self.model.parameters()).device
        prompt_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(device)
        full_ids = self.tokenizer(full, return_tensors="pt").input_ids.to(device)
        prompt_len = prompt_ids.shape[1]
        outputs = self.model(full_ids, labels=full_ids.clone())
        logits = outputs.logits
        log_probs = torch.log_softmax(logits, dim=-1)
        answer_tokens = full_ids[:, prompt_len:]
        n_answer = answer_tokens.shape[1]
        if n_answer == 0:
            return torch.tensor(0.0, device=device)
        token_log_probs = log_probs[0, prompt_len - 1:prompt_len - 1 + n_answer, :]
        answer_log_prob = token_log_probs.gather(
            1, answer_tokens[0, :n_answer].unsqueeze(1)
        ).squeeze(1).mean()
        return answer_log_prob

    def forward(self, question: str, correct_answer: str,
                wrong_answer: str) -> tuple[torch.Tensor, float]:
        log_correct = self._get_answer_log_prob(question, correct_answer)
        log_wrong = self._get_answer_log_prob(question, wrong_answer)
        margin = log_correct - log_wrong
        loss = -F.logsigmoid(self.beta * margin)
        return loss, float(margin.detach())


# ============================================================================
# Data loading
# ============================================================================

def load_training_data(tokenizer, max_samples: int = 500):
    ds = load_dataset("truthful_qa", "multiple_choice", split="validation")
    prompts, questions, correct_answers, wrong_answers = [], [], [], []
    device = "cpu"  # tokenize on CPU, move later

    for i, row in enumerate(ds):
        if i >= max_samples:
            break
        q = row["question"]
        choices = row["mc1_targets"]["choices"]
        labels = row["mc1_targets"]["labels"]
        correct_idx = labels.index(1) if 1 in labels else 0
        correct = choices[correct_idx]
        wrong = [c for j, c in enumerate(choices) if j != correct_idx]

        prompt_text = f"Q: {q}\nA:"
        encoded = tokenizer(prompt_text, return_tensors="pt",
                            truncation=True, max_length=256)
        prompts.append(encoded)
        questions.append(q)
        correct_answers.append(correct)
        wrong_answers.append(wrong)

    return prompts, questions, correct_answers, wrong_answers


# ============================================================================
# TruthfulQA MC1 Evaluation
# ============================================================================

def evaluate_mc1(model, tokenizer, max_samples: int | None = None) -> dict:
    ds = load_dataset("truthful_qa", "multiple_choice", split="validation")
    device = next(model.parameters()).device
    model.eval()

    correct, total = 0, 0
    samples = list(ds)
    if max_samples:
        samples = samples[:max_samples]

    with torch.no_grad():
        for row in tqdm(samples, desc="MC1 eval", leave=False):
            q = row["question"]
            choices = row["mc1_targets"]["choices"]
            labels = row["mc1_targets"]["labels"]

            scores = []
            for choice in choices:
                prompt = f"Q: {q}\nA:"
                full = f"Q: {q}\nA: {choice}"
                prompt_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
                full_ids = tokenizer(full, return_tensors="pt").input_ids.to(device)
                p_len = prompt_ids.shape[1]

                out = model(full_ids)
                logits = out.logits
                lp = torch.log_softmax(logits, dim=-1)
                ans_tokens = full_ids[:, p_len:]
                n_ans = ans_tokens.shape[1]
                if n_ans == 0:
                    scores.append(-999.0)
                    continue
                token_lp = lp[0, p_len - 1:p_len - 1 + n_ans, :]
                score = token_lp.gather(
                    1, ans_tokens[0, :n_ans].unsqueeze(1)
                ).squeeze(1).mean().item()
                scores.append(score)

            pred = int(np.argmax(scores))
            if labels[pred] == 1:
                correct += 1
            total += 1

    model.train()
    return {"mc1_accuracy": correct / total if total > 0 else 0.0,
            "mc1_total": total}


# ============================================================================
# WikiText-2 Perplexity
# ============================================================================

def evaluate_perplexity(model, tokenizer, max_chunks: int = 50) -> dict:
    ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
    text = "\n".join([t for t in ds["text"] if t.strip()])
    device = next(model.parameters()).device
    model.eval()

    tokens = tokenizer(text, return_tensors="pt",
                       truncation=True, max_length=100000).input_ids[0]
    chunk_size = 512
    nlls = []
    for i in range(0, min(len(tokens) - 1, chunk_size * max_chunks), chunk_size):
        chunk = tokens[i:i + chunk_size].unsqueeze(0).to(device)
        with torch.no_grad():
            out = model(chunk, labels=chunk)
            nlls.append(out.loss.item())

    model.train()
    ppl = math.exp(np.mean(nlls)) if nlls else float("inf")
    return {"perplexity": ppl, "num_chunks": len(nlls)}


# ============================================================================
# Hallucination Evaluation (TriviaQA / NQ-Open)
# ============================================================================

def normalize_answer(text: str) -> str:
    if not text:
        return ""
    text = text.lower().strip()
    text = re.sub(r"\b(a|an|the)\b", " ", text)
    text = re.sub(r"[^\w\s]", "", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text


def is_abstention(response: str) -> bool:
    if not response:
        return True
    resp_lower = response.lower().strip()
    for phrase in ["i don't know", "i do not know", "i'm not sure",
                   "i am not sure", "i cannot answer", "i can't answer",
                   "unknown", "not available", "no answer",
                   "cannot determine", "insufficient information"]:
        if phrase in resp_lower:
            return True
    return False


def generate_answer(model, tokenizer, question: str, device: torch.device,
                    max_new_tokens: int = 64) -> str:
    prompt = f"Q: {question}\nA:"
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
    with torch.no_grad():
        output_ids = model.generate(
            input_ids, max_new_tokens=max_new_tokens,
            do_sample=False, pad_token_id=tokenizer.pad_token_id,
        )
    generated = output_ids[0, input_ids.shape[1]:]
    return tokenizer.decode(generated, skip_special_tokens=True).strip()


def check_answer(response: str, gold_answers: list[str]) -> str:
    if not response or is_abstention(response):
        return "abstained"
    norm_response = normalize_answer(response)
    for gold in gold_answers:
        norm_gold = normalize_answer(gold)
        if not norm_gold:
            continue
        if norm_gold in norm_response:
            return "correct"
        if len(norm_response) > 2 and norm_response in norm_gold:
            return "correct"
    return "incorrect"


def eval_hallucination(model, tokenizer, benchmark: str,
                       max_samples: int = 100) -> dict:
    device = next(model.parameters()).device
    model.eval()

    if benchmark == "triviaqa":
        ds = load_dataset("trivia_qa", "rc.nocontext", split="validation")
    elif benchmark == "nq_open":
        ds = load_dataset("nq_open", split="validation")
    else:
        raise ValueError(f"Unknown benchmark: {benchmark}")

    if max_samples:
        ds = ds.select(range(min(max_samples, len(ds))))

    correct, incorrect, abstained = 0, 0, 0
    for example in tqdm(ds, desc=f"Halluc:{benchmark}", leave=False):
        question = example["question"]

        if benchmark == "triviaqa":
            answer_obj = example.get("answer", {})
            aliases = answer_obj.get("aliases", [])
            value = answer_obj.get("value", "")
            normalized = answer_obj.get("normalized_aliases", [])
            gold_answers = list(set(
                [a.strip().lower() for a in ([value] + aliases + normalized) if a]
            ))
        else:
            answers_raw = example.get("answer", [])
            if isinstance(answers_raw, str):
                answers_raw = [answers_raw]
            gold_answers = list(set(
                a.strip().lower() for a in answers_raw if a.strip()
            ))

        if not gold_answers:
            continue

        response = generate_answer(model, tokenizer, question, device)
        verdict = check_answer(response, gold_answers)

        if verdict == "correct":
            correct += 1
        elif verdict == "abstained":
            abstained += 1
        else:
            incorrect += 1

    total = correct + incorrect + abstained
    answered = correct + incorrect
    model.train()

    return {
        "benchmark": benchmark,
        "total": total,
        "correct": correct,
        "incorrect": incorrect,
        "abstained": abstained,
        "accuracy": round(correct / total, 4) if total > 0 else 0.0,
        "hallucination_rate": round(
            incorrect / answered, 4
        ) if answered > 0 else 0.0,
        "abstention_rate": round(abstained / total, 4) if total > 0 else 0.0,
    }


# ============================================================================
# Single condition runner
# ============================================================================

def run_condition(buffer_size: int, args, seed: int) -> dict:
    """Run DPO + Karmonic with a specific buffer size and seed.

    buffer_size=0 means DPO-only baseline (no Karmonic at all).
    """

    # Reproducibility
    torch.manual_seed(seed)
    np.random.seed(seed)
    _rng.seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    use_karmonic = buffer_size > 0
    dev = get_device()
    print(f"\n{'='*60}")
    if use_karmonic:
        print(f"Buffer size: {buffer_size} | Seed: {seed} | Device: {dev}")
        print(f"Expected fires in {args.n_steps} steps: ~{args.n_steps // buffer_size}")
    else:
        print(f"DPO-ONLY BASELINE | Seed: {seed} | Device: {dev}")
    print(f"{'='*60}\n")

    # Load model
    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        args.model, trust_remote_code=True,
        torch_dtype=torch.float32
    ).to(dev)

    # LoRA
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM, r=16, lora_alpha=32,
        lora_dropout=0.1, target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]
    )
    model = get_peft_model(model, lora_config)

    # Karmonic components (only if needed)
    hidden_dim = model.config.hidden_size
    torus_head = None
    karmonic_loss_fn = None
    if use_karmonic:
        torus_head = FourierTorusHeadLLM(
            input_dim=hidden_dim, hidden_dim=128,
            torus_dim=2, n_modes=args.n_modes
        ).to(dev)
        karmonic_loss_fn = KarmonicFilterLoss(
            torus_dim=2, n_modes=args.n_modes,
            grid_size=args.grid_size, t_uniformity=2.0
        ).to(dev)

    # DPO
    dpo = PreferenceLoss(model, tokenizer, beta=args.dpo_beta)

    # Optimizer
    trainable_params = list(model.parameters())
    if torus_head is not None:
        trainable_params += list(torus_head.parameters())
    optimizer = torch.optim.AdamW(
        [p for p in trainable_params if p.requires_grad],
        lr=args.lr, weight_decay=0.01
    )

    # Hidden state capture hook (only needed for Karmonic)
    captured_hidden = [None]
    hook_handle = None
    if use_karmonic:
        base = model.base_model.model if hasattr(model, "base_model") else model

        def hook_fn(module, inp, out):
            if isinstance(out, tuple):
                captured_hidden[0] = out[0]
            else:
                captured_hidden[0] = out

        if hasattr(base, "model") and hasattr(base.model, "norm"):
            hook_handle = base.model.norm.register_forward_hook(hook_fn)
        else:
            for name, mod in base.named_modules():
                if "norm" in name.lower() and isinstance(mod, nn.LayerNorm):
                    hook_handle = mod.register_forward_hook(hook_fn)
                    break

    # Load data
    prompts, questions, correct_answers, wrong_answers = load_training_data(
        tokenizer, max_samples=args.train_samples
    )

    # Training loop
    karmonic_buffer = []
    logs = []
    karmonic_fire_count = 0
    t_start = time.time()

    model.train()
    if torus_head is not None:
        torus_head.train()

    for step in tqdm(range(args.n_steps),
                     desc=f"buf={buffer_size} seed={seed}"):
        optimizer.zero_grad()
        idx = step % len(prompts)
        prompt_input = {k: v.to(dev) for k, v in prompts[idx].items()}
        attention_mask = prompt_input["attention_mask"]

        total_loss = torch.tensor(0.0, device=dev, requires_grad=True)
        log = {"step": step}

        # DPO loss
        wrong_ans = (
            _rng.choice(wrong_answers[idx]) if wrong_answers[idx] else ""
        )
        dpo_loss, margin = dpo(questions[idx], correct_answers[idx], wrong_ans)
        total_loss = total_loss + dpo_loss
        log["dpo_loss"] = float(dpo_loss.detach())
        log["preference_margin"] = margin

        # Forward for Karmonic (skip if DPO-only baseline)
        if use_karmonic:
            captured_hidden[0] = None
            outputs = model(**prompt_input)

        if use_karmonic and captured_hidden[0] is not None:
            hidden = captured_hidden[0]
            mask_float = attention_mask.unsqueeze(-1).float()
            pooled = (
                (hidden * mask_float).sum(dim=1) /
                mask_float.sum(dim=1).clamp(min=1)
            )

            karmonic_buffer.append(pooled.detach())
            if len(karmonic_buffer) >= buffer_size:
                context = torch.cat(karmonic_buffer, dim=0)
                live_batch = torch.cat([context, pooled], dim=0)
                pooled_scaled = gradient_scale(live_batch, args.grad_scale)
                angles, fourier = torus_head(pooled_scaled)
                k_loss = karmonic_loss_fn(angles, fourier)
                total_loss = total_loss + args.lambda_karmonic * k_loss
                log["karmonic_loss"] = float(k_loss.detach())
                karmonic_buffer = []
                karmonic_fire_count += 1
            else:
                log["karmonic_loss"] = 0.0

        total_loss.backward()
        torch.nn.utils.clip_grad_norm_(trainable_params, 1.0)
        optimizer.step()

        log["total_loss"] = float(total_loss.detach())
        logs.append(log)

    train_time = time.time() - t_start

    # Evaluation
    label = f"buf={buffer_size}" if buffer_size > 0 else "DPO-only"
    print(f"\nEvaluating {label}, seed={seed}...")
    mc1_results = evaluate_mc1(model, tokenizer, max_samples=args.eval_samples)
    ppl_results = evaluate_perplexity(model, tokenizer)

    # Hallucination benchmarks
    triviaqa_results = {}
    nq_results = {}
    if args.halluc_samples > 0:
        print(f"  TriviaQA hallucination eval ({args.halluc_samples} samples)...")
        triviaqa_results = eval_hallucination(
            model, tokenizer, "triviaqa", max_samples=args.halluc_samples
        )
        print(f"  NQ-Open hallucination eval ({args.halluc_samples} samples)...")
        nq_results = eval_hallucination(
            model, tokenizer, "nq_open", max_samples=args.halluc_samples
        )

    # Cleanup hook
    if hook_handle:
        hook_handle.remove()

    result = {
        "buffer_size": buffer_size,
        "seed": seed,
        "n_steps": args.n_steps,
        "karmonic_fires": karmonic_fire_count,
        "firing_frequency": karmonic_fire_count / args.n_steps if args.n_steps > 0 else 0,
        "effective_duty_cycle": 1.0 / buffer_size if buffer_size > 0 else 0.0,
        "mc1_accuracy": mc1_results["mc1_accuracy"],
        "mc1_total": mc1_results["mc1_total"],
        "perplexity": ppl_results["perplexity"],
        "triviaqa_halluc_rate": triviaqa_results.get("hallucination_rate", None),
        "triviaqa_accuracy": triviaqa_results.get("accuracy", None),
        "nq_halluc_rate": nq_results.get("hallucination_rate", None),
        "nq_accuracy": nq_results.get("accuracy", None),
        "train_time_s": train_time,
        "final_loss": logs[-1]["total_loss"] if logs else None,
        "mean_karmonic_loss": np.mean([
            l["karmonic_loss"] for l in logs
            if l.get("karmonic_loss", 0) != 0
        ]) if any(l.get("karmonic_loss", 0) != 0 for l in logs) else 0.0,
        "training_logs": logs,
    }

    # Free GPU memory
    to_del = [model, dpo, optimizer]
    if torus_head is not None:
        to_del.append(torus_head)
    if karmonic_loss_fn is not None:
        to_del.append(karmonic_loss_fn)
    del to_del
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return result


# ============================================================================
# Main
# ============================================================================

def main():
    parser = argparse.ArgumentParser(
        description="Buffer Frequency Sweep: Topological Coherence as Resonant Filter"
    )

    parser.add_argument(
        "--buffer_sizes", type=str, default="0,1,2,4,8,16,32,64",
        help="Comma-separated buffer sizes to test. 0 = DPO-only baseline (no Karmonic)"
    )
    parser.add_argument(
        "--seeds", type=str, default="42",
        help="Comma-separated seeds for replication (default: 42)"
    )
    parser.add_argument("--model", type=str,
                        default="Qwen/Qwen2.5-0.5B-Instruct")
    parser.add_argument("--n_steps", type=int, default=200,
                        help="Training steps per condition (default: 200)")
    parser.add_argument("--train_samples", type=int, default=500)
    parser.add_argument("--lr", type=float, default=2e-5)
    parser.add_argument("--dpo_beta", type=float, default=0.1)
    parser.add_argument("--lambda_karmonic", type=float, default=0.01)
    parser.add_argument("--grad_scale", type=float, default=0.1)
    parser.add_argument("--n_modes", type=int, default=6)
    parser.add_argument("--grid_size", type=int, default=12)
    parser.add_argument("--eval_samples", type=int, default=None,
                        help="Limit MC1 eval samples (None=all 817)")
    parser.add_argument("--halluc_samples", type=int, default=100,
                        help="TriviaQA/NQ-Open samples per benchmark (0=skip)")
    parser.add_argument("--output_dir", type=str,
                        default="results/buffer-frequency-sweep")

    args = parser.parse_args()

    buffer_sizes = [int(x) for x in args.buffer_sizes.split(",")]
    seeds = [int(x) for x in args.seeds.split(",")]

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

    print(f"Buffer Frequency Sweep")
    print(f"  Buffer sizes: {buffer_sizes}")
    print(f"  Seeds: {seeds}")
    print(f"  Steps: {args.n_steps}")
    print(f"  Model: {args.model}")
    print(f"  Output: {output_dir}")
    print()

    all_results = []

    for buf_size in buffer_sizes:
        for seed in seeds:
            result = run_condition(buf_size, args, seed)

            # Strip verbose logs for summary
            summary = {k: v for k, v in result.items() if k != "training_logs"}
            all_results.append(summary)

            # Save per-condition result (with full logs)
            cond_file = output_dir / f"buf{buf_size}_seed{seed}.json"
            with open(cond_file, "w") as f:
                json.dump(result, f, indent=2, default=str)

            label = f"buf={buf_size}" if buf_size > 0 else "DPO-only"
            halluc_str = ""
            if summary.get('triviaqa_halluc_rate') is not None:
                halluc_str = (f" TQA_H={summary['triviaqa_halluc_rate']:.3f}"
                              f" NQ_H={summary['nq_halluc_rate']:.3f}")
            print(f"\n  {label} seed={seed}: "
                  f"MC1={summary['mc1_accuracy']:.4f} "
                  f"PPL={summary['perplexity']:.2f}"
                  f"{halluc_str} "
                  f"fires={summary['karmonic_fires']} "
                  f"time={summary['train_time_s']:.0f}s")

    # Save combined summary
    summary_file = output_dir / f"sweep_summary_{timestamp}.json"
    with open(summary_file, "w") as f:
        json.dump({
            "experiment": "buffer_frequency_sweep",
            "timestamp": timestamp,
            "hypothesis": "Karmonic regularization exhibits DSP resonance: "
                         "optimal buffer size produces bandpass peak in MC1/PPL",
            "parameters": {
                "buffer_sizes": buffer_sizes,
                "seeds": seeds,
                "n_steps": args.n_steps,
                "model": args.model,
                "lambda_karmonic": args.lambda_karmonic,
                "grad_scale": args.grad_scale,
                "n_modes": args.n_modes,
                "grid_size": args.grid_size,
            },
            "results": all_results,
        }, f, indent=2)

    # Print summary table
    print(f"\n{'='*90}")
    print(f"BUFFER FREQUENCY SWEEP RESULTS")
    print(f"{'='*90}")
    has_halluc = any(r.get('triviaqa_halluc_rate') is not None for r in all_results)
    if has_halluc:
        print(f"{'Buffer':>8} {'Fires':>8} {'MC1':>8} {'TQA_H%':>8} {'NQ_H%':>8} {'PPL':>8} {'Time':>8}")
    else:
        print(f"{'Buffer':>8} {'Fires':>8} {'Duty%':>8} {'MC1':>8} {'PPL':>8} {'Time':>8}")
    print(f"{'-'*90}")

    for r in all_results:
        if has_halluc:
            tqa = f"{r.get('triviaqa_halluc_rate', 0)*100:>7.1f}%" if r.get('triviaqa_halluc_rate') is not None else "    N/A"
            nq = f"{r.get('nq_halluc_rate', 0)*100:>7.1f}%" if r.get('nq_halluc_rate') is not None else "    N/A"
            print(f"{r['buffer_size']:>8} "
                  f"{r['karmonic_fires']:>8} "
                  f"{r['mc1_accuracy']:>8.4f} "
                  f"{tqa:>8} "
                  f"{nq:>8} "
                  f"{r['perplexity']:>8.2f} "
                  f"{r['train_time_s']:>7.0f}s")
        else:
            print(f"{r['buffer_size']:>8} "
                  f"{r['karmonic_fires']:>8} "
                  f"{r['effective_duty_cycle']*100:>7.1f}% "
                  f"{r['mc1_accuracy']:>8.4f} "
                  f"{r['perplexity']:>8.2f} "
                  f"{r['train_time_s']:>7.0f}s")

    print(f"\nResults saved to: {summary_file}")
    print(f"\nLook for a PEAK in MC1 around buffer=8.")
    print(f"If MC1 degrades for both buffer=1 and buffer=64,")
    print(f"that's resonance — a bandpass characteristic in training dynamics.")


if __name__ == "__main__":
    main()