import ctypes
import platform
import os
from typing import List, Union

# Detect operating system and load appropriate library
# Detect operating system and load appropriate library
def detect_library():
    script_dir = os.path.dirname(os.path.abspath(__file__))
    system = platform.system()
    
    # Determine library name and platform directory
    if system == 'Windows':
        lib_name = 'activations.dll'
        platform_dir = 'Windows'
    elif system == 'Darwin':  # macOS
        lib_name = 'activations.dylib'
        platform_dir = 'Mac'
    elif system == 'Linux':
        lib_name = 'activations.so'
        platform_dir = 'Linux'
    else:
        raise OSError(f"Sistema operativo no soportado: {system}")
    
    # Try multiple possible locations
    possible_paths = [
        # Installed package location (site-packages/grnexus/exports/...)
        os.path.join(script_dir, '..', 'exports', platform_dir, lib_name),
        # Development location (python/lib/../../exports/...)
        os.path.join(script_dir, '..', '..', 'exports', platform_dir, lib_name),
    ]
    
    for lib_path in possible_paths:
        normalized_path = os.path.normpath(lib_path)
        if os.path.exists(normalized_path):
            return normalized_path
    
    # If not found, raise an error with helpful message
    raise FileNotFoundError(
        f"No se pudo encontrar la biblioteca nativa {lib_name} para {system}. "
        f"Buscado en: {', '.join(possible_paths)}"
    )

# Load the shared library
_lib = ctypes.CDLL(detect_library())

# Define GRNexusData structure
class GRNexusData(ctypes.Structure):
    _fields_ = [
        ('data', ctypes.POINTER(ctypes.c_double)),
        ('type', ctypes.c_int),
        ('size', ctypes.c_size_t),
        ('stride', ctypes.c_size_t),
        ('dims', ctypes.c_size_t * 3)
    ]

# Helper functions for creating and reading GRNexusData
def create_grnexus_data(array_or_scalar: Union[List[float], float]):
    """Create a GRNexusData structure from a Python list or scalar."""
    # Convert to array if it's a number
    values = array_or_scalar if isinstance(array_or_scalar, list) else [array_or_scalar]
    size = len(values)
    
    # Create buffer with data
    buffer = (ctypes.c_double * size)()
    for i, val in enumerate(values):
        buffer[i] = val
    
    # Create GRNexusData structure
    data = GRNexusData()
    data.data = ctypes.cast(buffer, ctypes.POINTER(ctypes.c_double))
    data.type = 0 if size == 1 else 1  # SCALAR=0, ARRAY=1
    data.size = size
    data.stride = 1
    data.dims[0] = size
    data.dims[1] = 0
    data.dims[2] = 0
    
    return data, buffer

def create_output_grnexus_data(size: int):
    """Create an empty GRNexusData structure for output."""
    buffer = (ctypes.c_double * size)()
    
    data = GRNexusData()
    data.data = ctypes.cast(buffer, ctypes.POINTER(ctypes.c_double))
    data.type = 1  # ARRAY
    data.size = size
    data.stride = 1
    data.dims[0] = size
    data.dims[1] = 0
    data.dims[2] = 0
    
    return data, buffer

def read_grnexus_data(buffer, size: int) -> List[float]:
    """Read data from a GRNexusData buffer."""
    return list(buffer[:size])

# Define function signatures
FUNCTIONS = {
    'Linear': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'Step': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'Sigmoid': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'Tanh': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'ReLU': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'LeakyReLU': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_double]),
    'PReLU': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.POINTER(ctypes.c_double)]),
    'ELU': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_double]),
    'SELU': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'Softplus': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'Softsign': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'HardSigmoid': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'HardTanh': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'ThresholdedReLU': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_double]),
    'GELU': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'Swish': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'Mish': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'LiSHT': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'ReLUSquared': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'SquaredReLU': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'CELU': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_double]),
    'HardShrink': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_double]),
    'SoftShrink': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_double]),
    'TanhShrink': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'ReLU6': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'HardSwish': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'SiLU': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'GLU': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_size_t]),
    'BReLU': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool]),
    'ARelu': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_double, ctypes.c_double]),
    'FReLU': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_double]),
    'Snake': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_double]),
    'SnakeBeta': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_double, ctypes.c_double]),
    'ISRU': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_double]),
    'ISRLU': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_double]),
    'Maxout': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_size_t]),
    'Minout': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_size_t]),
}

# Attach functions to library
for func_name, (restype, argtypes) in FUNCTIONS.items():
    func = getattr(_lib, func_name)
    func.restype = restype
    func.argtypes = argtypes

# Base activation function class
class ActivationFunction:
    def __call__(self, x: Union[List[float], float], derivative: bool = False) -> List[float]:
        raise NotImplementedError("Debes implementar la función de activación")

class BaseActivation(ActivationFunction):
    def _execute_activation(self, func_name: str, input_values: Union[List[float], float], 
                           derivative: bool = False, **kwargs) -> List[float]:
        input_data, input_buffer = create_grnexus_data(input_values)
        output_size = len(input_values) if isinstance(input_values, list) else 1
        output_data, output_buffer = create_output_grnexus_data(output_size)
        
        args = [ctypes.byref(input_data), ctypes.byref(output_data), derivative]
        
        # Add function-specific arguments
        if func_name in ['LeakyReLU', 'ELU', 'ThresholdedReLU', 'CELU', 'HardShrink', 
                         'SoftShrink', 'FReLU', 'Snake', 'ISRU', 'ISRLU']:
            args.append(kwargs.get('param1', 0.01))
        elif func_name == 'PReLU':
            alpha = ctypes.c_double(kwargs.get('param1', 0.01))
            args.append(ctypes.byref(alpha))
        elif func_name == 'ARelu':
            args.append(kwargs.get('param1', 0.01))
            args.append(kwargs.get('param2', 1.0))
        elif func_name == 'SnakeBeta':
            args.append(kwargs.get('param1', 1.0))
            args.append(kwargs.get('param2', 1.0))
        elif func_name == 'GLU':
            args.append(kwargs.get('dim', 0))
        elif func_name in ['Maxout', 'Minout']:
            args.append(kwargs.get('num_pieces', 2))
        
        # Call the C function
        func = getattr(_lib, func_name)
        result = func(*args)
        
        if result == 0:
            return list(output_buffer[:output_size])
        else:
            raise RuntimeError(f"Función {func_name} falló con código: {result}")

# Activation function classes
class Linear(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('Linear', x, derivative=derivative)

class Step(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('Step', x, derivative=derivative)

class Sigmoid(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('Sigmoid', x, derivative=derivative)

class Tanh(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('Tanh', x, derivative=derivative)

class ReLU(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('ReLU', x, derivative=derivative)

class LeakyReLU(BaseActivation):
    def __init__(self, alpha=0.01):
        self.alpha = alpha
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('LeakyReLU', x, derivative=derivative, param1=self.alpha)

class PReLU(BaseActivation):
    def __init__(self, alpha=0.01):
        self.alpha = alpha
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('PReLU', x, derivative=derivative, param1=self.alpha)

class ELU(BaseActivation):
    def __init__(self, alpha=1.0):
        self.alpha = alpha
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('ELU', x, derivative=derivative, param1=self.alpha)

class SELU(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('SELU', x, derivative=derivative)

class Softplus(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('Softplus', x, derivative=derivative)

class Softsign(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('Softsign', x, derivative=derivative)

class HardSigmoid(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('HardSigmoid', x, derivative=derivative)

class HardTanh(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('HardTanh', x, derivative=derivative)

class ThresholdedReLU(BaseActivation):
    def __init__(self, theta=1.0):
        self.theta = theta
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('ThresholdedReLU', x, derivative=derivative, param1=self.theta)

class GELU(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('GELU', x, derivative=derivative)

class Swish(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('Swish', x, derivative=derivative)

class Mish(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('Mish', x, derivative=derivative)

class LiSHT(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('LiSHT', x, derivative=derivative)

class ReLUSquared(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('ReLUSquared', x, derivative=derivative)

class SquaredReLU(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('SquaredReLU', x, derivative=derivative)

class CELU(BaseActivation):
    def __init__(self, alpha=1.0):
        self.alpha = alpha
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('CELU', x, derivative=derivative, param1=self.alpha)

class HardShrink(BaseActivation):
    def __init__(self, lambda_=0.5):
        self.lambda_ = lambda_
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('HardShrink', x, derivative=derivative, param1=self.lambda_)

class SoftShrink(BaseActivation):
    def __init__(self, lambda_=0.5):
        self.lambda_ = lambda_
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('SoftShrink', x, derivative=derivative, param1=self.lambda_)

class TanhShrink(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('TanhShrink', x, derivative=derivative)

class ReLU6(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('ReLU6', x, derivative=derivative)

class HardSwish(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('HardSwish', x, derivative=derivative)

class SiLU(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('SiLU', x, derivative=derivative)

class GLU(BaseActivation):
    def __init__(self, dim=1):
        self.dim = dim
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('GLU', x, derivative=derivative, dim=self.dim)

class BReLU(BaseActivation):
    def __call__(self, x, derivative=False):
        return self._execute_activation('BReLU', x, derivative=derivative)

class ARelu(BaseActivation):
    def __init__(self, alpha=0.0, beta=1.0):
        self.alpha = alpha
        self.beta = beta
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('ARelu', x, derivative=derivative, param1=self.alpha, param2=self.beta)

class FReLU(BaseActivation):
    def __init__(self, alpha=1.0):
        self.alpha = alpha
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('FReLU', x, derivative=derivative, param1=self.alpha)

class Snake(BaseActivation):
    def __init__(self, alpha=1.0):
        self.alpha = alpha
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('Snake', x, derivative=derivative, param1=self.alpha)

class SnakeBeta(BaseActivation):
    def __init__(self, alpha=1.0, beta=1.0):
        self.alpha = alpha
        self.beta = beta
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('SnakeBeta', x, derivative=derivative, param1=self.alpha, param2=self.beta)

class ISRU(BaseActivation):
    def __init__(self, alpha=1.0):
        self.alpha = alpha
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('ISRU', x, derivative=derivative, param1=self.alpha)

class ISRLU(BaseActivation):
    def __init__(self, alpha=1.0):
        self.alpha = alpha
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('ISRLU', x, derivative=derivative, param1=self.alpha)

class Maxout(BaseActivation):
    def __init__(self, num_pieces=2):
        self.num_pieces = num_pieces
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('Maxout', x, derivative=derivative, num_pieces=self.num_pieces)

class Minout(BaseActivation):
    def __init__(self, num_pieces=2):
        self.num_pieces = num_pieces
    
    def __call__(self, x, derivative=False):
        return self._execute_activation('Minout', x, derivative=derivative, num_pieces=self.num_pieces)