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

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from dno.core.organism import OrganismManager
from dno.core.network import DynamicNetwork
from dno.core.base import BaseEvolvableModule
from dno.config import DnoConfig
from dno.core.survival import SurvivalEngine

def test_phase4_survival():
    print("--- Starting Phase 4 Verification ---")
    
    manager = OrganismManager()
    network = DynamicNetwork(manager)
    
    # Initialize Survival Engine with fast decay for testing
    config = DnoConfig(
        utility_threshold=0.5, 
        pruning_decay_constant=0.5, 
        enable_negative_cleanup=True
    )
    survival = SurvivalEngine(manager, config, debug=True)
    
    # 1. Setup Graph: Input(A) -> Hidden(B) -> Output(C)
    #               -> DeadBranch(D)
    
    layer_a = BaseEvolvableModule(nn.Linear(10, 10)) # Input
    layer_a.is_immortal = True # Protected
    
    layer_b = BaseEvolvableModule(nn.Linear(10, 10)) # Useful Hidden
    
    layer_c = BaseEvolvableModule(nn.Linear(10, 10)) # Output
    layer_c.is_immortal = True
    
    layer_d = BaseEvolvableModule(nn.Linear(10, 10)) # Useless Branch
    
    network.add_layer(layer_a)
    network.add_layer(layer_b, parents=[layer_a.dynamic_id])
    network.add_layer(layer_c, parents=[layer_b.dynamic_id])
    network.add_layer(layer_d, parents=[layer_a.dynamic_id])
    
    # 2. Test Information Gain
    print("\n[Test] Information Gain")
    
    # Layer B: Useful (High Entropy change)
    x = torch.randn(1, 10)
    out_b = torch.randn(1, 10) * 10 # Big change
    survival.calculate_information_gain(layer_b, x, out_b)
    print(f"Layer B Utility: {layer_b.utility_score:.4f}")
    if layer_b.utility_score > 0.5:
        print("[PASS] High utility detected.")
    else:
        print("[FAIL] Utility too low.")

    # Layer D: Useless (Identity)
    # Simulate identity behavior
    out_d = x.clone()
    survival.calculate_information_gain(layer_d, x, out_d)
    print(f"Layer D Utility: {layer_d.utility_score:.4f}")
    
    # Force D to zero utility for test speed
    layer_d.utility_score = 0.0
    survival._update_status(layer_d)
    
    if layer_d.is_dying:
        print("[PASS] Layer D marked as dying.")
    else:
        print("[FAIL] Layer D not dying.")

    # 3. Test Selective Decay & Grace Period
    print("\n[Test] Selective Decay")
    
    # D is dying, grace period 100
    optimizer = optim.AdamW(network.parameters())
    
    # Fast forward grace period
    layer_d.grace_period_steps = 1
    survival.apply_selective_decay(optimizer) # Step 1: Grace becomes 0
    
    # Check weights (should not decay yet)
    w_before = layer_d.child_module.weight.clone()
    
    survival.apply_selective_decay(optimizer) # Step 2: Grace 0 -> Decay happens!
    w_after = layer_d.child_module.weight
    
    # Manual decay check
    # w_after should be w_before * (1 - 0.5) roughly
    if w_after.norm() < w_before.norm():
        print("[PASS] Weight decay applied to dying layer.")
    else:
        print("[FAIL] Weights did not decay.")
        
    # 4. Test Garbage Collection & Safety
    print("\n[Test] Garbage Collection")
    
    # Force D weights to near zero to allow pruning
    with torch.no_grad():
        layer_d.child_module.weight.fill_(0.0)
        layer_d.child_module.bias.fill_(0.0)
        
    # Try to collect D (Should succeed, it's a dead leaf)
    removed = survival.garbage_collect(network)
    
    if removed == 1 and layer_d.dynamic_id not in manager.registry:
        print("[PASS] Dead branch D collected.")
    else:
        print(f"[FAIL] GC failed for D. Removed: {removed}")
        
    # Try to collect B (Should FAIL, it's a bridge)
    # Even if we force B to die
    layer_b.is_immortal = False
    layer_b.utility_score = 0.0
    layer_b.is_dying = True
    layer_b.grace_period_steps = 0
    with torch.no_grad():
        layer_b.child_module.weight.fill_(0.0)
        
    removed_b = survival.garbage_collect(network)
    if removed_b == 0 and layer_b.dynamic_id in manager.registry:
        print("[PASS] Critical bridge B protected by Redundancy Check.")
    else:
        print(f"[FAIL] Critical bridge B removed! Graph broken.")

    print("\n--- Phase 4 Verification Complete ---")

if __name__ == "__main__":
    test_phase4_survival()
