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.base import BaseEvolvableModule
from dno.core.network import DynamicNetwork
from dno.config import DnoConfig
from dno.core.growth import GrowthEngine

def test_phase3_growth():
    print("--- Starting Phase 3 Verification ---")
    
    manager = OrganismManager()
    
    # Create config for testing
    config = DnoConfig(
        max_param_count=1000, 
        stabilization_steps=5,
        entropy_threshold=0.8
    )
    network = DynamicNetwork(manager, config)
    engine = GrowthEngine(network, config)
    
    # 1. Test Growth Trigger (Wisdom Threshold)
    print("\n[Test] Wisdom Threshold")
    # Low Entropy -> No Growth
    hist_calm = [0.1, 0.09, 0.08]
    if engine.check_growth_trigger(hist_calm, manager.step_counter + 200): # Ensure cooldown pass
        print("[FAIL] Growth triggered on calm state.")
    else:
        print("[PASS] Calm state ignored.")
        
    # High Entropy but decreasing (Learning) -> No Growth
    hist_learning = [0.9, 0.8, 0.7] # Dropping fast
    if engine.check_growth_trigger(hist_learning, manager.step_counter + 200):
        print("[FAIL] Growth triggered while learning.")
    else:
        print("[PASS] Learning state ignored.")
        
    # High Entropy and Stuck (Confusion) -> Growth!
    hist_stuck = [0.9, 0.89, 0.91] 
    # Force cooldown pass
    engine.last_growth_step = -200
    if engine.check_growth_trigger(hist_stuck, manager.step_counter):
        print("[PASS] Growth triggered on confusion.")
    else:
        print("[FAIL] Confusion ignored.")

    # 2. Test Smart Cloning & Optimizer Bridge
    print("\n[Test] Smart Cloning & Check Resources")
    
    # Create Parent
    parent = BaseEvolvableModule(nn.Linear(10, 10))
    parent.dynamic_id = "parent_1"
    manager.register_module(parent)
    
    # Initialize Optimizer
    optimizer = optim.AdamW(parent.parameters(), lr=0.01)
    
    # Fake some optimizer state (variance)
    loss = parent(torch.randn(1, 10)).sum()
    loss.backward()
    optimizer.step()
    
    # Inject high variance into weight
    for group in optimizer.param_groups:
        for p in group['params']:
            if p.shape == parent.child_module.weight.shape:
                state = optimizer.state[p]
                state['exp_avg_sq'] = torch.ones_like(p) * 100.0 # Huge variance
    
    # Clone
    clone = engine.smart_clone("parent_1", optimizer)
    
    # Verify Noise
    # Unwrap IdentityWrapper if present (v0.1.5 Update)
    from dno.core.blueprints import IdentityWrapper
    child = clone
    while isinstance(child, (IdentityWrapper, BaseEvolvableModule)):
        if hasattr(child, 'child_module'):
            if isinstance(child.child_module, nn.Linear):
                 real_linear = child.child_module
                 break
            child = child.child_module
        else:
            break
            
    # Also find parent linear
    parent_linear = parent.child_module
            
    diff = (real_linear.weight - parent_linear.weight).abs().mean().item()
    print(f"Clone Mutation Magnitude: {diff:.4f}")
    if diff > 0.1: 
        print("[PASS] Smart noise applied.")
    else:
        print(f"[FAIL] Noise too low: {diff}")
        
    # 3. Test Cold Start & Stabilization
    print("\n[Test] Cold Start & Stabilization")
    base_lr = 0.01
    engine.register_clone_to_optimizer(clone, optimizer, base_lr)
    
    # Check if new group has high LR
    new_group_lr = optimizer.param_groups[-1]['lr']
    print(f"Initial LR: {new_group_lr} (Base: {base_lr})")
    if new_group_lr > base_lr * 5:
        print("[PASS] Cold start High LR set.")
    else:
        print("[FAIL] Cold start LR not high enough.")
        
    # Simulate Steps
    manager.step_counter = 0
    engine.stabilize_optimizer(optimizer) # Step 0
    
    manager.step_counter = 3
    engine.stabilize_optimizer(optimizer) # Step 3 (Midway)
    
    mid_lr = optimizer.param_groups[-1]['lr']
    print(f"Midway LR: {mid_lr}")
    
    manager.step_counter = 6
    engine.stabilize_optimizer(optimizer) # Step 6 (Done)
    final_lr = optimizer.param_groups[-1]['lr']
    print(f"Final LR: {final_lr}")
    
    if abs(final_lr - base_lr) < 1e-6:
        print("[PASS] Stabilization complete. LR reset.")
    else:
        print(f"[FAIL] LR did not return to base. Final: {final_lr}")
        
    # 4. Resource Limit
    print("\n[Test] Resource Limit")
    engine.config.max_param_count = 10 # Tiny limit
    current = engine.calculate_current_params()
    if engine.check_resources(current):
         print(f"[FAIL] Allowed growth despite limit (Current: {current} > Max: 10)")
    else:
         print(f"[PASS] Growth blocked by max_params.")
         
    print("\n--- Phase 3 Verification Complete ---")

if __name__ == "__main__":
    test_phase3_growth()
