#!/bin/bash
#SBATCH --job-name=snn_test
#SBATCH --account=bfqt-delta-gpu
#SBATCH --partition=gpuA40x4-interactive
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4
#SBATCH --gpus-per-node=1
#SBATCH --mem=16G
#SBATCH --time=00:30:00
#SBATCH --output=runs/slurm_logs/%j_test.out
#SBATCH --error=runs/slurm_logs/%j_test.err

# ============================================================
# Quick Test: Verify snnTorch model and Lyapunov computation
# ============================================================
# Use interactive partition for fast turnaround
#
# Usage:
#   sbatch scripts/test_interactive.sbatch
# ============================================================

set -e

PROJECT_DIR="/projects/bfqt/users/yurenh2/ml-projects/snn-training"
cd "$PROJECT_DIR"

mkdir -p runs/slurm_logs

echo "============================================================"
echo "Quick Test: SNN with Lyapunov Regularization"
echo "Job ID: $SLURM_JOB_ID | Node: $SLURM_NODELIST"
echo "============================================================"
nvidia-smi --query-gpu=name,memory.total --format=csv,noheader
echo "============================================================"

# Test 1: Model creation and forward pass
echo "Test 1: Model and Lyapunov computation..."
python -c "
import torch
import sys
sys.path.insert(0, '.')
from files.models.snn_snntorch import LyapunovSNN
from files.analysis.stability_monitor import StabilityMonitor

device = torch.device('cuda')
model = LyapunovSNN(input_dim=100, hidden_dims=[128, 64], num_classes=10).to(device)
x = torch.randn(8, 50, 100, device=device)

logits, lyap, recordings = model(x, compute_lyapunov=True, record_states=True)
print(f'  Logits shape: {logits.shape}')
print(f'  Lyapunov exponent: {lyap.item():.4f}')
print(f'  Spikes shape: {recordings[\"spikes\"].shape}')
print('  PASSED')
"

# Test 2: Training loop
echo ""
echo "Test 2: Training loop with Lyapunov regularization..."
python -c "
import torch
import torch.nn as nn
import torch.optim as optim
import sys
sys.path.insert(0, '.')
from files.models.snn_snntorch import LyapunovSNN

device = torch.device('cuda')
model = LyapunovSNN(input_dim=100, hidden_dims=[128, 64, 32], num_classes=10).to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
ce_loss = nn.CrossEntropyLoss()

x = torch.randn(16, 50, 100, device=device)
y = torch.randint(0, 10, (16,), device=device)

for step in range(5):
    optimizer.zero_grad()
    logits, lyap, _ = model(x, compute_lyapunov=True, record_states=False)
    loss = ce_loss(logits, y) + 0.1 * (lyap - 0.0) ** 2
    loss.backward()
    optimizer.step()
    print(f'  Step {step+1}: loss={loss.item():.4f}, lyap={lyap.item():.4f}')

print('  PASSED')
"

# Test 3: Depth comparison (minimal)
echo ""
echo "Test 3: Depth comparison (2 epochs, depths 1,2,4)..."
python files/experiments/depth_comparison.py \
    --synthetic \
    --epochs 2 \
    --depths 1 2 4 \
    --hidden_dim 64 \
    --out_dir runs/test_output \
    --device cuda \
    --no-progress

echo ""
echo "============================================================"
echo "All tests PASSED"
echo "============================================================"