virtual_hardware/enhanced_cpu.py

"""
Enhanced CPU Module with GPU Integration

This module extends the original CPU implementation to include GPU communication
capabilities and enhanced threading support for the 50 cores / 100 threads configuration.
"""

import multiprocessing
import threading
import time
import queue
from typing import Dict, Any, Optional, List
from dataclasses import dataclass

# Import original CPU components
from virtual_hardware_display_system.src.cpu import Core, MultiCoreCPU, CPULogger

# Import GPU driver interface
from virtual_gpu_driver import VirtualGPUDriver, CPUGPUInterface, GPUCommandType


@dataclass
class VirtualThread:
    """Represents a virtual thread running on a CPU core."""
    thread_id: int
    core_id: int
    program_counter: int = 0
    stack_pointer: int = 255
    registers: Dict[str, int] = None
    status: str = "ready"  # ready, running, waiting, terminated
    priority: int = 1
    
    def __post_init__(self):
        if self.registers is None:
            self.registers = {"AX": 0, "BX": 0, "CX": 0, "DX": 0}


class ThreadScheduler:
    """Simple round-robin thread scheduler for virtual threads."""
    
    def __init__(self, max_threads_per_core: int = 2):
        self.max_threads_per_core = max_threads_per_core
        self.threads: Dict[int, List[VirtualThread]] = {}  # core_id -> list of threads
        self.current_thread_index: Dict[int, int] = {}  # core_id -> current thread index
        self.thread_counter = 0
        
    def create_thread(self, core_id: int, program_counter: int = 0) -> int:
        """Create a new virtual thread on the specified core."""
        if core_id not in self.threads:
            self.threads[core_id] = []
            self.current_thread_index[core_id] = 0
            
        if len(self.threads[core_id]) >= self.max_threads_per_core:
            return -1  # Core is at thread capacity
            
        thread_id = self.thread_counter
        self.thread_counter += 1
        
        thread = VirtualThread(
            thread_id=thread_id,
            core_id=core_id,
            program_counter=program_counter
        )
        
        self.threads[core_id].append(thread)
        return thread_id
        
    def get_current_thread(self, core_id: int) -> Optional[VirtualThread]:
        """Get the currently scheduled thread for a core."""
        if core_id not in self.threads or not self.threads[core_id]:
            return None
            
        threads = self.threads[core_id]
        current_index = self.current_thread_index[core_id]
        
        if current_index < len(threads):
            return threads[current_index]
        return None
        
    def schedule_next_thread(self, core_id: int) -> Optional[VirtualThread]:
        """Schedule the next thread for execution on a core."""
        if core_id not in self.threads or not self.threads[core_id]:
            return None
            
        threads = self.threads[core_id]
        if not threads:
            return None
            
        # Round-robin scheduling
        self.current_thread_index[core_id] = (self.current_thread_index[core_id] + 1) % len(threads)
        return self.get_current_thread(core_id)
        
    def terminate_thread(self, thread_id: int) -> bool:
        """Terminate a virtual thread."""
        for core_id, threads in self.threads.items():
            for i, thread in enumerate(threads):
                if thread.thread_id == thread_id:
                    thread.status = "terminated"
                    threads.pop(i)
                    # Adjust current thread index if necessary
                    if self.current_thread_index[core_id] >= len(threads):
                        self.current_thread_index[core_id] = 0
                    return True
        return False
        
    def get_thread_count(self, core_id: int) -> int:
        """Get the number of active threads on a core."""
        return len(self.threads.get(core_id, []))
        
    def get_total_thread_count(self) -> int:
        """Get the total number of active threads across all cores."""
        return sum(len(threads) for threads in self.threads.values())


class EnhancedCore(Core):
    """Enhanced CPU core with GPU integration and threading support."""
    
    def __init__(self, core_id: int, gpu_interface: Optional[CPUGPUInterface] = None):
        super().__init__(core_id)
        self.gpu_interface = gpu_interface
        self.thread_scheduler = ThreadScheduler(max_threads_per_core=2)  # 2 threads per core for 100 total
        self.gpu_command_queue = queue.Queue()
        self.gpu_results = {}
        
        # Enhanced instruction set for GPU operations
        self.gpu_instructions = {
            'GPU_CLEAR', 'GPU_RECT', 'GPU_TRANSFER', 'GPU_ALLOC', 'GPU_AI_INFER',
            'GPU_MATRIX_MUL', 'GPU_WAIT', 'GPU_STATUS'
        }
        
    def connect_gpu_interface(self, gpu_interface: CPUGPUInterface):
        """Connect the GPU interface to this core."""
        self.gpu_interface = gpu_interface
        
    def create_virtual_thread(self, program_counter: int = 0) -> int:
        """Create a new virtual thread on this core."""
        return self.thread_scheduler.create_thread(self.core_id, program_counter)
        
    def execute_with_threading(self, instruction):
        """Execute instruction with threading support."""
        current_thread = self.thread_scheduler.get_current_thread(self.core_id)
        
        if current_thread is None:
            # No threads, execute normally
            return self.execute(instruction)
            
        # Save current core state to thread
        current_thread.registers["AX"] = self.AX
        current_thread.registers["BX"] = self.BX
        current_thread.registers["CX"] = self.CX
        current_thread.registers["DX"] = self.DX
        current_thread.program_counter = self.PC
        current_thread.stack_pointer = self.SP
        
        # Execute instruction
        result = self.execute(instruction)
        
        # Restore thread state to core
        self.AX = current_thread.registers["AX"]
        self.BX = current_thread.registers["BX"]
        self.CX = current_thread.registers["CX"]
        self.DX = current_thread.registers["DX"]
        self.PC = current_thread.program_counter
        self.SP = current_thread.stack_pointer
        
        return result
        
    def execute(self, instruction):
        """Enhanced execute method with GPU instruction support."""
        op = instruction.get("op")
        
        # Handle GPU instructions
        if op in self.gpu_instructions:
            return self._execute_gpu_instruction(instruction)
        
        # Handle regular CPU instructions
        return super().execute(instruction)
        
    def _execute_gpu_instruction(self, instruction):
        """Execute GPU-specific instructions."""
        if not self.gpu_interface:
            print(f"Core {self.core_id} Error: GPU interface not connected")
            return
            
        op = instruction.get("op")
        
        try:
            if op == 'GPU_CLEAR':
                color = instruction.get('color', (0, 0, 0))
                cmd_id = self.gpu_interface.gpu_clear_screen(color, self.core_id)
                self.gpu_results[cmd_id] = "pending"
                self.AX = hash(cmd_id) & 0xFFFF  # Store command ID hash in AX
                
            elif op == 'GPU_RECT':
                x = instruction.get('x', 0)
                y = instruction.get('y', 0)
                width = instruction.get('width', 100)
                height = instruction.get('height', 100)
                color = instruction.get('color', (255, 255, 255))
                cmd_id = self.gpu_interface.gpu_draw_rect(x, y, width, height, color, self.core_id)
                self.gpu_results[cmd_id] = "pending"
                self.AX = hash(cmd_id) & 0xFFFF
                
            elif op == 'GPU_TRANSFER':
                data = instruction.get('data', b'')
                name = instruction.get('name', f'transfer_{self.core_id}')
                cmd_id = self.gpu_interface.gpu_transfer_data(data, name, self.core_id)
                self.gpu_results[cmd_id] = "pending"
                self.AX = hash(cmd_id) & 0xFFFF
                
            elif op == 'GPU_ALLOC':
                width = instruction.get('width', 1920)
                height = instruction.get('height', 1080)
                channels = instruction.get('channels', 3)
                name = instruction.get('name')
                cmd_id = self.gpu_interface.gpu_alloc_framebuffer(width, height, channels, name, self.core_id)
                self.gpu_results[cmd_id] = "pending"
                self.AX = hash(cmd_id) & 0xFFFF
                
            elif op == 'GPU_AI_INFER':
                model_data = instruction.get('model_data')
                input_data = instruction.get('input_data')
                cmd_id = self.gpu_interface.gpu_ai_inference(model_data, input_data, self.core_id)
                self.gpu_results[cmd_id] = "pending"
                self.AX = hash(cmd_id) & 0xFFFF
                
            elif op == 'GPU_MATRIX_MUL':
                matrix_a = instruction.get('matrix_a')
                matrix_b = instruction.get('matrix_b')
                cmd_id = self.gpu_interface.gpu_matrix_multiply(matrix_a, matrix_b, self.core_id)
                self.gpu_results[cmd_id] = "pending"
                self.AX = hash(cmd_id) & 0xFFFF
                
            elif op == 'GPU_WAIT':
                cmd_id_hash = instruction.get('cmd_id_hash', self.AX)
                timeout = instruction.get('timeout', 10.0)
                
                # Find command ID by hash (simplified)
                cmd_id = None
                for cid in self.gpu_results:
                    if (hash(cid) & 0xFFFF) == cmd_id_hash:
                        cmd_id = cid
                        break
                        
                if cmd_id:
                    success = self.gpu_interface.wait_for_gpu_task(cmd_id, timeout)
                    self.ZF = 1 if success else 0
                    if success:
                        self.gpu_results[cmd_id] = "completed"
                else:
                    self.ZF = 0
                    
            elif op == 'GPU_STATUS':
                cmd_id_hash = instruction.get('cmd_id_hash', self.AX)
                
                # Find command ID by hash and check status
                for cid in self.gpu_results:
                    if (hash(cid) & 0xFFFF) == cmd_id_hash:
                        status = self.gpu_interface.gpu_driver.get_command_status(cid)
                        if status == "completed":
                            self.ZF = 1
                        elif status == "error":
                            self.ZF = 0
                            self.CF = 1
                        else:
                            self.ZF = 0
                            self.CF = 0
                        break
                        
        except Exception as e:
            print(f"Core {self.core_id} GPU instruction error: {e}")
            self.CF = 1  # Set carry flag to indicate error
            
    def run_with_threading(self):
        """Enhanced run method with threading support."""
        # Create initial threads if none exist
        if self.thread_scheduler.get_total_thread_count() == 0:
            self.create_virtual_thread(0)  # Create at least one thread
            
        time_slice = 0.01  # 10ms time slice per thread
        
        while True:
            current_thread = self.thread_scheduler.get_current_thread(self.core_id)
            
            if current_thread is None:
                break  # No threads to execute
                
            if current_thread.status == "terminated":
                self.thread_scheduler.schedule_next_thread(self.core_id)
                continue
                
            # Execute instructions for current thread
            start_time = time.time()
            instruction_count = 0
            
            while (time.time() - start_time) < time_slice and instruction_count < 100:
                try:
                    instruction = self.fetch()
                    decoded_instruction = self.decode(instruction)
                    self.execute_with_threading(decoded_instruction)
                    
                    if decoded_instruction and decoded_instruction.get('op') == 'HLT':
                        current_thread.status = "terminated"
                        break
                        
                    instruction_count += 1
                    
                except Exception as e:
                    print(f"Core {self.core_id} Thread {current_thread.thread_id} error: {e}")
                    current_thread.status = "terminated"
                    break
                    
            # Schedule next thread
            self.thread_scheduler.schedule_next_thread(self.core_id)
            
            # Small delay to prevent busy waiting
            time.sleep(0.001)


class EnhancedMultiCoreCPU(MultiCoreCPU):
    """Enhanced multi-core CPU with GPU integration and threading support."""
    
    def __init__(self, num_cores: int = 50, gpu_driver: Optional[VirtualGPUDriver] = None):
        # Initialize with enhanced cores
        self.num_cores = num_cores
        self.total_cores = 50000  # Virtual GPU cores, not CPU cores
        self.cores_per_sm = self.total_cores // num_cores
        
        # Create enhanced cores
        self.cores = []
        for i in range(num_cores):
            core = EnhancedCore(i)
            if gpu_driver:
                gpu_interface = CPUGPUInterface(gpu_driver)
                core.connect_gpu_interface(gpu_interface)
            self.cores.append(core)
            
        self.shared_ram = None
        self.shared_interrupt_handler = None
        self.gpu_driver = gpu_driver
        
        # Threading statistics
        self.total_threads_created = 0
        
    def create_threads_on_all_cores(self, threads_per_core: int = 2):
        """Create virtual threads on all cores to achieve 100 total threads."""
        total_threads = 0
        for core in self.cores:
            for _ in range(threads_per_core):
                thread_id = core.create_virtual_thread()
                if thread_id != -1:
                    total_threads += 1
                    self.total_threads_created += 1
                    
        print(f"Created {total_threads} virtual threads across {self.num_cores} cores")
        return total_threads
        
    def get_threading_stats(self) -> Dict[str, Any]:
        """Get threading statistics across all cores."""
        stats = {
            "total_cores": self.num_cores,
            "total_threads_created": self.total_threads_created,
            "active_threads_per_core": {},
            "total_active_threads": 0
        }
        
        for core in self.cores:
            thread_count = core.thread_scheduler.get_total_thread_count()
            stats["active_threads_per_core"][core.core_id] = thread_count
            stats["total_active_threads"] += thread_count
            
        return stats
        
    def get_gpu_stats(self) -> Dict[str, Any]:
        """Get GPU-related statistics."""
        if self.gpu_driver:
            return self.gpu_driver.get_driver_stats()
        return {"error": "No GPU driver connected"}
        
    def __str__(self):
        threading_stats = self.get_threading_stats()
        return (f"EnhancedMultiCoreCPU with {self.num_cores} cores, "
                f"{threading_stats['total_active_threads']} active threads, "
                f"GPU {'connected' if self.gpu_driver else 'not connected'}")


if __name__ == "__main__":
    # Test the enhanced CPU with GPU integration
    print("Testing Enhanced CPU with GPU Integration...")
    
    # This would normally be connected to actual GPU components
    # For testing, we'll create a mock setup
    
    # Create enhanced CPU
    enhanced_cpu = EnhancedMultiCoreCPU(num_cores=4)  # Use 4 cores for testing
    
    # Create threads
    threads_created = enhanced_cpu.create_threads_on_all_cores(threads_per_core=2)
    print(f"Created {threads_created} threads")
    
    # Get stats
    threading_stats = enhanced_cpu.get_threading_stats()
    print(f"Threading stats: {threading_stats}")
    
    print(f"Enhanced CPU: {enhanced_cpu}")
    print("Enhanced CPU test completed")