"""Integration test for Grid-JEPA."""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

import torch
import numpy as np
from data.arc_dataset import ARCDataset, collate_fn
from data.masking import GridMaskingStrategy
from models.grid_jepa import GridJEPA
from torch.utils.data import DataLoader

def test():
    B, G, C = 4, 10, 10
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Device: {device}")
    
    import json, shutil
    d = Path("./test_data")
    d.mkdir(exist_ok=True)
    (d / "training").mkdir(exist_ok=True)
    (d / "evaluation").mkdir(exist_ok=True)
    task = {"train": [{"input": [[0,0,0],[0,1,0],[0,0,0]], "output": [[0,0,0],[0,2,0],[0,0,0]]}],
            "test": [{"input": [[0,0,0],[0,1,0],[0,0,0]], "output": [[0,0,0],[0,2,0],[0,0,0]]}]}
    for s in ["training", "evaluation"]:
        with open(d / s / "task.json", "w") as f:
            json.dump(task, f)
    
    train_ds = ARCDataset(d, "training", G, C, augment=True)
    val_ds = ARCDataset(d, "evaluation", G, C, augment=False)
    train_loader = DataLoader(train_ds, batch_size=B, shuffle=True, collate_fn=collate_fn)
    val_loader = DataLoader(val_ds, batch_size=B, shuffle=False, collate_fn=collate_fn)
    print(f"Train: {len(train_ds)}, Val: {len(val_ds)}")
    
    model = GridJEPA(C, 128, 4, 4, 4, max_grid_size=G).to(device)
    params = list(model.patch_embed.parameters()) + list(model.context_encoder.parameters()) + list(model.predictor.parameters())
    if model.use_action_conditioning:
        params += list(model.action_embed.parameters())
    optimizer = torch.optim.AdamW(params, lr=1e-3, weight_decay=0.04)
    mask_strategy = GridMaskingStrategy(G, 2, (0.1, 0.3))
    
    for epoch in range(1, 3):
        model.train()
        epoch_loss = 0.0
        nb = 0
        for batch in train_loader:
            grid = batch["input_grids"].to(device)
            Bb = grid.shape[0]
            masks = mask_strategy.sample_masks_batch(Bb)
            ctx = masks["context_mask"].to(device)
            tgt = masks["target_mask"].to(device)
            act = torch.randint(0, 6, (Bb,), device=device)
            pos = torch.randint(0, G * G, (Bb,), device=device)
            loss, _ = model(grid, ctx, tgt, act, pos)
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            model.update_ema()
            epoch_loss += loss.item()
            nb += 1
        avg = epoch_loss / nb
        
        model.eval()
        vl = 0.0
        vb = 0
        with torch.no_grad():
            for batch in val_loader:
                grid = batch["input_grids"].to(device)
                Bb = grid.shape[0]
                masks = mask_strategy.sample_masks_batch(Bb)
                loss, _ = model(grid, masks["context_mask"].to(device), masks["target_mask"].to(device),
                                torch.randint(0, 6, (Bb,), device=device), torch.randint(0, G * G, (Bb,), device=device))
                vl += loss.item()
                vb += 1
        print(f"Epoch {epoch}: train={avg:.2f}, val={vl/vb:.2f}")
    
    shutil.rmtree(d)
    print("\nIntegration test PASSED!")

if __name__ == "__main__":
    test()