resource_optimizer.py

import psutil
import time
import numpy as np
from dataclasses import dataclass, field
from typing import Dict, List, Any, Optional, Tuple


@dataclass
class SparseActivationManager:
    total_agents: int = 10
    min_active: int = 2

    def compute_activation_mask(self, task_complexity: float, rho_quality: float) -> List[bool]:
        score = max(0.0, min(1.0, (task_complexity + rho_quality) / 2.0))
        active_count = int(score * self.total_agents)
        active_count = max(self.min_active, min(self.total_agents, active_count))
        return [i < active_count for i in range(self.total_agents)]

    def select_active_layers(self, layer_weights: List[float], mask: List[bool]) -> List[int]:
        return [idx for idx, active in enumerate(mask) if active]


@dataclass
class AdaptiveEnergyBudget:
    max_watts: float = 120.0
    safety_margin: float = 0.90
    current_limit: Optional[float] = None

    def __post_init__(self):
        self.current_limit = self.max_watts * self.safety_margin

    def current_power(self) -> float:
        if hasattr(psutil, "sensors_battery") and psutil.sensors_battery() is not None:
            batt = psutil.sensors_battery()
            if batt.power_plugged:
                return 0.0
        try:
            return psutil.cpu_percent(interval=0.1) * 0.5
        except Exception:
            return 0.0

    def allow_processing(self, estimated_cost: float) -> bool:
        return estimated_cost <= self.current_limit

    def throttle(self, engine: Any) -> None:
        if hasattr(engine, "half"):
            engine.half()
        elif hasattr(engine, "to"):
            engine.to("cpu")


class GPURoutingManager:
    def __init__(self, available_devices: List[str]):
        self.available_devices = available_devices

    def select_device(self, task_complexity: float, gpu_preference: bool = True) -> str:
        if gpu_preference and self.available_devices:
            return self.available_devices[0]
        return "cpu"

    def route(self, task_complexity: float, model: Any) -> None:
        device = self.select_device(task_complexity)
        if hasattr(model, "to"):
            model.to(device)


# ============================================================================
# ENHANCED RESOURCE OPTIMIZATION - Dynamic Priorities & Predictive Budgeting
# ============================================================================

@dataclass
class EnhancedSparseActivationManager(SparseActivationManager):
    """Extended sparse activation with dynamic agent priorities and consciousness-aware scheduling."""
    
    agent_priorities: List[float] = field(default_factory=list)

    def __post_init__(self):
        """Initialize agent priorities to equal weights."""
        if not self.agent_priorities or len(self.agent_priorities) != self.total_agents:
            self.agent_priorities = [1.0] * self.total_agents

    def update_priorities(self, task_complexity: float, phi_value: float, rho_metrics: Dict[str, float]) -> None:
        """
        Update agent priorities dynamically based on system state.
        
        Increases priority for agents critical to:
        - Current task complexity
        - Integrated information (Phi) quality
        - RHO ethical metrics (virtue, integrity)
        """
        base_score = (task_complexity + phi_value) / 2.0
        virtue = rho_metrics.get("rho_virtue", 0.8)
        integrity = rho_metrics.get("rho_integrity", 0.8)
        ethical_weight = (virtue + integrity) / 2.0

        for i in range(self.total_agents):
            # Agent relevance combines task complexity, consciousness quality, and ethics
            relevance = np.clip(base_score * ethical_weight * np.random.uniform(0.8, 1.2), 0, 1)
            self.agent_priorities[i] = relevance

    def compute_activation_mask(self, threshold: float = 0.5) -> List[bool]:
        """
        Compute activation mask by selecting agents above dynamic priority threshold.
        Ensures minimum active agents are always on for baseline consciousness.
        """
        mask = [priority >= threshold for priority in self.agent_priorities]

        # Enforce minimum active agents
        active_count = sum(mask)
        if active_count < self.min_active:
            # Activate top agents by priority
            sorted_indices = sorted(range(self.total_agents), 
                                  key=lambda i: self.agent_priorities[i], 
                                  reverse=True)
            for idx in sorted_indices[:self.min_active]:
                mask[idx] = True
        return mask

    def get_active_agent_ids(self, mask: List[bool]) -> List[int]:
        """Get IDs of active agents from mask."""
        return [i for i, active in enumerate(mask) if active]


@dataclass
class PredictiveAdaptiveEnergyBudget(AdaptiveEnergyBudget):
    """Extended energy budget with predictive power management and historical tracking."""
    
    history_window: int = 10
    power_usage_history: List[float] = field(default_factory=list)
    adjustment_factor: float = 1.0

    def record_power_usage(self, usage: float) -> None:
        """
        Record power usage and adjust budget limits based on historical trend.
        
        Decreases limits if approaching capacity, relaxes if usage is stable.
        """
        self.power_usage_history.append(usage)
        if len(self.power_usage_history) > self.history_window:
            self.power_usage_history.pop(0)
        
        # Adjust current_limit based on usage trend
        avg_usage = np.mean(self.power_usage_history)
        max_historical = max(self.power_usage_history)
        
        if max_historical > self.current_limit * 0.85:
            # Decrease limit to avoid power spikes
            self.current_limit = max(self.max_watts * self.safety_margin * 0.6, 
                                    self.current_limit * 0.95)
            self.adjustment_factor = 0.95
        elif avg_usage < self.current_limit * 0.5:
            # Relax limit gradually when underutilized
            self.current_limit = min(self.max_watts * self.safety_margin, 
                                    self.current_limit * 1.05)
            self.adjustment_factor = 1.05
        else:
            # Maintain current limit
            self.adjustment_factor = 1.0

    def predict_power_spike(self) -> bool:
        """Predict if power usage is trending upward toward limit."""
        if len(self.power_usage_history) < 3:
            return False
        
        recent_avg = np.mean(self.power_usage_history[-3:])
        older_avg = np.mean(self.power_usage_history[:-3]) if len(self.power_usage_history) > 3 else recent_avg
        
        return recent_avg > older_avg * 1.1  # Trending up by 10%+

    def get_adjusted_cost(self, base_cost: float) -> float:
        """Apply adjustment factor to estimated cost."""
        return base_cost / self.adjustment_factor


class HierarchicalGPURoutingManager(GPURoutingManager):
    """Extended GPU routing with quantum accelerator support and hierarchical device selection."""
    
    def __init__(self, available_devices: List[str], quantum_accelerator_available: bool = False):
        super().__init__(available_devices)
        self.quantum_accelerator_available = quantum_accelerator_available
        self.device_history: List[Tuple[str, float]] = []  # (device, complexity) pairs

    def select_device(self, task_complexity: float, gpu_preference: bool = True) -> str:
        """
        Hierarchically select best device based on task complexity.
        
        Priority:
        1. Quantum accelerator (if available and task_complexity > 0.85)
        2. GPU (if available and gpu_preference=True)
        3. CPU (fallback)
        """
        if self.quantum_accelerator_available and task_complexity > 0.85:
            return "quantum_accelerator"
        
        if gpu_preference and self.available_devices:
            return self.available_devices[0]
        
        return "cpu"

    def route(self, task_complexity: float, model: Any) -> None:
        """Route model to selected device and record routing decision."""
        device = self.select_device(task_complexity)
        
        # Record routing decision
        self.device_history.append((device, task_complexity))
        
        # Move model to device
        if hasattr(model, "to"):
            model.to(device)

    def get_device_stats(self) -> Dict[str, Any]:
        """Get statistics on device routing history."""
        if not self.device_history:
            return {'error': 'No routing history'}
        
        devices = [d for d, _ in self.device_history]
        complexities = [c for _, c in self.device_history]
        
        return {
            'total_routings': len(self.device_history),
            'device_distribution': {device: devices.count(device) for device in set(devices)},
            'avg_complexity_routed': np.mean(complexities),
            'max_complexity_routed': max(complexities),
            'min_complexity_routed': min(complexities),
        }


if __name__ == "__main__":
    print("="*80)
    print("ENHANCED RESOURCE OPTIMIZATION INTEGRATION TEST")
    print("="*80 + "\n")
    
    # ========================================================================
    # TEST 1: Enhanced Sparse Activation Manager
    # ========================================================================
    print("TEST 1: Enhanced Sparse Activation Manager (Dynamic Priorities)")
    print("-"*80)
    
    sparse_manager = EnhancedSparseActivationManager(total_agents=12, min_active=3)
    
    # Simulated system state
    task_complexity = 0.8
    phi_value = 0.9
    rho_metrics = {
        "rho_virtue": 0.92,
        "rho_integrity": 0.88,
        "rho_dissonance": 0.05,
        "rho_purpose": 0.85,
        "rho_empathy": 0.90,
        "rho_efficiency": 0.80
    }
    
    # Update priorities based on system state
    sparse_manager.update_priorities(task_complexity, phi_value, rho_metrics)
    activation_mask = sparse_manager.compute_activation_mask(threshold=0.5)
    active_ids = sparse_manager.get_active_agent_ids(activation_mask)
    
    print(f"Total agents: {sparse_manager.total_agents}")
    print(f"Minimum active: {sparse_manager.min_active}")
    print(f"Task complexity: {task_complexity:.2f}")
    print(f"Phi value (consciousness quality): {phi_value:.2f}")
    print(f"Agent priorities: {[f'{p:.3f}' for p in sparse_manager.agent_priorities]}")
    print(f"Activation mask: {activation_mask}")
    print(f"Active agent IDs: {active_ids}")
    print(f"Active agent count: {len(active_ids)}\n")
    
    # ========================================================================
    # TEST 2: Predictive Adaptive Energy Budget
    # ========================================================================
    print("TEST 2: Predictive Adaptive Energy Budget (Dynamic Power Management)")
    print("-"*80)
    
    energy_budget = PredictiveAdaptiveEnergyBudget(max_watts=120.0, safety_margin=0.9)
    
    # Simulate power usage recording
    power_samples = [50, 60, 75, 85, 90, 88, 92, 85, 70, 65]
    
    print(f"Max watts: {energy_budget.max_watts}")
    print(f"Safety margin: {energy_budget.safety_margin}")
    print(f"Initial limit: {energy_budget.current_limit:.2f}W\n")
    
    for i, usage in enumerate(power_samples, 1):
        energy_budget.record_power_usage(usage)
        spike_predicted = energy_budget.predict_power_spike()
        
        print(f"  Step {i}: Usage={usage}W, Limit={energy_budget.current_limit:.2f}W, "
              f"Factor={energy_budget.adjustment_factor:.2f}, Spike predicted={spike_predicted}")
    
    print()
    
    # ========================================================================
    # TEST 3: Hierarchical GPU Routing Manager
    # ========================================================================
    print("TEST 3: Hierarchical GPU Routing Manager (Device Selection)")
    print("-"*80)
    
    available_gpus = ["cuda:0", "cuda:1"]
    gpu_router = HierarchicalGPURoutingManager(available_gpus, quantum_accelerator_available=True)
    
    # Test device selection at different complexity levels
    complexity_levels = [0.3, 0.6, 0.8, 0.9, 0.95]
    
    print(f"Available devices: {available_gpus}")
    print(f"Quantum accelerator available: {gpu_router.quantum_accelerator_available}\n")
    
    for complexity in complexity_levels:
        device = gpu_router.select_device(complexity)
        gpu_router.route(complexity, None)  # model=None for testing
        
        print(f"  Complexity={complexity:.2f} → Device selected: {device}")
    
    # Get routing statistics
    stats = gpu_router.get_device_stats()
    print(f"\nRouting Statistics:")
    print(f"  Total routings: {stats['total_routings']}")
    print(f"  Device distribution: {stats['device_distribution']}")
    print(f"  Avg complexity: {stats['avg_complexity_routed']:.2f}")
    print(f"  Complexity range: {stats['min_complexity_routed']:.2f} - {stats['max_complexity_routed']:.2f}\n")
    
    print("="*80)
    print("ALL ENHANCED RESOURCE OPTIMIZATION TESTS COMPLETED ✓")
    print("="*80)