#!/usr/bin/env python3
"""
Test script for WrinkleBrane dataset creation (small version)
"""

import os
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))

import numpy as np
from wrinklebrane_dataset_builder import WrinkleBraneDatasetBuilder

def test_small_dataset():
    print("🧪 Testing WrinkleBrane Dataset Builder...")
    
    # Create builder with your token
    hf_token = "os.environ.get('HF_TOKEN', 'your-token-here')"
    repo_id = "WrinkleBrane"
    
    builder = WrinkleBraneDatasetBuilder(hf_token, repo_id)
    
    # Test visual memory generation
    print("👁️ Testing visual memory pairs...")
    visual_samples = builder.generate_visual_memory_pairs(5, H=32, W=32)
    print(f"✅ Generated {len(visual_samples)} visual memory samples")
    
    # Test synthetic maps
    print("🗺️ Testing synthetic maps...")
    map_samples = builder.generate_synthetic_maps(3, H=32, W=32)
    print(f"✅ Generated {len(map_samples)} synthetic map samples")
    
    # Test interference studies
    print("⚡ Testing interference studies...")
    interference_samples = builder.generate_interference_studies(4, H=16, W=16)
    print(f"✅ Generated {len(interference_samples)} interference samples")
    
    # Test orthogonality benchmarks
    print("📐 Testing orthogonality benchmarks...")
    orthogonal_samples = builder.generate_orthogonality_benchmarks(2, L=16, K=16)
    print(f"✅ Generated {len(orthogonal_samples)} orthogonality samples")
    
    # Test persistence traces
    print("⏰ Testing persistence traces...")
    persistence_samples = builder.generate_persistence_traces(3, H=16, W=16)
    print(f"✅ Generated {len(persistence_samples)} persistence samples")
    
    # Show sample structure
    print("\n📊 Visual Memory Sample Structure:")
    if visual_samples:
        sample = visual_samples[0]
        for key, value in sample.items():
            if key in ["key_pattern", "value_pattern"]:
                arr = np.array(value)
                print(f"  {key}: shape {arr.shape}, range [{arr.min():.3f}, {arr.max():.3f}]")
            else:
                print(f"  {key}: {value}")
    
    print("\n📊 Interference Sample Structure:")
    if interference_samples:
        sample = interference_samples[0]
        for key, value in sample.items():
            if key in ["key_pattern", "value_pattern"]:
                arr = np.array(value)
                print(f"  {key}: shape {arr.shape}, range [{arr.min():.3f}, {arr.max():.3f}]")
            elif key not in ["key_pattern", "value_pattern"] and value is not None:
                print(f"  {key}: {value}")
    
    print("\n🎉 All tests passed! WrinkleBrane dataset builder is working correctly.")
    return True

if __name__ == "__main__":
    test_small_dataset()