import os
import pathlib
from typing import Dict, Any, List, Optional, Tuple
import numpy as np

from gpu_chip import GPUChip, OpticalInterconnect
from streaming_multiprocessor import StreamingMultiprocessor
from virtual_vram import VirtualVRAM
from tensor_core import TensorCore
from core import AdvancedCore
from multicore import MultiCoreSystem

class HardwareAbstractionLayer:
    def __init__(self):
        self.chips: Dict[int, GPUChip] = {}
        self.optical_links: Dict[str, OpticalInterconnect] = {}
        
    def get_chip(self, chip_id: int) -> GPUChip:
        """Get or create a GPU chip"""
        if chip_id not in self.chips:
            self.chips[chip_id] = GPUChip(chip_id)
        return self.chips[chip_id]
        
    def connect_chips(self, chip_id_a: int, chip_id_b: int, bandwidth_tbps: float = 800, latency_ns: float = 1) -> None:
        """Connect two chips with an optical link"""
        chip_a = self.get_chip(chip_id_a)
        chip_b = self.get_chip(chip_id_b)
        
        link_id = f"link_{chip_id_a}_{chip_id_b}"
        if link_id not in self.optical_links:
            self.optical_links[link_id] = OpticalInterconnect(bandwidth_tbps, latency_ns)
            
        chip_a.connect_chip(chip_b, self.optical_links[link_id])
        chip_b.connect_chip(chip_a, self.optical_links[link_id])
        
    def execute_tensor_core_matmul(self, chip_id: int, sm_id: int, A: np.ndarray, B: np.ndarray) -> Optional[np.ndarray]:
        """Execute matrix multiplication on tensor core"""
        chip = self.get_chip(chip_id)
        if sm_id >= len(chip.sms):
            return None
        return chip.sms[sm_id].tensor_core_matmul(A, B)
        
    def v2_vertex_shader(self, chip_id: int, vertex_data: List[float], shader_program: Dict[str, Any]) -> List[float]:
        """
        Run vertex shader using provided instructions.
        Supports AI/ML ops: matmul, activation, softmax, etc.
        """
        chip = self.get_chip(chip_id)
        if not chip.sms:
            return vertex_data
            
        sm = chip.sms[0]  # Use first SM for shader execution
        registers = list(vertex_data)
        
        for instr in shader_program.get('instructions', []):
            op = instr.get('opcode')
            args = instr.get('args', [])
            
            if op == 'load_vertex_data':
                continue
            elif op == 'transform_vertex':
                registers = [v * 2 for v in registers]
            elif op == 'matmul':
                A = args[0] if args else [[v] for v in registers]
                B = args[1] if len(args) > 1 else [[1.0] * len(registers)]
                result = sm.tensor_core_matmul(np.array(A), np.array(B))
                if result is not None:
                    registers = result.flatten().tolist()
            elif op == 'activation':
                registers = [max(0, v) for v in registers]  # ReLU
            elif op == 'softmax':
                import math
                exp_vals = [math.exp(v) for v in registers]
                s = sum(exp_vals)
                registers = [v / s for v in exp_vals]
                
        return registers
        
    def v2_fragment_shader(self, chip_id: int, fragment_data: Dict[str, Any], 
                          shader_program: Dict[str, Any]) -> Tuple[float, float, float, float]:
        """
        Run fragment shader using provided instructions.
        Supports AI/ML ops: matmul, activation, softmax, etc.
        """
        chip = self.get_chip(chip_id)
        if not chip.sms:
            return (1.0, 1.0, 1.0, 1.0)
            
        sm = chip.sms[0]  # Use first SM for shader execution
        color = [1.0, 1.0, 1.0, 1.0]  # Default white
        
        for instr in shader_program.get('instructions', []):
            op = instr.get('opcode')
            args = instr.get('args', [])
            
            if op == 'load_fragment_data':
                continue
            elif op == 'compute_color':
                x = fragment_data.get('x', 0)
                y = fragment_data.get('y', 0)
                color = [x % 256 / 255.0, y % 256 / 255.0, 0.5, 1.0]
            elif op == 'matmul':
                A = args[0] if args else [[c] for c in color]
                B = args[1] if len(args) > 1 else [[1.0] * len(color)]
                result = sm.tensor_core_matmul(np.array(A), np.array(B))
                if result is not None:
                    color = result.flatten().tolist()
            elif op == 'activation':
                color = [max(0, v) for v in color]  # ReLU
            elif op == 'softmax':
                import math
                exp_vals = [math.exp(v) for v in color]
                s = sum(exp_vals)
                color = [v / s for v in exp_vals]
                
        return tuple(color[:4])  # Ensure RGBA output
        
    def allocate_vram(self, chip_id: int, size_bytes: int) -> Optional[str]:
        """Allocate VRAM on specified chip"""
        chip = self.get_chip(chip_id)
        return chip.allocate_memory(size_bytes)
        
    def transfer_data(self, src_chip_id: int, dst_chip_id: int, size_bytes: int) -> float:
        """Transfer data between chips, returns transfer time"""
        src_chip = self.get_chip(src_chip_id)
        dst_chip = self.get_chip(dst_chip_id)
        return src_chip.transfer_data(dst_chip, size_bytes)