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 = 'normalization.dll'
        platform_dir = 'Windows'
    elif system == 'Darwin':  # macOS
        lib_name = 'normalization.dylib'
        platform_dir = 'Mac'
    elif system == 'Linux':
        lib_name = 'normalization.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):
    """Create a GRNexusData structure from a Python list, nested list, or scalar."""
    # Flatten the array and determine dimensions
    def flatten(lst):
        """Flatten nested lists and return dimensions."""
        if not isinstance(lst, list):
            return [lst], []
        
        # Check if it's a 1D list
        if not any(isinstance(item, list) for item in lst):
            return lst, [len(lst)]
        
        # Check if it's a 2D list
        if all(isinstance(item, list) and not any(isinstance(x, list) for x in item) for item in lst):
            flat = []
            for row in lst:
                flat.extend(row)
            return flat, [len(lst), len(lst[0]) if lst else 0]
        
        # Check if it's a 3D list
        if all(isinstance(item, list) for item in lst):
            if all(isinstance(row, list) for item in lst for row in item):
                flat = []
                dim0 = len(lst)
                dim1 = len(lst[0]) if lst else 0
                dim2 = len(lst[0][0]) if lst and lst[0] else 0
                for matrix in lst:
                    for row in matrix:
                        flat.extend(row)
                return flat, [dim0, dim1, dim2]
        
        # Fallback: just flatten everything
        flat = []
        def recursive_flatten(l):
            for item in l:
                if isinstance(item, list):
                    recursive_flatten(item)
                else:
                    flat.append(item)
        recursive_flatten(lst)
        return flat, [len(flat)]
    
    values, dims = flatten(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
    
    # Set dimensions
    if len(dims) == 0:
        data.dims[0] = 1
        data.dims[1] = 0
        data.dims[2] = 0
    elif len(dims) == 1:
        data.dims[0] = dims[0]
        data.dims[1] = 0
        data.dims[2] = 0
    elif len(dims) == 2:
        data.dims[0] = dims[0]
        data.dims[1] = dims[1]
        data.dims[2] = 0
    elif len(dims) >= 3:
        data.dims[0] = dims[0]
        data.dims[1] = dims[1]
        data.dims[2] = dims[2]
    
    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 = {
    'Softmax': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_size_t]),
    'LogSoftmax': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_size_t]),
    'Sparsemax': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_size_t]),
    'TsallisSoftmax': (ctypes.c_int, [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_bool, ctypes.c_size_t, ctypes.c_double]),
}

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

# Base normalization layer class
class NormalizationLayer:
    def normalize(self, data: Union[List[float], float], derivative: bool = False) -> List[float]:
        raise NotImplementedError("Must implement normalization method")

class BaseNormalization(NormalizationLayer):
    def _execute_normalization(self, func_name: str, input_values, 
                               derivative: bool = False, **kwargs):
        # Store original shape for reshaping output
        original_shape = None
        if isinstance(input_values, list):
            # Detect shape
            if input_values and isinstance(input_values[0], list):
                if input_values[0] and isinstance(input_values[0][0], list):
                    # 3D array
                    original_shape = (len(input_values), len(input_values[0]), len(input_values[0][0]))
                else:
                    # 2D array
                    original_shape = (len(input_values), len(input_values[0]))
        
        input_data, input_buffer = create_grnexus_data(input_values)
        output_size = input_data.size
        output_data, output_buffer = create_output_grnexus_data(output_size)
        
        args = [ctypes.byref(input_data), ctypes.byref(output_data), derivative]
        
        # Add axis parameter
        axis = kwargs.get('axis', output_size)  # Use size as "global" axis by default
        
        # Handle negative axis (Python-style indexing)
        if axis < 0:
            if original_shape:
                axis = len(original_shape) + axis
            else:
                axis = 0
        
        args.append(axis)
        
        # Add function-specific arguments
        if func_name == 'TsallisSoftmax':
            args.append(kwargs.get('q', 1.0))
        
        # Call the C function
        func = getattr(_lib, func_name)
        result = func(*args)
        
        if result != 0:
            raise RuntimeError(f"Función {func_name} falló con código: {result}")
        
        # Get flat result
        flat_result = list(output_buffer[:output_size])
        
        # Reshape if needed
        if original_shape:
            if len(original_shape) == 2:
                # Reshape to 2D
                reshaped = []
                idx = 0
                for i in range(original_shape[0]):
                    row = []
                    for j in range(original_shape[1]):
                        row.append(flat_result[idx])
                        idx += 1
                    reshaped.append(row)
                return reshaped
            elif len(original_shape) == 3:
                # Reshape to 3D
                reshaped = []
                idx = 0
                for i in range(original_shape[0]):
                    matrix = []
                    for j in range(original_shape[1]):
                        row = []
                        for k in range(original_shape[2]):
                            row.append(flat_result[idx])
                            idx += 1
                        matrix.append(row)
                    reshaped.append(matrix)
                return reshaped
        
        return flat_result

# Normalization layer classes
class Softmax(BaseNormalization):
    def __init__(self, axis=-1):
        self.axis = axis
    
    def normalize(self, data, derivative=False):
        return self._execute_normalization('Softmax', data, derivative=derivative, axis=self.axis)
    
    def __call__(self, data, derivative=False):
        return self.normalize(data, derivative)

class LogSoftmax(BaseNormalization):
    def __init__(self, axis=-1):
        self.axis = axis
    
    def normalize(self, data, derivative=False):
        return self._execute_normalization('LogSoftmax', data, derivative=derivative, axis=self.axis)
    
    def __call__(self, data, derivative=False):
        return self.normalize(data, derivative)

class Sparsemax(BaseNormalization):
    def __init__(self, axis=-1):
        self.axis = axis
    
    def normalize(self, data, derivative=False):
        return self._execute_normalization('Sparsemax', data, derivative=derivative, axis=self.axis)
    
    def __call__(self, data, derivative=False):
        return self.normalize(data, derivative)

class TsallisSoftmax(BaseNormalization):
    def __init__(self, axis=-1, q=1.0):
        self.axis = axis
        self.q = q
    
    def normalize(self, data, derivative=False):
        return self._execute_normalization('TsallisSoftmax', data, derivative=derivative, axis=self.axis, q=self.q)
    
    def __call__(self, data, derivative=False):
        return self.normalize(data, derivative)
