#!/usr/bin/env python3
import ctypes
import platform
import os
import numpy as np

# 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 = 'machine_learning.dll'
        platform_dir = 'Windows'
    elif system == 'Darwin':  # macOS
        lib_name = 'machine_learning.dylib'
        platform_dir = 'Mac'
    elif system == 'Linux':
        lib_name = 'machine_learning.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
try:
    _lib = ctypes.CDLL(detect_library())
except (OSError, FileNotFoundError) as e:
    print(f"Warning: Could not load machine learning library: {e}")
    print("KNN and other ML algorithms will not be available")
    _lib = None

# 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 function to create GRNexusData from numpy array
def create_data(array):
    if not isinstance(array, np.ndarray):
        array = np.array(array, dtype=np.float64)
    
    # Ensure C-contiguous array
    if not array.flags['C_CONTIGUOUS']:
        array = np.ascontiguousarray(array)
    
    data = GRNexusData()
    data.data = array.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
    data.type = 3  # GRNEXUS_MATRIX
    data.size = array.size
    data.stride = 1
    
    if array.ndim == 1:
        data.dims[0] = array.shape[0]
        data.dims[1] = 0
        data.dims[2] = 0
    elif array.ndim == 2:
        data.dims[0] = array.shape[0]
        data.dims[1] = array.shape[1]
        data.dims[2] = 0
    else:
        data.dims[0] = array.shape[0]
        data.dims[1] = array.shape[1] if len(array.shape) > 1 else 0
        data.dims[2] = array.shape[2] if len(array.shape) > 2 else 0
    
    return data, array

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

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

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

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

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

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

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

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

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


# K-Nearest Neighbors Classifier
class KNeighborsClassifier:
    def __init__(self, n_neighbors=5):
        self.n_neighbors = n_neighbors
        self.X_train = None
        self.y_train = None
    
    def fit(self, X, y):
        self.X_train = np.array(X, dtype=np.float64)
        self.y_train = np.array(y, dtype=np.float64)
        return self
    
    def save(self, filepath):
        import json
        import zlib
        
        if filepath.endswith('.nexus'):
            filepath = filepath.replace('.nexus', '.lnexus')
        elif not filepath.endswith('.lnexus'):
            filepath += '.lnexus'
        
        model_data = {
            'model_type': 'KNeighborsClassifier',
            'framework': 'GRNexus',
            'language': 'Python',
            'version': '1.0',
            'n_neighbors': self.n_neighbors,
            'x_train': self.X_train.tolist(),
            'y_train': self.y_train.tolist()
        }
        
        json_data = json.dumps(model_data)
        compressed_data = zlib.compress(json_data.encode('utf-8'))
        
        with open(filepath, 'wb') as f:
            f.write(compressed_data)
        print(f"Model saved to {filepath}")
    
    @classmethod
    def load(cls, filepath):
        import json
        import zlib
        
        if not os.path.exists(filepath):
            raise FileNotFoundError(f"File not found: {filepath}")
        if filepath.endswith('.nexus'):
            raise ValueError("Invalid file format. Expected .lnexus for ML models, got .nexus (Neural Network format)")
        
        with open(filepath, 'rb') as f:
            compressed_data = f.read()
        
        json_data = zlib.decompress(compressed_data).decode('utf-8')
        model_data = json.loads(json_data)
        
        if model_data['model_type'] != 'KNeighborsClassifier':
            raise ValueError(f"Invalid model type. Expected KNeighborsClassifier, got {model_data['model_type']}")
        
        model = cls(n_neighbors=model_data['n_neighbors'])
        model.X_train = np.array(model_data['x_train'], dtype=np.float64)
        model.y_train = np.array(model_data['y_train'], dtype=np.float64)
        return model
    
    def __repr__(self):
        trained = 'trained' if self.X_train is not None else 'not trained'
        n_samples = len(self.X_train) if self.X_train is not None else 0
        n_features = self.X_train.shape[1] if self.X_train is not None and len(self.X_train) > 0 else 0
        return f"<KNeighborsClassifier n_neighbors={self.n_neighbors}, status={trained}, samples={n_samples}, features={n_features}>"
    
    def predict(self, X):
        if self.X_train is None or self.y_train is None:
            raise ValueError("Model not trained")
        
        X = np.array(X, dtype=np.float64)
        n_samples = X.shape[0]
        
        # Keep references to prevent garbage collection
        x_train_arr = self.X_train.flatten()
        y_train_arr = self.y_train.reshape(-1, 1).flatten()
        x_test_arr = X.flatten()
        output = np.zeros(n_samples, dtype=np.float64)
        
        x_train_data, _ = create_data(self.X_train)
        y_train_data, _ = create_data(self.y_train.reshape(-1, 1))
        x_test_data, _ = create_data(X)
        output_data, _ = create_data(output)
        
        result = _lib.knn_predict(ctypes.byref(x_train_data), ctypes.byref(y_train_data),
                                   ctypes.byref(x_test_data), self.n_neighbors, ctypes.byref(output_data))
        
        if result != 0:
            raise RuntimeError("KNN prediction failed")
        
        return output


# K-Means Clustering
class KMeans:
    def __init__(self, n_clusters=3, max_iter=300):
        self.n_clusters = n_clusters
        self.max_iter = max_iter
        self.centroids = None
        self.labels_ = None
    
    def save(self, filepath):
        import json
        import zlib
        
        if filepath.endswith('.nexus'):
            filepath = filepath.replace('.nexus', '.lnexus')
        elif not filepath.endswith('.lnexus'):
            filepath += '.lnexus'
        
        model_data = {
            'model_type': 'KMeans',
            'framework': 'GRNexus',
            'language': 'Python',
            'version': '1.0',
            'n_clusters': self.n_clusters,
            'centroids': self.centroids.tolist() if self.centroids is not None else None,
            'labels': self.labels_.tolist() if self.labels_ is not None else None
        }
        
        json_data = json.dumps(model_data)
        compressed_data = zlib.compress(json_data.encode('utf-8'))
        
        with open(filepath, 'wb') as f:
            f.write(compressed_data)
        print(f"Model saved to {filepath}")
    
    @classmethod
    def load(cls, filepath):
        import json
        import zlib
        
        if not os.path.exists(filepath):
            raise FileNotFoundError(f"File not found: {filepath}")
        if filepath.endswith('.nexus'):
            raise ValueError("Invalid file format. Expected .lnexus for ML models, got .nexus (Neural Network format)")
        
        with open(filepath, 'rb') as f:
            compressed_data = f.read()
        
        json_data = zlib.decompress(compressed_data).decode('utf-8')
        model_data = json.loads(json_data)
        
        if model_data['model_type'] != 'KMeans':
            raise ValueError(f"Invalid model type. Expected KMeans, got {model_data['model_type']}")
        
        model = cls(n_clusters=model_data['n_clusters'])
        if model_data['centroids'] is not None:
            model.centroids = np.array(model_data['centroids'], dtype=np.float64)
        if model_data['labels'] is not None:
            model.labels_ = np.array(model_data['labels'], dtype=np.int32)
        return model
    
    def __repr__(self):
        trained = 'trained' if self.centroids is not None else 'not trained'
        n_features = self.centroids.shape[1] if self.centroids is not None and len(self.centroids) > 0 else 0
        return f"<KMeans n_clusters={self.n_clusters}, status={trained}, features={n_features}>"
    
    def fit(self, X):
        X = np.array(X, dtype=np.float64)
        n_samples, n_features = X.shape
        
        # Keep arrays in scope
        centroids = np.zeros((self.n_clusters, n_features), dtype=np.float64)
        labels = np.zeros(n_samples, dtype=np.float64)
        
        x_data, x_arr = create_data(X)
        centroids_data, centroids_arr = create_data(centroids)
        labels_data, labels_arr = create_data(labels)
        
        result = _lib.kmeans_fit(ctypes.byref(x_data), self.n_clusters, self.max_iter,
                                  ctypes.byref(centroids_data), ctypes.byref(labels_data))
        
        if result != 0:
            raise RuntimeError("KMeans fit failed")
        
        self.centroids = centroids
        self.labels_ = labels.astype(int)
        
        return self
    
    def predict(self, X):
        if self.centroids is None:
            raise ValueError("Model not trained")
        
        X = np.array(X, dtype=np.float64)
        n_samples = X.shape[0]
        
        labels = np.zeros(n_samples, dtype=np.float64)
        
        x_data, x_arr = create_data(X)
        centroids_data, centroids_arr = create_data(self.centroids)
        labels_data, labels_arr = create_data(labels)
        
        result = _lib.kmeans_predict(ctypes.byref(x_data), ctypes.byref(centroids_data), 
                                       ctypes.byref(labels_data))
        
        if result != 0:
            raise RuntimeError("KMeans prediction failed")
        
        return labels.astype(int)


# Linear Regression
class LinearRegression:
    def __init__(self):
        self.weights = None
    
    def save(self, filepath):
        import json
        import zlib
        
        if filepath.endswith('.nexus'):
            filepath = filepath.replace('.nexus', '.lnexus')
        elif not filepath.endswith('.lnexus'):
            filepath += '.lnexus'
        
        model_data = {
            'model_type': 'LinearRegression',
            'framework': 'GRNexus',
            'language': 'Python',
            'version': '1.0',
            'weights': self.weights.tolist() if self.weights is not None else None
        }
        
        json_data = json.dumps(model_data)
        compressed_data = zlib.compress(json_data.encode('utf-8'))
        
        with open(filepath, 'wb') as f:
            f.write(compressed_data)
        print(f"Model saved to {filepath}")
    
    @classmethod
    def load(cls, filepath):
        import json
        import zlib
        
        if not os.path.exists(filepath):
            raise FileNotFoundError(f"File not found: {filepath}")
        if filepath.endswith('.nexus'):
            raise ValueError("Invalid file format. Expected .lnexus for ML models, got .nexus (Neural Network format)")
        
        with open(filepath, 'rb') as f:
            compressed_data = f.read()
        
        json_data = zlib.decompress(compressed_data).decode('utf-8')
        model_data = json.loads(json_data)
        
        if model_data['model_type'] != 'LinearRegression':
            raise ValueError(f"Invalid model type. Expected LinearRegression, got {model_data['model_type']}")
        
        model = cls()
        if model_data['weights'] is not None:
            model.weights = np.array(model_data['weights'], dtype=np.float64)
        return model
    
    def __repr__(self):
        trained = 'trained' if self.weights is not None else 'not trained'
        n_features = len(self.weights) - 1 if self.weights is not None else 0
        return f"<LinearRegression status={trained}, features={n_features}, weights={len(self.weights) if self.weights is not None else 0}>"
    
    def fit(self, X, y):
        X = np.array(X, dtype=np.float64)
        y = np.array(y, dtype=np.float64)
        n_features = X.shape[1]
        
        weights = np.zeros(n_features + 1, dtype=np.float64)
        y_reshaped = y.reshape(-1, 1)
        
        x_data, x_arr = create_data(X)
        y_data, y_arr = create_data(y_reshaped)
        weights_data, weights_arr = create_data(weights)
        
        result = _lib.linear_regression_fit(ctypes.byref(x_data), ctypes.byref(y_data), 
                                             ctypes.byref(weights_data))
        
        if result != 0:
            raise RuntimeError("Linear regression fit failed")
        
        self.weights = weights
        return self
    
    def predict(self, X):
        if self.weights is None:
            raise ValueError("Model not trained")
        
        X = np.array(X, dtype=np.float64)
        n_samples = X.shape[0]
        
        output = np.zeros(n_samples, dtype=np.float64)
        weights_reshaped = self.weights.reshape(-1, 1)
        
        x_data, x_arr = create_data(X)
        weights_data, weights_arr = create_data(weights_reshaped)
        output_data, output_arr = create_data(output)
        
        result = _lib.linear_regression_predict(ctypes.byref(x_data), ctypes.byref(weights_data), 
                                                  ctypes.byref(output_data))
        
        if result != 0:
            raise RuntimeError("Linear regression prediction failed")
        
        return output


# Logistic Regression
class LogisticRegression:
    def __init__(self, learning_rate=0.01, max_iter=1000):
        self.learning_rate = learning_rate
        self.max_iter = max_iter
        self.weights = None
    
    def save(self, filepath):
        import json
        import zlib
        
        if filepath.endswith('.nexus'):
            filepath = filepath.replace('.nexus', '.lnexus')
        elif not filepath.endswith('.lnexus'):
            filepath += '.lnexus'
        
        model_data = {
            'model_type': 'LogisticRegression',
            'framework': 'GRNexus',
            'language': 'Python',
            'version': '1.0',
            'learning_rate': self.learning_rate,
            'max_iter': self.max_iter,
            'weights': self.weights.tolist() if self.weights is not None else None
        }
        
        json_data = json.dumps(model_data)
        compressed_data = zlib.compress(json_data.encode('utf-8'))
        
        with open(filepath, 'wb') as f:
            f.write(compressed_data)
        print(f"Model saved to {filepath}")
    
    @classmethod
    def load(cls, filepath):
        import json
        import zlib
        
        if not os.path.exists(filepath):
            raise FileNotFoundError(f"File not found: {filepath}")
        if filepath.endswith('.nexus'):
            raise ValueError("Invalid file format. Expected .lnexus for ML models, got .nexus (Neural Network format)")
        
        with open(filepath, 'rb') as f:
            compressed_data = f.read()
        
        json_data = zlib.decompress(compressed_data).decode('utf-8')
        model_data = json.loads(json_data)
        
        if model_data['model_type'] != 'LogisticRegression':
            raise ValueError(f"Invalid model type. Expected LogisticRegression, got {model_data['model_type']}")
        
        model = cls(learning_rate=model_data['learning_rate'], max_iter=model_data['max_iter'])
        if model_data['weights'] is not None:
            model.weights = np.array(model_data['weights'], dtype=np.float64)
        return model
    
    def __repr__(self):
        trained = 'trained' if self.weights is not None else 'not trained'
        n_features = len(self.weights) - 1 if self.weights is not None else 0
        return f"<LogisticRegression lr={self.learning_rate}, max_iter={self.max_iter}, status={trained}, features={n_features}>"
    
    def fit(self, X, y):
        X = np.array(X, dtype=np.float64)
        y = np.array(y, dtype=np.float64)
        n_features = X.shape[1]
        
        weights = np.zeros(n_features + 1, dtype=np.float64)
        y_reshaped = y.reshape(-1, 1)
        
        x_data, x_arr = create_data(X)
        y_data, y_arr = create_data(y_reshaped)
        weights_data, weights_arr = create_data(weights)
        
        result = _lib.logistic_regression_fit(ctypes.byref(x_data), ctypes.byref(y_data), 
                                               ctypes.byref(weights_data), self.learning_rate, self.max_iter)
        
        if result != 0:
            raise RuntimeError("Logistic regression fit failed")
        
        self.weights = weights
        return self
    
    def predict_proba(self, X):
        if self.weights is None:
            raise ValueError("Model not trained")
        
        X = np.array(X, dtype=np.float64)
        n_samples = X.shape[0]
        
        output = np.zeros(n_samples, dtype=np.float64)
        weights_reshaped = self.weights.reshape(-1, 1)
        
        x_data, x_arr = create_data(X)
        weights_data, weights_arr = create_data(weights_reshaped)
        output_data, output_arr = create_data(output)
        
        result = _lib.logistic_regression_predict(ctypes.byref(x_data), ctypes.byref(weights_data), 
                                                    ctypes.byref(output_data))
        
        if result != 0:
            raise RuntimeError("Logistic regression prediction failed")
        
        return output
    
    def predict(self, X):
        proba = self.predict_proba(X)
        return (proba >= 0.5).astype(int)


# Gaussian Naive Bayes
class GaussianNB:
    def __init__(self):
        self.means = None
        self.variances = None
        self.priors = None
        self.n_classes = None
    
    def save(self, filepath):
        import json
        import zlib
        
        if filepath.endswith('.nexus'):
            filepath = filepath.replace('.nexus', '.lnexus')
        elif not filepath.endswith('.lnexus'):
            filepath += '.lnexus'
        
        model_data = {
            'model_type': 'GaussianNB',
            'framework': 'GRNexus',
            'language': 'Python',
            'version': '1.0',
            'n_classes': self.n_classes,
            'means': self.means.tolist() if self.means is not None else None,
            'variances': self.variances.tolist() if self.variances is not None else None,
            'priors': self.priors.tolist() if self.priors is not None else None
        }
        
        json_data = json.dumps(model_data)
        compressed_data = zlib.compress(json_data.encode('utf-8'))
        
        with open(filepath, 'wb') as f:
            f.write(compressed_data)
        print(f"Model saved to {filepath}")
    
    @classmethod
    def load(cls, filepath):
        import json
        import zlib
        
        if not os.path.exists(filepath):
            raise FileNotFoundError(f"File not found: {filepath}")
        if filepath.endswith('.nexus'):
            raise ValueError("Invalid file format. Expected .lnexus for ML models, got .nexus (Neural Network format)")
        
        with open(filepath, 'rb') as f:
            compressed_data = f.read()
        
        json_data = zlib.decompress(compressed_data).decode('utf-8')
        model_data = json.loads(json_data)
        
        if model_data['model_type'] != 'GaussianNB':
            raise ValueError(f"Invalid model type. Expected GaussianNB, got {model_data['model_type']}")
        
        model = cls()
        model.n_classes = model_data['n_classes']
        if model_data['means'] is not None:
            model.means = np.array(model_data['means'], dtype=np.float64)
        if model_data['variances'] is not None:
            model.variances = np.array(model_data['variances'], dtype=np.float64)
        if model_data['priors'] is not None:
            model.priors = np.array(model_data['priors'], dtype=np.float64)
        return model
    
    def __repr__(self):
        trained = 'trained' if self.means is not None else 'not trained'
        n_features = self.means.shape[1] if self.means is not None and len(self.means) > 0 else 0
        return f"<GaussianNB n_classes={self.n_classes or 0}, status={trained}, features={n_features}>"
    
    def fit(self, X, y):
        X = np.array(X, dtype=np.float64)
        y = np.array(y, dtype=np.float64)
        n_samples, n_features = X.shape
        self.n_classes = int(y.max()) + 1
        
        means = np.zeros((self.n_classes, n_features), dtype=np.float64)
        variances = np.zeros((self.n_classes, n_features), dtype=np.float64)
        priors = np.zeros(self.n_classes, dtype=np.float64)
        y_reshaped = y.reshape(-1, 1)
        
        x_data, x_arr = create_data(X)
        y_data, y_arr = create_data(y_reshaped)
        means_data, means_arr = create_data(means)
        vars_data, vars_arr = create_data(variances)
        priors_data, priors_arr = create_data(priors)
        
        result = _lib.naive_bayes_fit(ctypes.byref(x_data), ctypes.byref(y_data), self.n_classes,
                                       ctypes.byref(means_data), ctypes.byref(vars_data), ctypes.byref(priors_data))
        
        if result != 0:
            raise RuntimeError("Naive Bayes fit failed")
        
        self.means = means
        self.variances = variances
        self.priors = priors
        
        return self
    
    def predict(self, X):
        if self.means is None or self.variances is None or self.priors is None:
            raise ValueError("Model not trained")
        
        X = np.array(X, dtype=np.float64)
        n_samples = X.shape[0]
        
        output = np.zeros(n_samples, dtype=np.float64)
        priors_reshaped = self.priors.reshape(-1, 1)
        
        x_data, x_arr = create_data(X)
        means_data, means_arr = create_data(self.means)
        vars_data, vars_arr = create_data(self.variances)
        priors_data, priors_arr = create_data(priors_reshaped)
        output_data, output_arr = create_data(output)
        
        result = _lib.naive_bayes_predict(ctypes.byref(x_data), self.n_classes,
                                           ctypes.byref(means_data), ctypes.byref(vars_data), 
                                           ctypes.byref(priors_data), ctypes.byref(output_data))
        
        if result != 0:
            raise RuntimeError("Naive Bayes prediction failed")
        
        return output.astype(int)
