import torch
import torch.nn as nn
import os
import sys

# Add project root to path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from dno.core.base import BaseEvolvableModule
from dno.core.organism import OrganismManager
from dno.core.network import DynamicNetwork

def test_phase1_flow():
    print("--- Starting Phase 1 Verification ---")
    
    # 1. Setup Manager & Network
    manager = OrganismManager()
    network = DynamicNetwork(manager)
    print("[OK] Network initialized.")

    # 2. Create Layers (Cells)
    # Layer A: Input -> Hidden
    layer_a = BaseEvolvableModule(nn.Linear(10, 20), creation_step=1)
    # Layer B: Hidden -> Output
    layer_b = BaseEvolvableModule(nn.Linear(20, 5), creation_step=2)
    # Layer C (Branch): Hidden -> Side Output
    layer_c = BaseEvolvableModule(nn.Linear(20, 5), creation_step=2)

    # 3. Register & Connect
    # Register A as Input
    network.add_layer(layer_a, parents=None)
    # Register B and C as children of A
    network.add_layer(layer_b, parents=[layer_a.dynamic_id])
    network.add_layer(layer_c, parents=[layer_a.dynamic_id])
    
    print(f"[OK] Layers registered. Graph: A({layer_a.dynamic_id}) -> [B({layer_b.dynamic_id}), C({layer_c.dynamic_id})]")

    # 4. Forward Pass
    input_tensor = torch.randn(1, 10)
    output = network(input_tensor)
    
    print(f"[OK] Forward pass complete. Output shape: {output.shape}")
    
    # Check Energy Scores
    print(f"Layer A Energy: {layer_a.energy_score:.4f}")
    if layer_a.energy_score <= 0:
        print("[FAIL] Layer A energy score did not increase.")
        return
    else:
        print("[PASS] Energy tracking active.")

    # 5. Serialization Test
    print("--- Testing Serialization ---")
    network.save_dno("test_organism")
    
    if os.path.exists("test_organism.dno"):
        print("[PASS] Serialization files created (.dno).")
    else:
        print("[FAIL] Serialization files missing.")

    # Cleanup
    if os.path.exists("test_organism.dno"):
        os.remove("test_organism.dno")
    print("--- Phase 1 Verification Complete ---")

if __name__ == "__main__":
    test_phase1_flow()
