#!/usr/bin/env python3
# ==============================================================================
# 🌉 ChiasmBridge: Universal Bi-Directional Cross-Modal Verification Test
# Tests Forward (N -> M) and Inverse (M -> N) Projections across Modalities
# ==============================================================================

import sys
import os
import torch

# Add root ChiasmBridge path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from chiasm import ChiasmBridge, ChiasmConfig

def run_test():
    print("======================================================================")
    print(" 🌉 CHIASMBRIDGE: UNIVERSAL BI-DIRECTIONAL (N <-> M) TEST SUITE")
    print("======================================================================")

    test_cases = [
        ("👁️ Vision <-> Text LLM", 3584, 5120, "7B Vision Encoder <-> 24B/32B LLM"),
        ("🎙️ Speech/Audio <-> LLM", 1024, 4096, "Whisper STT <-> 8B Llama-3 LLM"),
        ("🖐️ SNN Haptics <-> LLM", 256, 3584, "256-dim LIF Spikes <-> 7B Qwen LLM"),
        ("🧠 Small LLM <-> Large LLM", 5120, 8192, "24B Model Embeddings <-> 72B Qwen LLM")
    ]

    for label, src_dim, tgt_dim, desc in test_cases:
        print(f"\n[Test] {label} ({desc}):")
        print(f"       Source Dimension (N): {src_dim} <--> Target Dimension (M): {tgt_dim}")

        # 1. Initialize Bi-Directional Config
        config = ChiasmConfig(source_dim=src_dim, target_dim=tgt_dim)
        bridge = ChiasmBridge(config)

        # 2. Forward Pass: Sensory Input (N) -> LLM Embedding (M)
        sensory_input = torch.randn(2, 16, src_dim)
        llm_embedding = bridge.project_forward(sensory_input)

        # 3. Inverse Pass: LLM Embedding (M) -> Reconstructed Sensory Space (N)
        reconstructed_sensory = bridge.project_inverse(llm_embedding)

        # 4. Verify Shapes & Energy Preservation
        print(f"       1. Sensory Input (N):         {list(sensory_input.shape)} (Mean Norm: {torch.norm(sensory_input, dim=-1).mean().item():.4f})")
        print(f"       2. Forward LLM Embedding (M): {list(llm_embedding.shape)} (Mean Norm: {torch.norm(llm_embedding, dim=-1).mean().item():.4f})")
        print(f"       3. Inverse Sensory Rec (N):   {list(reconstructed_sensory.shape)} (Mean Norm: {torch.norm(reconstructed_sensory, dim=-1).mean().item():.4f})")

        assert llm_embedding.shape == (2, 16, tgt_dim), f"Forward shape mismatch! Expected (2, 16, {tgt_dim})"
        assert reconstructed_sensory.shape == (2, 16, src_dim), f"Inverse shape mismatch! Expected (2, 16, {src_dim})"
        print(f"       ✅ SUCCESS! Bi-directional round-trip ({src_dim} -> {tgt_dim} -> {src_dim}) complete!")

    print("\n======================================================================")
    print(" 🎉 ALL BI-DIRECTIONAL (N <-> M) TESTS PASSED 100% SUCCESSFULLY!")
    print("======================================================================")

if __name__ == "__main__":
    run_test()