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

# Detect operating system and load appropriate library
def detect_library():
    script_dir = os.path.dirname(os.path.abspath(__file__))
    system = platform.system()
    
    # Try to find the library in the package structure
    # First try the installed package location
    if system == 'Windows':
        lib_name = 'grnexus_core.dll'
        platform_dir = 'Windows'
    elif system == 'Darwin':  # macOS
        lib_name = 'grnexus_core.dylib'
        platform_dir = 'Mac'
    elif system == 'Linux':
        lib_name = 'grnexus_core.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)
    ]

# Define function signatures
_lib.matrix_multiply.argtypes = [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData)]
_lib.matrix_multiply.restype = ctypes.c_int

_lib.matrix_add.argtypes = [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData)]
_lib.matrix_add.restype = ctypes.c_int

_lib.matrix_transpose.argtypes = [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData)]
_lib.matrix_transpose.restype = ctypes.c_int

_lib.matrix_hadamard.argtypes = [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData)]
_lib.matrix_hadamard.restype = ctypes.c_int

_lib.matrix_scale.argtypes = [ctypes.POINTER(GRNexusData), ctypes.c_double, ctypes.POINTER(GRNexusData)]
_lib.matrix_scale.restype = ctypes.c_int

_lib.dense_forward.argtypes = [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData)]
_lib.dense_forward.restype = ctypes.c_int

_lib.dense_backward.argtypes = [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_double]
_lib.dense_backward.restype = ctypes.c_int

_lib.cross_entropy_loss.argtypes = [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData)]
_lib.cross_entropy_loss.restype = ctypes.c_double

_lib.cross_entropy_gradient.argtypes = [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData)]
_lib.cross_entropy_gradient.restype = ctypes.c_int

_lib.mse_loss.argtypes = [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData)]
_lib.mse_loss.restype = ctypes.c_double

_lib.mse_gradient.argtypes = [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData)]
_lib.mse_gradient.restype = ctypes.c_int

_lib.sgd_update.argtypes = [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_double]
_lib.sgd_update.restype = ctypes.c_int

_lib.adam_update.argtypes = [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_int]
_lib.adam_update.restype = ctypes.c_int

_lib.rmsprop_update.argtypes = [ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.POINTER(GRNexusData), ctypes.c_double, ctypes.c_double, ctypes.c_double]
_lib.rmsprop_update.restype = ctypes.c_int

_lib.xavier_init.argtypes = [ctypes.POINTER(GRNexusData), ctypes.c_size_t, ctypes.c_size_t]
_lib.xavier_init.restype = ctypes.c_int

_lib.he_init.argtypes = [ctypes.POINTER(GRNexusData), ctypes.c_size_t]
_lib.he_init.restype = ctypes.c_int

_lib.clip_gradients.argtypes = [ctypes.POINTER(GRNexusData), ctypes.c_double]
_lib.clip_gradients.restype = ctypes.c_int

# Helper functions
def create_matrix_data(matrix: List[List[float]]) -> Tuple[GRNexusData, ctypes.Array]:
    """Create a GRNexusData structure from a 2D Python list."""
    rows = len(matrix)
    cols = len(matrix[0])
    size = rows * cols
    
    # Flatten matrix
    flat_data = [val for row in matrix for val in row]
    buffer = (ctypes.c_double * size)(*flat_data)
    
    data = GRNexusData()
    data.data = ctypes.cast(buffer, ctypes.POINTER(ctypes.c_double))
    data.type = 3  # MATRIX
    data.size = size
    data.stride = cols
    data.dims[0] = rows
    data.dims[1] = cols
    data.dims[2] = 0
    
    return data, buffer

def create_vector_data(vector: List[float]) -> Tuple[GRNexusData, ctypes.Array]:
    """Create a GRNexusData structure from a 1D Python list."""
    size = len(vector)
    buffer = (ctypes.c_double * size)(*vector)
    
    data = GRNexusData()
    data.data = ctypes.cast(buffer, ctypes.POINTER(ctypes.c_double))
    data.type = 2  # VECTOR
    data.size = size
    data.stride = 1
    data.dims[0] = size
    data.dims[1] = 0
    data.dims[2] = 0
    
    return data, buffer

def create_output_matrix_data(rows: int, cols: int) -> Tuple[GRNexusData, ctypes.Array]:
    """Create an empty GRNexusData structure for output."""
    size = rows * cols
    buffer = (ctypes.c_double * size)()
    
    data = GRNexusData()
    data.data = ctypes.cast(buffer, ctypes.POINTER(ctypes.c_double))
    data.type = 3  # MATRIX
    data.size = size
    data.stride = cols
    data.dims[0] = rows
    data.dims[1] = cols
    data.dims[2] = 0
    
    return data, buffer

def read_matrix_data(buffer: ctypes.Array, rows: int, cols: int) -> List[List[float]]:
    """Read data from a buffer into a 2D Python list."""
    matrix = []
    for i in range(rows):
        row = []
        for j in range(cols):
            row.append(buffer[i * cols + j])
        matrix.append(row)
    return matrix

def read_vector_data(buffer: ctypes.Array, size: int) -> List[float]:
    """Read data from a buffer into a 1D Python list."""
    return list(buffer[:size])

# Export library and helper functions
__all__ = [
    '_lib', 'GRNexusData',
    'create_matrix_data', 'create_vector_data', 'create_output_matrix_data',
    'read_matrix_data', 'read_vector_data'
]
