#!/usr/bin/env python3
"""Verify all models can be imported and instantiated."""

import sys
import os

# Add each model directory to path
base_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, base_dir)

print("Testing model imports and basic functionality...")
print("="*60)

# Test 7B
print("\n[1] Testing Zenith-7B...")
try:
    from Zenith.V1_7B.configs import zenith_config as cfg_7b
    from Zenith.V1_7B.modeling_zenith import ZenithModel as Model7B, ZenithConfig as Config7B
    config = Config7B(
        use_eq_adapter=True,
        use_eq_attention_bias=True,
        use_eq_gated_ffn=True,
        use_eq_recurrence=True,
        eq_consistency_weight=0.02,
        eq_state_dim=256,
        num_layers=2,
        hidden_size=512,
        num_heads=8,
        head_dim=64,
        intermediate_size=2048
    )
    model = Model7B(config)
    print(f"    [OK] 7B model instantiated: {sum(p.numel() for p in model.parameters()):,} parameters")
except Exception as e:
    print(f"    [FAIL] 7B: {e}")

# Test 28B-p300
print("\n[2] Testing Zenith-28B-p300...")
try:
    p300_28b_dir = os.path.join(base_dir, '..', 'V1-Tenstorrent-Blackhole-p300', '28B')
    sys.path.insert(0, p300_28b_dir)
    from Zenith.V1_Tenstorrent_Blackhole_p300_28B.configs import zenith_config as cfg_28b
    from Zenith.V1_Tenstorrent_Blackhole_p300_28B.modeling_zenith import ZenithModel as Model28B, ZenithConfig as Config28B
    config = Config28B(
        use_eq_adapter=True,
        use_eq_attention_bias=True,
        use_eq_gated_ffn=True,
        use_eq_recurrence=True,
        eq_consistency_weight=0.02,
        eq_state_dim=256,
        num_layers=2,
        hidden_size=3072,
        num_heads=24,
        head_dim=128,
        intermediate_size=8192
    )
    model = Model28B(config)
    print(f"    [OK] 28B-p300 model instantiated: {sum(p.numel() for p in model.parameters()):,} parameters")
except Exception as e:
    print(f"    [FAIL] 28B-p300: {e}")

# Test 32B-p300
print("\n[3] Testing Zenith-32B-p300...")
try:
    p300_32b_dir = os.path.join(base_dir, '..', 'V1-Tenstorrent-Blackhole-p300', '32B')
    sys.path.insert(0, p300_32b_dir)
    from Zenith.V1_Tenstorrent_Blackhole_p300_32B.configs import zenith_config as cfg_32b
    from Zenith.V1_Tenstorrent_Blackhole_p300_32B.modeling_zenith import ZenithModel as Model32B, ZenithConfig as Config32B
    config = Config32B(
        use_eq_adapter=True,
        use_eq_attention_bias=True,
        use_eq_gated_ffn=True,
        use_eq_recurrence=True,
        eq_consistency_weight=0.02,
        eq_state_dim=256,
        num_layers=2,
        hidden_size=4096,
        num_heads=32,
        head_dim=128,
        intermediate_size=11008
    )
    model = Model32B(config)
    print(f"    [OK] 32B-p300 model instantiated: {sum(p.numel() for p in model.parameters()):,} parameters")
except Exception as e:
    print(f"    [FAIL] 32B-p300: {e}")

# Test 70B-p300
print("\n[4] Testing Zenith-70B-p300...")
try:
    p300_70b_dir = os.path.join(base_dir, '..', 'V1-Tenstorrent-Blackhole-p300', '70B')
    sys.path.insert(0, p300_70b_dir)
    from Zenith.V1_Tenstorrent_Blackhole_p300_70B.configs import zenith_config as cfg_70b
    from Zenith.V1_Tenstorrent_Blackhole_p300_70B.modeling_zenith import ZenithModel as Model70B, ZenithConfig as Config70B
    config = Config70B(
        use_eq_adapter=True,
        use_eq_attention_bias=True,
        use_eq_gated_ffn=True,
        use_eq_recurrence=True,
        eq_consistency_weight=0.02,
        eq_state_dim=256,
        num_layers=2,
        hidden_size=8192,
        num_heads=64,
        head_dim=128,
        intermediate_size=28672
    )
    model = Model70B(config)
    print(f"    [OK] 70B-p300 model instantiated: {sum(p.numel() for p in model.parameters()):,} parameters")
except Exception as e:
    print(f"    [FAIL] 70B-p300: {e}")

print("\n" + "="*60)
print("EQ ENGINE IMPLEMENTATION VERIFICATION COMPLETE")
print("="*60)