from aiforge.core.registry import register
from aiforge.core.context import BaseOperator

@register("train_rf")
class TrainRF(BaseOperator):
    """Train a scikit-learn Random Forest Classifier."""
    def run(self, data, labels=None, n_estimators=100):
        # We handle data containing both X and y, or passed via kwargs
        from sklearn.ensemble import RandomForestClassifier
        model = RandomForestClassifier(n_estimators=n_estimators)
        
        # If the input data is a tuple (X, y)
        if isinstance(data, tuple) and len(data) == 2:
            X, y = data
            model.fit(X, y)
        else:
            if labels is None:
                raise ValueError("Labels must be provided either as a tuple (X, y) or as a kwarg 'labels'.")
            model.fit(data, labels)
            
        return model

@register("predict_rf")
class PredictRF(BaseOperator):
    """Make predictions using a trained Random Forest Classifier."""
    def run(self, data, model):
        # Data is the input features (X)
        return model.predict(data)

@register("flatten_data")
class FlattenData(BaseOperator):
    """Flatten multi-dimensional arrays into 1D arrays (Bridge: Vision/DL -> ML)."""
    def run(self, data):
        import numpy as np
        if not isinstance(data, np.ndarray):
            data = np.array(data)
        return data.flatten()

@register("train_svm")
class TrainSVM(BaseOperator):
    """Train a scikit-learn Support Vector Classifier."""
    def run(self, data, labels=None, C=1.0, kernel='rbf'):
        from sklearn.svm import SVC
        model = SVC(C=C, kernel=kernel)
        if isinstance(data, tuple) and len(data) == 2:
            X, y = data
            model.fit(X, y)
        else:
            if labels is None:
                raise ValueError("Labels must be provided either as a tuple (X, y) or as a kwarg 'labels'.")
            model.fit(data, labels)
        return model

@register("predict_svm")
class PredictSVM(BaseOperator):
    """Make predictions using a trained SVM."""
    def run(self, data, model):
        return model.predict(data)

@register("train_kmeans")
class TrainKMeans(BaseOperator):
    """Train a scikit-learn KMeans clustering model."""
    def run(self, data, n_clusters=8):
        from sklearn.cluster import KMeans
        model = KMeans(n_clusters=n_clusters, n_init='auto')
        model.fit(data)
        return model

@register("predict_kmeans")
class PredictKMeans(BaseOperator):
    """Predict the closest cluster each sample in data belongs to."""
    def run(self, data, model):
        return model.predict(data)

@register("train_pca")
class TrainPCA(BaseOperator):
    """Train a Principal Component Analysis (PCA) model."""
    def run(self, data, n_components=None):
        from sklearn.decomposition import PCA
        model = PCA(n_components=n_components)
        model.fit(data)
        return model

@register("transform_pca")
class TransformPCA(BaseOperator):
    """Apply dimensionality reduction to data using a trained PCA model."""
    def run(self, data, model):
        return model.transform(data)
