"""
Phase 5A: Vector Credit Field Audit.

Verify that the vector field's gains are real, not implementation artifacts.

4 mandatory sanity checks:
A. Train/eval direction split (independent random directions)
B. Shuffled-target control (permute g_j within batch)
C. No-terminal ablation (L_term = 0)
D. One-sided vs symmetric finite difference
"""
import os
import sys
import json
import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import copy

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from models.value_net import ValueNet, SinusoidalTimeEmbed, create_ema_model, update_ema
from metrics.credit_metrics import (
    cosine_similarity_batch, perturbation_correlation, nudging_test
)


# =============================================================================
# Synthetic teacher-student
# =============================================================================
class TeacherNet(nn.Module):
    def __init__(self, d_hidden, num_classes, num_blocks, alpha=1.0, seed=0):
        super().__init__()
        self.d_hidden = d_hidden
        self.num_blocks = num_blocks
        self.alpha = alpha
        rng = torch.Generator().manual_seed(seed)
        self.Ws = nn.ParameterList()
        for _ in range(num_blocks):
            W = torch.randn(d_hidden, d_hidden, generator=rng) * 0.3 / (d_hidden ** 0.5)
            U, S, Vh = torch.linalg.svd(W, full_matrices=False)
            S_clamped = S.clamp(max=0.3)
            W = U @ torch.diag(S_clamped) @ Vh
            self.Ws.append(nn.Parameter(W, requires_grad=False))
        self.U = nn.Parameter(
            torch.randn(num_classes, d_hidden, generator=rng) / (d_hidden ** 0.5),
            requires_grad=False)

    def phi(self, z):
        return (1 - self.alpha) * z + self.alpha * torch.tanh(z)

    def forward(self, x):
        h = x
        for W in self.Ws:
            h = h + self.phi(h @ W.T)
        return h @ self.U.T


class StudentBlock(nn.Module):
    def __init__(self, d_hidden, alpha=1.0):
        super().__init__()
        self.ln = nn.LayerNorm(d_hidden)
        self.w = nn.Linear(d_hidden, d_hidden, bias=False)
        nn.init.normal_(self.w.weight, std=0.01)
        self.alpha = alpha

    def phi(self, z):
        return (1 - self.alpha) * z + self.alpha * torch.tanh(z)

    def forward(self, h):
        return self.w(self.phi(self.ln(h)))


class StudentNet(nn.Module):
    def __init__(self, d_hidden, num_classes, num_blocks, alpha=1.0):
        super().__init__()
        self.blocks = nn.ModuleList([StudentBlock(d_hidden, alpha) for _ in range(num_blocks)])
        self.out_head = nn.Linear(d_hidden, num_classes)
        self.d_hidden = d_hidden
        self.num_blocks = num_blocks

    def forward(self, x, return_hidden=False):
        h = x
        hiddens = [h] if return_hidden else None
        for block in self.blocks:
            f = block(h)
            h = h + f
            if return_hidden:
                hiddens.append(h)
        logits = self.out_head(h)
        if return_hidden:
            return logits, hiddens
        return logits

    def forward_from_layer(self, h, start_layer):
        for i in range(start_layer, self.num_blocks):
            f = self.blocks[i](h)
            h = h + f
        return self.out_head(h)


class VectorCreditNet(nn.Module):
    """Direct vector credit field: a_phi(h_l, t_l, s) -> R^d."""
    def __init__(self, d_hidden, s_dim, time_embed_dim=32, hidden_dim=256, num_layers=3):
        super().__init__()
        self.ln = nn.LayerNorm(d_hidden)
        self.time_embed = SinusoidalTimeEmbed(time_embed_dim)
        input_dim = d_hidden + time_embed_dim + s_dim
        layers = []
        for i in range(num_layers):
            in_d = input_dim if i == 0 else hidden_dim
            layers.append(nn.Linear(in_d, hidden_dim))
            layers.append(nn.GELU())
        layers.append(nn.Linear(hidden_dim, d_hidden))
        self.net = nn.Sequential(*layers)

    def forward(self, h, t, s):
        h_normed = self.ln(h)
        t_emb = self.time_embed(t)
        inp = torch.cat([h_normed, t_emb, s], dim=-1)
        return self.net(inp)


def generate_batch(teacher, d_hidden, num_classes, batch_size, device):
    x = torch.randn(batch_size, d_hidden, device=device)
    with torch.no_grad():
        teacher_logits = teacher(x)
        y = teacher_logits.argmax(dim=-1)
    return x, y


# =============================================================================
# Training: vector field with audit controls
# =============================================================================
def train_vector_field_audit(model, teacher, device, args, M=4,
                              use_terminal=True,
                              shuffle_targets=False,
                              use_central_diff=True,
                              tag='vec'):
    """
    Train vector credit field with configurable audit controls.

    Args:
        use_terminal: if False, L_term = 0 (no-terminal ablation)
        shuffle_targets: if True, permute g_j within batch (leak check)
        use_central_diff: if True, central difference; if False, one-sided
        tag: label for printing
    """
    d = model.d_hidden
    L = model.num_blocks
    num_classes = args.num_classes

    vector_net = VectorCreditNet(d_hidden=d, s_dim=num_classes, time_embed_dim=32,
                                  hidden_dim=256, num_layers=3).to(device)

    Bs = [torch.randn(d, num_classes, device=device) / np.sqrt(num_classes) for _ in range(L)]

    block_opts = [optim.AdamW(b.parameters(), lr=args.lr, weight_decay=0.01) for b in model.blocks]
    head_opt = optim.AdamW(model.out_head.parameters(), lr=args.lr, weight_decay=0.01)
    vec_opt = optim.Adam(vector_net.parameters(), lr=args.lr_fb)

    warmup_epochs = max(1, int(args.epochs * args.warmup_ratio))
    eps = args.pert_eps
    beta = args.pert_beta

    for epoch in range(1, args.epochs + 1):
        model.train()
        vector_net.train()

        if epoch <= warmup_epochs:
            credit_blend = 0.0
        else:
            credit_blend = min(1.0, (epoch - warmup_epochs) / max(1, warmup_epochs))

        total_loss, correct, total = 0, 0, 0
        total_vloss = 0

        for _ in range(args.steps_per_epoch):
            x, y = generate_batch(teacher, d, num_classes, args.batch_size, device)
            batch = x.size(0)

            with torch.no_grad():
                logits, hiddens = model(x, return_hidden=True)
                loss_val = F.cross_entropy(logits, y)
                e_T = logits.softmax(dim=-1)
                e_T[torch.arange(batch), y] -= 1
                s = e_T.detach()

            hL_det = hiddens[-1].detach()

            # --- Terminal matching ---
            loss_term = torch.tensor(0.0, device=device)
            if use_terminal:
                t_L = torch.ones(batch, device=device)
                a_terminal = vector_net(hL_det, t_L, s)
                hL_req = hL_det.clone().requires_grad_(True)
                logits_tgt = model.out_head(hL_req)
                ce = F.cross_entropy(logits_tgt, y, reduction='sum')
                delta_L = torch.autograd.grad(ce, hL_req, create_graph=False)[0].detach()
                loss_term = ((a_terminal - delta_L) ** 2).sum(dim=-1).mean()

            # --- Perturbation directional targets ---
            # IMPORTANT: training directions are sampled fresh each step.
            # Evaluation uses independently sampled directions (see compute_diagnostics).
            loss_proj = torch.tensor(0.0, device=device)
            for l in range(L):
                h_l_det = hiddens[l].detach()
                t_l = torch.full((batch,), l / L, device=device)
                a_l = vector_net(h_l_det, t_l, s)

                layer_proj_loss = 0.0
                for _ in range(M):
                    v = torch.randn_like(h_l_det)
                    v = v / (v.norm(dim=-1, keepdim=True) + 1e-8)

                    with torch.no_grad():
                        if use_central_diff:
                            # Central difference: [loss(h+eps*v) - loss(h-eps*v)] / (2*eps)
                            logits_plus = model.forward_from_layer(h_l_det + eps * v, l)
                            loss_plus = F.cross_entropy(logits_plus, y, reduction='none')
                            logits_minus = model.forward_from_layer(h_l_det - eps * v, l)
                            loss_minus = F.cross_entropy(logits_minus, y, reduction='none')
                            g_j = (loss_plus - loss_minus) / (2 * eps)
                        else:
                            # One-sided difference: [loss(h+eps*v) - loss(h)] / eps
                            logits_base = model.forward_from_layer(h_l_det, l)
                            loss_base = F.cross_entropy(logits_base, y, reduction='none')
                            logits_plus = model.forward_from_layer(h_l_det + eps * v, l)
                            loss_plus = F.cross_entropy(logits_plus, y, reduction='none')
                            g_j = (loss_plus - loss_base) / eps

                    # Shuffled-target control: permute g_j within batch
                    if shuffle_targets:
                        perm = torch.randperm(batch, device=device)
                        g_j = g_j[perm]

                    pred_j = (a_l * v).sum(dim=-1)
                    layer_proj_loss = layer_proj_loss + ((pred_j - g_j.detach()) ** 2).mean()

                loss_proj = loss_proj + layer_proj_loss / M
            loss_proj = loss_proj / L

            vec_loss = loss_term + beta * loss_proj
            vec_opt.zero_grad()
            vec_loss.backward()
            torch.nn.utils.clip_grad_norm_(vector_net.parameters(), 1.0)
            vec_opt.step()
            total_vloss += vec_loss.item() * batch

            # --- Block updates ---
            with torch.no_grad():
                vec_credits = []
                for l in range(L):
                    h_l_det = hiddens[l].detach()
                    t_l = torch.full((batch,), l / L, device=device)
                    a_l = vector_net(h_l_det, t_l, s)
                    vec_credits.append(a_l.detach())

            dfa_credits = [(e_T @ Bs[l].T).detach() for l in range(L)]

            credits = []
            for l in range(L):
                if credit_blend >= 1.0:
                    credits.append(vec_credits[l])
                elif credit_blend <= 0.0:
                    credits.append(dfa_credits[l])
                else:
                    vc_rms = (vec_credits[l] ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6
                    dfa_rms = (dfa_credits[l] ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6
                    credits.append(credit_blend * vec_credits[l] / vc_rms +
                                   (1 - credit_blend) * dfa_credits[l] / dfa_rms)

            logits_out = model.out_head(hL_det)
            loss_out = F.cross_entropy(logits_out, y)
            head_opt.zero_grad()
            loss_out.backward()
            head_opt.step()

            for l in range(L):
                h_l = hiddens[l].detach()
                a = credits[l]
                rms = (a ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6
                f_l = model.blocks[l](h_l)
                local_loss = (f_l * (a / rms)).sum(dim=-1).mean()
                block_opts[l].zero_grad()
                local_loss.backward()
                torch.nn.utils.clip_grad_norm_(model.blocks[l].parameters(), 1.0)
                block_opts[l].step()

            total_loss += loss_val.item() * batch
            correct += (logits.argmax(1) == y).sum().item()
            total += batch

        if epoch % 20 == 0 or epoch == 1:
            acc = correct / total
            print(f"    [{tag}] Ep {epoch}: loss={total_loss/total:.4f}, acc={acc:.4f}, "
                  f"vloss={total_vloss/total:.6f}")

    return vector_net


def train_scalar_cb(model, teacher, device, args):
    """Scalar credit bridge baseline."""
    d = model.d_hidden
    L = model.num_blocks
    num_classes = args.num_classes

    value_net = ValueNet(d_hidden=d, s_dim=num_classes, time_embed_dim=32,
                         hidden_dim=256, num_layers=3).to(device)
    value_net_ema = create_ema_model(value_net)

    Bs = [torch.randn(d, num_classes, device=device) / np.sqrt(num_classes) for _ in range(L)]

    block_opts = [optim.AdamW(b.parameters(), lr=args.lr, weight_decay=0.01) for b in model.blocks]
    head_opt = optim.AdamW(model.out_head.parameters(), lr=args.lr, weight_decay=0.01)
    value_opt = optim.Adam(value_net.parameters(), lr=args.lr_fb)

    warmup_epochs = max(1, int(args.epochs * args.warmup_ratio))

    for epoch in range(1, args.epochs + 1):
        model.train()
        value_net.train()

        if epoch <= warmup_epochs:
            credit_blend = 0.0
        else:
            credit_blend = min(1.0, (epoch - warmup_epochs) / max(1, warmup_epochs))

        total_loss, correct, total = 0, 0, 0
        for _ in range(args.steps_per_epoch):
            x, y = generate_batch(teacher, d, num_classes, args.batch_size, device)
            batch = x.size(0)

            with torch.no_grad():
                logits, hiddens = model(x, return_hidden=True)
                loss_val = F.cross_entropy(logits, y)
                e_T = logits.softmax(dim=-1)
                e_T[torch.arange(batch), y] -= 1
                s = e_T.detach()
                true_loss = F.cross_entropy(logits, y, reduction='none').detach()

            hL_det = hiddens[-1].detach()
            t_L = torch.ones(batch, device=device)
            V_term = value_net(hL_det, t_L, s)
            loss_term = ((V_term - true_loss) ** 2).mean()

            hL_req = hL_det.clone().requires_grad_(True)
            V_at_L = value_net(hL_req, t_L, s)
            grad_V_L = torch.autograd.grad(V_at_L.sum(), hL_req, create_graph=True)[0]
            hL_req2 = hL_det.clone().requires_grad_(True)
            logits_tgt = model.out_head(hL_req2)
            ce = F.cross_entropy(logits_tgt, y, reduction='sum')
            a_L_exact = torch.autograd.grad(ce, hL_req2, create_graph=False)[0].detach()
            loss_tgrad = ((grad_V_L - a_L_exact) ** 2).sum(dim=-1).mean()

            loss_bridge = 0.0
            for l in range(L):
                h_l_det = hiddens[l].detach()
                t_l = torch.full((batch,), l / L, device=device)
                t_next = torch.full((batch,), (l + 1) / L, device=device)
                V_l = value_net(h_l_det, t_l, s)
                with torch.no_grad():
                    h_next = hiddens[l + 1].detach()
                    log_terms = []
                    for k in range(args.K):
                        noise = args.sigma_bridge * torch.randn_like(h_next)
                        V_next = value_net_ema(h_next + noise, t_next, s)
                        log_terms.append(-V_next / args.lam)
                    log_stack = torch.stack(log_terms, dim=-1)
                    V_target = -args.lam * (torch.logsumexp(log_stack, dim=-1) - np.log(args.K))
                loss_bridge += ((V_l - V_target.detach()) ** 2).mean()
            loss_bridge /= L

            vloss = loss_term + loss_bridge + args.term_grad_weight * loss_tgrad
            value_opt.zero_grad()
            vloss.backward()
            torch.nn.utils.clip_grad_norm_(value_net.parameters(), 1.0)
            value_opt.step()
            update_ema(value_net, value_net_ema, args.ema_momentum)

            cb_credits = []
            for l in range(L):
                h_l_det = hiddens[l].detach().requires_grad_(True)
                t_l = torch.full((batch,), l / L, device=device)
                V_l = value_net(h_l_det, t_l, s)
                a_l = torch.autograd.grad(V_l.sum(), h_l_det, create_graph=False)[0]
                cb_credits.append(a_l.detach())

            dfa_credits = [(e_T @ Bs[l].T).detach() for l in range(L)]

            credits = []
            for l in range(L):
                if credit_blend >= 1.0:
                    credits.append(cb_credits[l])
                elif credit_blend <= 0.0:
                    credits.append(dfa_credits[l])
                else:
                    cb_rms = (cb_credits[l] ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6
                    dfa_rms = (dfa_credits[l] ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6
                    credits.append(credit_blend * cb_credits[l] / cb_rms +
                                   (1 - credit_blend) * dfa_credits[l] / dfa_rms)

            logits_out = model.out_head(hL_det)
            loss_out = F.cross_entropy(logits_out, y)
            head_opt.zero_grad()
            loss_out.backward()
            head_opt.step()

            for l in range(L):
                h_l = hiddens[l].detach()
                a = credits[l]
                rms = (a ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6
                f_l = model.blocks[l](h_l)
                local_loss = (f_l * (a / rms)).sum(dim=-1).mean()
                block_opts[l].zero_grad()
                local_loss.backward()
                torch.nn.utils.clip_grad_norm_(model.blocks[l].parameters(), 1.0)
                block_opts[l].step()

            total_loss += loss_val.item() * batch
            correct += (logits.argmax(1) == y).sum().item()
            total += batch

        if epoch % 20 == 0 or epoch == 1:
            print(f"    [scalar_cb] Ep {epoch}: loss={total_loss/total:.4f}, acc={correct/total:.4f}")

    return value_net


def train_dfa(model, teacher, device, args):
    """DFA baseline."""
    d = model.d_hidden
    L = model.num_blocks
    num_classes = args.num_classes
    Bs = [torch.randn(d, num_classes, device=device) / np.sqrt(num_classes) for _ in range(L)]

    block_opts = [optim.AdamW(b.parameters(), lr=args.lr, weight_decay=0.01) for b in model.blocks]
    head_opt = optim.AdamW(model.out_head.parameters(), lr=args.lr, weight_decay=0.01)

    for epoch in range(1, args.epochs + 1):
        model.train()
        total_loss, correct, total = 0, 0, 0
        for _ in range(args.steps_per_epoch):
            x, y = generate_batch(teacher, d, num_classes, args.batch_size, device)
            batch = x.size(0)
            with torch.no_grad():
                logits, hiddens = model(x, return_hidden=True)
                loss_val = F.cross_entropy(logits, y)
                e_T = logits.softmax(dim=-1)
                e_T[torch.arange(batch), y] -= 1
            hL_det = hiddens[-1].detach()
            logits_out = model.out_head(hL_det)
            loss_out = F.cross_entropy(logits_out, y)
            head_opt.zero_grad()
            loss_out.backward()
            head_opt.step()
            for l in range(L):
                h_l = hiddens[l].detach()
                a = (e_T @ Bs[l].T).detach()
                rms = (a ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6
                f_l = model.blocks[l](h_l)
                local_loss = (f_l * (a / rms)).sum(dim=-1).mean()
                block_opts[l].zero_grad()
                local_loss.backward()
                torch.nn.utils.clip_grad_norm_(model.blocks[l].parameters(), 1.0)
                block_opts[l].step()
            total_loss += loss_val.item() * batch
            correct += (logits.argmax(1) == y).sum().item()
            total += batch
        if epoch % 20 == 0 or epoch == 1:
            print(f"    [DFA] Ep {epoch}: loss={total_loss/total:.4f}, acc={correct/total:.4f}")
    return Bs


# =============================================================================
# Diagnostics — uses INDEPENDENT eval directions (check A)
# =============================================================================
def compute_diagnostics(model, teacher, device, method_name, args,
                        value_net=None, vector_net=None, dfa_Bs=None):
    """
    Compute Gamma, rho, nudging per layer.
    IMPORTANT: perturbation_correlation uses its own freshly-sampled directions,
    completely independent of any training directions. This ensures check A.
    """
    model.eval()
    if value_net is not None:
        value_net.eval()
    if vector_net is not None:
        vector_net.eval()

    d = model.d_hidden
    L = model.num_blocks
    num_classes = args.num_classes

    # Use a fixed eval seed different from training
    eval_rng = torch.Generator(device=device)
    eval_rng.manual_seed(99999)

    x = torch.randn(512, d, device=device, generator=eval_rng)
    with torch.no_grad():
        teacher_logits = teacher(x)
        y = teacher_logits.argmax(dim=-1)
    batch = x.size(0)

    # BP gradients (evaluation only — never used for training)
    h = x.detach().requires_grad_(True)
    hiddens_bp = [h]
    for block in model.blocks:
        f = block(hiddens_bp[-1])
        h_next = hiddens_bp[-1] + f
        hiddens_bp.append(h_next)
    logits_bp = model.out_head(hiddens_bp[-1])
    loss_bp = F.cross_entropy(logits_bp, y)
    grads = torch.autograd.grad(loss_bp, hiddens_bp, retain_graph=False)
    bp_grads = {l: grads[l].detach().clone() for l in range(L + 1)}

    with torch.no_grad():
        logits, hiddens = model(x, return_hidden=True)
        e_T = logits.softmax(dim=-1)
        e_T[torch.arange(batch), y] -= 1
        s = e_T.detach()

    results = {'bp_cosine': [], 'perturbation_rho': [], 'nudging': []}

    for l in range(L):
        h_l = hiddens[l].detach()
        t_l = torch.full((batch,), l / L, device=device)

        if method_name == 'dfa':
            a_l = (s @ dfa_Bs[l].T).detach()
        elif method_name == 'scalar_cb':
            h_l_req = h_l.clone().requires_grad_(True)
            V_l = value_net(h_l_req, t_l, s)
            a_l = torch.autograd.grad(V_l.sum(), h_l_req, create_graph=False)[0].detach()
        elif method_name.startswith('vec'):
            a_l = vector_net(h_l, t_l, s).detach()
        else:
            raise ValueError(f"Unknown: {method_name}")

        bp_cos = cosine_similarity_batch(a_l, bp_grads[l])
        results['bp_cosine'].append(float(bp_cos))

        # perturbation_correlation uses its own random directions internally
        # (from metrics/credit_metrics.py — independent of training directions)
        def make_fwd_fn(start_l):
            def fwd_fn(h):
                with torch.no_grad():
                    logits = model.forward_from_layer(h, start_l)
                    return F.cross_entropy(logits, y, reduction='none')
            return fwd_fn

        fwd_fn = make_fwd_fn(l)
        rho = perturbation_correlation(h_l, a_l, fwd_fn, epsilon=1e-3, M=32)
        results['perturbation_rho'].append(float(rho))

        nud = nudging_test(h_l, a_l, fwd_fn, eta=0.003)
        results['nudging'].append(float(nud))

    return results


# =============================================================================
# Main
# =============================================================================
def run_experiment(args):
    device = torch.device(f'cuda:{args.gpu}' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    os.makedirs(args.output_dir, exist_ok=True)

    all_results = []

    for L in args.depths:
        for seed in args.seeds:
            print(f"\n{'='*60}")
            print(f"L={L}, seed={seed}")
            print(f"{'='*60}")

            teacher = TeacherNet(args.d_hidden, args.num_classes, L,
                                 alpha=args.alpha, seed=seed * 1000).to(device)

            # --- DFA ---
            print("\n  --- DFA ---")
            torch.manual_seed(seed)
            np.random.seed(seed)
            torch.cuda.manual_seed_all(seed)
            model_dfa = StudentNet(args.d_hidden, args.num_classes, L, alpha=args.alpha).to(device)
            Bs = train_dfa(model_dfa, teacher, device, args)
            diag = compute_diagnostics(model_dfa, teacher, device, 'dfa', args, dfa_Bs=Bs)
            r = {'method': 'dfa', 'L': L, 'seed': seed,
                 'mean_gamma': float(np.mean(diag['bp_cosine'])),
                 'mean_rho': float(np.mean(diag['perturbation_rho'])),
                 'mean_nudge': float(np.mean(diag['nudging'])),
                 'per_layer': diag}
            print(f"    Result: Gamma={r['mean_gamma']:.4f}, rho={r['mean_rho']:.4f}, nudge={r['mean_nudge']:.6f}")
            all_results.append(r)

            # --- Scalar CB ---
            print("\n  --- Scalar CB ---")
            torch.manual_seed(seed)
            np.random.seed(seed)
            torch.cuda.manual_seed_all(seed)
            model_cb = StudentNet(args.d_hidden, args.num_classes, L, alpha=args.alpha).to(device)
            vnet = train_scalar_cb(model_cb, teacher, device, args)
            diag = compute_diagnostics(model_cb, teacher, device, 'scalar_cb', args, value_net=vnet)
            r = {'method': 'scalar_cb', 'L': L, 'seed': seed,
                 'mean_gamma': float(np.mean(diag['bp_cosine'])),
                 'mean_rho': float(np.mean(diag['perturbation_rho'])),
                 'mean_nudge': float(np.mean(diag['nudging'])),
                 'per_layer': diag}
            print(f"    Result: Gamma={r['mean_gamma']:.4f}, rho={r['mean_rho']:.4f}, nudge={r['mean_nudge']:.6f}")
            all_results.append(r)

            # --- Vector Field M4 (central diff, with terminal) ---
            print("\n  --- vec_eT_M4 (central, +term) ---")
            torch.manual_seed(seed)
            np.random.seed(seed)
            torch.cuda.manual_seed_all(seed)
            model_v4 = StudentNet(args.d_hidden, args.num_classes, L, alpha=args.alpha).to(device)
            vnet4 = train_vector_field_audit(model_v4, teacher, device, args, M=4,
                                              use_terminal=True, shuffle_targets=False,
                                              use_central_diff=True, tag='vec_eT_M4')
            diag = compute_diagnostics(model_v4, teacher, device, 'vec_eT_M4', args, vector_net=vnet4)
            r = {'method': 'vec_eT_M4', 'L': L, 'seed': seed,
                 'mean_gamma': float(np.mean(diag['bp_cosine'])),
                 'mean_rho': float(np.mean(diag['perturbation_rho'])),
                 'mean_nudge': float(np.mean(diag['nudging'])),
                 'per_layer': diag}
            print(f"    Result: Gamma={r['mean_gamma']:.4f}, rho={r['mean_rho']:.4f}, nudge={r['mean_nudge']:.6f}")
            all_results.append(r)

            # --- Vector Field M8 (central diff, with terminal) ---
            if 8 in args.M_values:
                print("\n  --- vec_eT_M8 (central, +term) ---")
                torch.manual_seed(seed)
                np.random.seed(seed)
                torch.cuda.manual_seed_all(seed)
                model_v8 = StudentNet(args.d_hidden, args.num_classes, L, alpha=args.alpha).to(device)
                vnet8 = train_vector_field_audit(model_v8, teacher, device, args, M=8,
                                                  use_terminal=True, shuffle_targets=False,
                                                  use_central_diff=True, tag='vec_eT_M8')
                diag = compute_diagnostics(model_v8, teacher, device, 'vec_eT_M8', args, vector_net=vnet8)
                r = {'method': 'vec_eT_M8', 'L': L, 'seed': seed,
                     'mean_gamma': float(np.mean(diag['bp_cosine'])),
                     'mean_rho': float(np.mean(diag['perturbation_rho'])),
                     'mean_nudge': float(np.mean(diag['nudging'])),
                     'per_layer': diag}
                print(f"    Result: Gamma={r['mean_gamma']:.4f}, rho={r['mean_rho']:.4f}, nudge={r['mean_nudge']:.6f}")
                all_results.append(r)

            # =================================================================
            # SANITY CHECKS (only for first seed to save time, unless full mode)
            # =================================================================
            if seed == args.seeds[0] or args.full_audit:
                # --- Check B: Shuffled-target control ---
                print("\n  --- vec_eT_M4_shuffleCtrl ---")
                torch.manual_seed(seed)
                np.random.seed(seed)
                torch.cuda.manual_seed_all(seed)
                model_shuf = StudentNet(args.d_hidden, args.num_classes, L, alpha=args.alpha).to(device)
                vnet_shuf = train_vector_field_audit(model_shuf, teacher, device, args, M=4,
                                                      use_terminal=True, shuffle_targets=True,
                                                      use_central_diff=True, tag='vec_shuffleCtrl')
                diag = compute_diagnostics(model_shuf, teacher, device, 'vec_shuffleCtrl', args, vector_net=vnet_shuf)
                r = {'method': 'vec_eT_M4_shuffleCtrl', 'L': L, 'seed': seed,
                     'mean_gamma': float(np.mean(diag['bp_cosine'])),
                     'mean_rho': float(np.mean(diag['perturbation_rho'])),
                     'mean_nudge': float(np.mean(diag['nudging'])),
                     'per_layer': diag}
                print(f"    Result: Gamma={r['mean_gamma']:.4f}, rho={r['mean_rho']:.4f}, nudge={r['mean_nudge']:.6f}")
                all_results.append(r)

                # --- Check C: No-terminal ablation ---
                print("\n  --- vec_eT_M4_noTerm ---")
                torch.manual_seed(seed)
                np.random.seed(seed)
                torch.cuda.manual_seed_all(seed)
                model_nt = StudentNet(args.d_hidden, args.num_classes, L, alpha=args.alpha).to(device)
                vnet_nt = train_vector_field_audit(model_nt, teacher, device, args, M=4,
                                                    use_terminal=False, shuffle_targets=False,
                                                    use_central_diff=True, tag='vec_noTerm')
                diag = compute_diagnostics(model_nt, teacher, device, 'vec_noTerm', args, vector_net=vnet_nt)
                r = {'method': 'vec_eT_M4_noTerm', 'L': L, 'seed': seed,
                     'mean_gamma': float(np.mean(diag['bp_cosine'])),
                     'mean_rho': float(np.mean(diag['perturbation_rho'])),
                     'mean_nudge': float(np.mean(diag['nudging'])),
                     'per_layer': diag}
                print(f"    Result: Gamma={r['mean_gamma']:.4f}, rho={r['mean_rho']:.4f}, nudge={r['mean_nudge']:.6f}")
                all_results.append(r)

                # --- Check D: One-sided difference ---
                print("\n  --- vec_eT_M4_onesided ---")
                torch.manual_seed(seed)
                np.random.seed(seed)
                torch.cuda.manual_seed_all(seed)
                model_os = StudentNet(args.d_hidden, args.num_classes, L, alpha=args.alpha).to(device)
                vnet_os = train_vector_field_audit(model_os, teacher, device, args, M=4,
                                                    use_terminal=True, shuffle_targets=False,
                                                    use_central_diff=False, tag='vec_onesided')
                diag = compute_diagnostics(model_os, teacher, device, 'vec_onesided', args, vector_net=vnet_os)
                r = {'method': 'vec_eT_M4_onesided', 'L': L, 'seed': seed,
                     'mean_gamma': float(np.mean(diag['bp_cosine'])),
                     'mean_rho': float(np.mean(diag['perturbation_rho'])),
                     'mean_nudge': float(np.mean(diag['nudging'])),
                     'per_layer': diag}
                print(f"    Result: Gamma={r['mean_gamma']:.4f}, rho={r['mean_rho']:.4f}, nudge={r['mean_nudge']:.6f}")
                all_results.append(r)

    # =================================================================
    # Summary
    # =================================================================
    print(f"\n{'='*80}")
    print("AUDIT SUMMARY")
    print(f"{'='*80}")
    print(f"{'Method':<30} {'L':>3} {'seed':>5} {'Gamma':>8} {'rho':>8} {'nudge':>10}")
    print("-" * 70)
    for r in all_results:
        print(f"{r['method']:<30} {r['L']:>3} {r['seed']:>5} "
              f"{r['mean_gamma']:>8.4f} {r['mean_rho']:>8.4f} {r['mean_nudge']:>10.6f}")

    # Check verdicts
    print(f"\n{'='*60}")
    print("SANITY CHECK VERDICTS")
    print(f"{'='*60}")

    for L in args.depths:
        seed0 = args.seeds[0]
        vec_main = [r for r in all_results if r['method'] == 'vec_eT_M4' and r['L'] == L and r['seed'] == seed0]
        scalar_cb = [r for r in all_results if r['method'] == 'scalar_cb' and r['L'] == L and r['seed'] == seed0]
        shuf = [r for r in all_results if r['method'] == 'vec_eT_M4_shuffleCtrl' and r['L'] == L and r['seed'] == seed0]
        noterm = [r for r in all_results if r['method'] == 'vec_eT_M4_noTerm' and r['L'] == L and r['seed'] == seed0]
        onesided = [r for r in all_results if r['method'] == 'vec_eT_M4_onesided' and r['L'] == L and r['seed'] == seed0]

        if not vec_main or not scalar_cb:
            continue
        v = vec_main[0]
        cb = scalar_cb[0]

        print(f"\n  L={L}:")
        delta_gamma = v['mean_gamma'] - cb['mean_gamma']
        delta_rho = v['mean_rho'] - cb['mean_rho']
        print(f"    vec_M4 vs scalar_cb: delta_Gamma={delta_gamma:+.4f}, delta_rho={delta_rho:+.4f}")

        if shuf:
            s = shuf[0]
            print(f"    Check B (shuffle): Gamma={s['mean_gamma']:.4f}, rho={s['mean_rho']:.4f}")
            if s['mean_gamma'] < v['mean_gamma'] * 0.5 and s['mean_rho'] < v['mean_rho'] * 0.5:
                print(f"      -> PASS: shuffled control collapses (Gamma dropped by {v['mean_gamma']-s['mean_gamma']:.3f})")
            else:
                print(f"      -> FAIL: shuffled control too close to main result!")

        if noterm:
            n = noterm[0]
            print(f"    Check C (noTerm):  Gamma={n['mean_gamma']:.4f}, rho={n['mean_rho']:.4f}")
            if n['mean_gamma'] < v['mean_gamma'] * 0.8:
                print(f"      -> PASS: terminal matching contributes (Gamma dropped by {v['mean_gamma']-n['mean_gamma']:.3f})")
            else:
                print(f"      -> NOTE: terminal removal didn't collapse result. Perturbation target alone is sufficient.")

        if onesided:
            o = onesided[0]
            print(f"    Check D (onesided): Gamma={o['mean_gamma']:.4f}, rho={o['mean_rho']:.4f}")
            if abs(o['mean_gamma'] - v['mean_gamma']) < 0.15:
                print(f"      -> PASS: one-sided ≈ central (difference = {abs(o['mean_gamma']-v['mean_gamma']):.3f})")
            else:
                print(f"      -> NOTE: one-sided differs from central by {abs(o['mean_gamma']-v['mean_gamma']):.3f}")

    # Final verdict
    print(f"\n{'='*60}")
    print("OVERALL AUDIT VERDICT")
    print(f"{'='*60}")
    all_pass = True
    for L in args.depths:
        for seed in args.seeds:
            v = [r for r in all_results if r['method'] == 'vec_eT_M4' and r['L'] == L and r['seed'] == seed]
            cb = [r for r in all_results if r['method'] == 'scalar_cb' and r['L'] == L and r['seed'] == seed]
            if v and cb:
                dg = v[0]['mean_gamma'] - cb[0]['mean_gamma']
                dr = v[0]['mean_rho'] - cb[0]['mean_rho']
                if dg < 0.2 or dr < 0.2:
                    print(f"  L={L} seed={seed}: delta_Gamma={dg:.3f}, delta_rho={dr:.3f} - BELOW THRESHOLD")
                    all_pass = False
                else:
                    print(f"  L={L} seed={seed}: delta_Gamma={dg:.3f}, delta_rho={dr:.3f} - PASS")

    shuf_results = [r for r in all_results if 'shuffleCtrl' in r['method']]
    for s in shuf_results:
        if s['mean_rho'] > 0.3:
            print(f"  SHUFFLE CONTROL WARNING: L={s['L']} rho={s['mean_rho']:.3f} too high!")
            all_pass = False

    if all_pass:
        print("\n  AUDIT PASSED. Vector field gains are real.")
    else:
        print("\n  AUDIT FAILED or INCOMPLETE. Investigate before proceeding.")

    # Save
    save_data = []
    for r in all_results:
        save_r = {k: v for k, v in r.items() if k != 'per_layer'}
        save_r['per_layer_gamma'] = r['per_layer']['bp_cosine']
        save_r['per_layer_rho'] = r['per_layer']['perturbation_rho']
        save_r['per_layer_nudge'] = r['per_layer']['nudging']
        save_data.append(save_r)

    out_path = os.path.join(args.output_dir, 'audit_results.json')
    with open(out_path, 'w') as f:
        json.dump(save_data, f, indent=2)
    print(f"\nResults saved to {out_path}")


def main():
    parser = argparse.ArgumentParser(description='Phase 5A: Vector Credit Field Audit')
    parser.add_argument('--d_hidden', type=int, default=128)
    parser.add_argument('--num_classes', type=int, default=10)
    parser.add_argument('--alpha', type=float, default=1.0)
    parser.add_argument('--depths', type=int, nargs='+', default=[4])
    parser.add_argument('--M_values', type=int, nargs='+', default=[4, 8])
    parser.add_argument('--epochs', type=int, default=80)
    parser.add_argument('--steps_per_epoch', type=int, default=50)
    parser.add_argument('--batch_size', type=int, default=256)
    parser.add_argument('--lr', type=float, default=1e-3)
    parser.add_argument('--lr_fb', type=float, default=1e-3)
    parser.add_argument('--warmup_ratio', type=float, default=0.05)
    parser.add_argument('--term_grad_weight', type=float, default=1.0)
    parser.add_argument('--lam', type=float, default=0.1)
    parser.add_argument('--K', type=int, default=4)
    parser.add_argument('--sigma_bridge', type=float, default=0.05)
    parser.add_argument('--ema_momentum', type=float, default=0.995)
    parser.add_argument('--pert_eps', type=float, default=1e-3)
    parser.add_argument('--pert_beta', type=float, default=1.0)
    parser.add_argument('--seeds', type=int, nargs='+', default=[42])
    parser.add_argument('--gpu', type=int, default=2)
    parser.add_argument('--output_dir', type=str, default='results/vector_audit')
    parser.add_argument('--full_audit', action='store_true',
                        help='Run sanity checks for all seeds (default: first seed only)')
    args = parser.parse_args()
    run_experiment(args)


if __name__ == '__main__':
    main()