#!/usr/bin/env python3
"""
Quick test to verify DynamisLab package structure and imports.
"""

import sys
from pathlib import Path

# Add src to path
sys.path.insert(0, str(Path(__file__).parent / 'src'))

print("=" * 70)
print("DynamisLab Package Structure Test")
print("=" * 70)

# Test 1: Import main package
print("\n[1] Testing package import...")
try:
    import config
    import environments
    print("✓ Package imports successful")
    print(f"  - config module: {config.__file__}")
    print(f"  - environments module: {environments.__file__}")
except ImportError as e:
    print(f"✗ Import failed: {e}")
    sys.exit(1)

# Test 2: Import config functions
print("\n[2] Testing config module...")
try:
    from config import (
        load_celeris_configs,
        get_model_path,
        get_tensorboard_logdir,
        get_output_path,
        CONFIG_DIR,
    )
    print("✓ Config functions imported")
    print(f"  - CONFIG_DIR: {CONFIG_DIR}")
except ImportError as e:
    print(f"✗ Config import failed: {e}")
    sys.exit(1)

# Test 3: Import environment
print("\n[3] Testing environments module...")
try:
    import environments
    print("✓ Environments module imported")
    print(f"  - Module: {environments.__file__}")
    
    # Try to import CFDFlowControlEnv (might fail if dependencies not installed)
    try:
        from environments import CFDFlowControlEnv
        print("✓ CFDFlowControlEnv imported")
        print(f"  - Class: {CFDFlowControlEnv}")
    except ImportError as e:
        print(f"⚠ CFDFlowControlEnv import failed (missing dependencies): {e}")
        print("  This is OK if gymnasium/torch not installed yet")
except ImportError as e:
    print(f"✗ Environments module import failed: {e}")
    sys.exit(1)

# Test 4: Check version
print("\n[4] Checking package info...")
try:
    # Try importing as installed package (might not work in src layout without install)
    # This is just to show what will work after pip install -e .
    print("  Note: Run 'pip install -e .' to enable 'import dynamis' style")
    print("  Current import style: 'from config import ...' (src in PYTHONPATH)")
except Exception as e:
    print(f"  {e}")

# Test 5: Config paths
print("\n[5] Testing path helpers...")
try:
    model_path = get_model_path("test_model")
    tb_logdir = get_tensorboard_logdir("test_run")
    output_path = get_output_path("test.pkl")
    
    print(f"✓ Path helpers working")
    print(f"  - Model path: {model_path}")
    print(f"  - TensorBoard logdir: {tb_logdir}")
    print(f"  - Output path: {output_path}")
except Exception as e:
    print(f"✗ Path helpers failed: {e}")
    sys.exit(1)

print("\n" + "=" * 70)
print("✅ All structure tests passed!")
print("=" * 70)
print("\nNext steps:")
print("  1. pip install -e .                  # Install in development mode")
print("  2. python scripts/train_ppo.py --help  # Test training script")
print("=" * 70)