"""
Frozen-random-blocks baseline for ViT-Mini: train BP and DFA where the 4
transformer blocks are randomly initialized and FROZEN (no parameter updates).
Only patch_embed + cls_token + pos_embed + out_ln + out_head are trainable.

This is the codex-round-6 control for the "DFA actually trains the transformer
blocks" claim. If frozen-blocks DFA gets ≈ 24% (matching the trainable-blocks
4-block ViT-Mini DFA acc), then the blocks are passengers — DFA's "24%" is
coming from patch_embed + head learning routed via untrained block mixing.
If frozen-blocks DFA stays much lower than 24%, then the trainable blocks
are doing learned work.

Usage:
    CUDA_VISIBLE_DEVICES=2 python experiments/vit_frozen_blocks_baseline.py
"""
import sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision
import torchvision.transforms as transforms
import numpy as np

from models.vit_mini import ViTMini


def get_loaders(batch_size=128):
    tv_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
    ])
    tv = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
    ])
    tr = torchvision.datasets.CIFAR10('./data', True, download=True, transform=tv_train)
    te = torchvision.datasets.CIFAR10('./data', False, download=True, transform=tv)
    return (
        DataLoader(tr, batch_size=batch_size, shuffle=True, num_workers=2),
        DataLoader(te, batch_size=batch_size, shuffle=False, num_workers=2),
    )


def evaluate(model, loader, dev):
    model.eval()
    n = c = 0
    with torch.no_grad():
        for x, y in loader:
            x, y = x.to(dev), y.to(dev)
            preds = model(x).argmax(-1)
            c += (preds == y).sum().item()
            n += x.size(0)
    return c / n


def freeze_blocks(model):
    for p in model.blocks.parameters():
        p.requires_grad_(False)
    model.blocks.eval()


def train_bp_frozen(train_loader, test_loader, dev, epochs=30, seed=42, lr=1e-3, wd=0.05):
    torch.manual_seed(seed); np.random.seed(seed); torch.cuda.manual_seed_all(seed)
    m = ViTMini(num_blocks=4, d_model=128, n_heads=4).to(dev)
    freeze_blocks(m)
    n_trainable = sum(p.numel() for p in m.parameters() if p.requires_grad)
    n_total = sum(p.numel() for p in m.parameters())
    print(f"BP-frozen-blocks: {n_trainable}/{n_total} params trainable", flush=True)
    opt = optim.AdamW(filter(lambda p: p.requires_grad, m.parameters()), lr=lr, weight_decay=wd)
    sch = optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
    for ep in range(1, epochs + 1):
        m.train()
        m.blocks.eval()  # keep blocks in eval mode (no dropout etc)
        for x, y in train_loader:
            x = x.to(dev); y = y.to(dev)
            loss = F.cross_entropy(m(x), y)
            opt.zero_grad(); loss.backward(); opt.step()
        sch.step()
        if ep % 5 == 0 or ep == 1 or ep == epochs:
            acc = evaluate(m, test_loader, dev)
            print(f"  BP-frozen ep {ep}: test_acc={acc:.4f}", flush=True)
    return m


def train_dfa_frozen(train_loader, test_loader, dev, epochs=30, seed=42, lr=1e-3, wd=0.05):
    """4 transformer blocks frozen at random init.
    Trainable: patch_embed, cls_token, pos_embed, out_ln, out_head.
    DFA-style: head with true CE on cls token; embed (patch+cls+pos) with random feedback."""
    torch.manual_seed(seed); np.random.seed(seed); torch.cuda.manual_seed_all(seed)
    m = ViTMini(num_blocks=4, d_model=128, n_heads=4).to(dev)
    freeze_blocks(m)
    n_trainable = sum(p.numel() for p in m.parameters() if p.requires_grad)
    n_total = sum(p.numel() for p in m.parameters())
    print(f"DFA-frozen-blocks: {n_trainable}/{n_total} params trainable", flush=True)

    d_model, C = 128, 10
    B0 = torch.randn(d_model, C, device=dev) / np.sqrt(C)
    embed_opt = optim.AdamW(
        list(m.patch_embed.parameters()) + [m.cls_token, m.pos_embed],
        lr=lr, weight_decay=wd
    )
    head_opt = optim.AdamW(
        list(m.out_head.parameters()) + list(m.out_ln.parameters()),
        lr=lr, weight_decay=wd
    )
    sch1 = optim.lr_scheduler.CosineAnnealingLR(embed_opt, T_max=epochs)
    sch2 = optim.lr_scheduler.CosineAnnealingLR(head_opt, T_max=epochs)
    for ep in range(1, epochs + 1):
        m.train()
        m.blocks.eval()
        for x, y in train_loader:
            x = x.to(dev); y = y.to(dev)
            with torch.no_grad():
                logits, hi = m(x, return_hidden=True)
                e_T = logits.softmax(-1); e_T[torch.arange(x.size(0)), y] -= 1
            hL_det = hi[-1].detach()
            # Head update via true CE on cls token
            h_cls = m.out_ln(hL_det[:, 0])
            head_opt.zero_grad()
            F.cross_entropy(m.out_head(h_cls), y).backward()
            head_opt.step()
            # Embed update via DFA feedback
            a0 = (e_T @ B0.T).detach()
            rms = (a0 ** 2).mean(-1, keepdim=True).sqrt() + 1e-6
            h0 = m.embed(x)
            a0_b = a0.unsqueeze(1).expand_as(h0)
            embed_loss = (h0 * (a0_b / rms.unsqueeze(1))).sum(-1).mean()
            embed_opt.zero_grad()
            embed_loss.backward()
            embed_opt.step()
        sch1.step(); sch2.step()
        if ep % 5 == 0 or ep == 1 or ep == epochs:
            acc = evaluate(m, test_loader, dev)
            print(f"  DFA-frozen ep {ep}: test_acc={acc:.4f}", flush=True)
    return m


def main():
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument('--seed', type=int, default=42)
    p.add_argument('--epochs', type=int, default=30)
    args = p.parse_args()

    dev = torch.device('cuda:0')
    print(f"Device: {dev}, seed={args.seed}, epochs={args.epochs}", flush=True)
    train_loader, test_loader = get_loaders(batch_size=128)

    print(f"\n=== BP frozen-blocks baseline (4 random-init transformer blocks, frozen), seed={args.seed} ===", flush=True)
    mb = train_bp_frozen(train_loader, test_loader, dev, epochs=args.epochs, seed=args.seed)
    bp_acc = evaluate(mb, test_loader, dev)
    print(f"FINAL BP-frozen-blocks acc: {bp_acc:.4f}", flush=True)

    print(f"\n=== DFA frozen-blocks baseline, seed={args.seed} ===", flush=True)
    md = train_dfa_frozen(train_loader, test_loader, dev, epochs=args.epochs, seed=args.seed)
    dfa_acc = evaluate(md, test_loader, dev)
    print(f"FINAL DFA-frozen-blocks acc: {dfa_acc:.4f}", flush=True)

    print(f"\n=== Summary ===")
    print(f"BP-frozen-blocks: {bp_acc:.4f}  (chance=0.10)")
    print(f"DFA-frozen-blocks: {dfa_acc:.4f}")
    print(f"Compare to ViT-Mini 4-block trainable (3-seed avg): BP=0.792, DFA=0.237")
    print(f"Compare to ViT-Mini 0-block (shallow baseline): BP=0.10, DFA=0.10")
    print()
    print("Interpretation:")
    print("  If DFA-frozen-blocks ≈ 0.237: blocks are passengers, DFA is just learning patch_embed+head")
    print("  If DFA-frozen-blocks << 0.237: trainable blocks ARE doing learned work")
    print("  If DFA-frozen-blocks ~ 0.10: untrained blocks add no useful mixing (less informative)")


if __name__ == '__main__':
    main()