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

@register("tensor_conversion")
class ToTensor(BaseOperator):
    """Convert numpy arrays or standard lists to PyTorch tensors."""
    def run(self, data):
        import torch
        import numpy as np
        if isinstance(data, np.ndarray):
            return torch.from_numpy(data)
        return torch.tensor(data)

@register("normalize_tensor")
class NormalizeTensor(BaseOperator):
    """Normalize a PyTorch tensor (data - mean) / std."""
    def run(self, data, mean=0.0, std=1.0):
        # Assumes data is already a torch tensor
        return (data - mean) / std

@register("tensor_to_numpy")
class TensorToNumpy(BaseOperator):
    """Convert a PyTorch tensor to a NumPy array."""
    def run(self, data):
        return data.detach().cpu().numpy()

@register("build_mlp")
class BuildMLP(BaseOperator):
    """Build a simple Multilayer Perceptron (MLP) using PyTorch."""
    def run(self, data=None, input_size=10, hidden_size=64, output_size=2):
        import torch.nn as nn
        model = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, output_size),
            nn.Softmax(dim=-1)
        )
        return model

@register("train_mlp")
class TrainMLP(BaseOperator):
    """Train a basic PyTorch model (like an MLP) for a few epochs."""
    def run(self, data, labels=None, model=None, epochs=5, lr=0.01):
        import torch
        import torch.nn as nn
        import torch.optim as optim
        
        if model is None:
            raise ValueError("A PyTorch model must be provided.")
        
        # If the input data is a tuple (X, y)
        if isinstance(data, tuple) and len(data) == 2:
            X, y = data
        else:
            if labels is None:
                raise ValueError("Labels must be provided either as a tuple (X, y) or as a kwarg 'labels'.")
            X, y = data, labels
            
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(model.parameters(), lr=lr)
        
        for epoch in range(epochs):
            optimizer.zero_grad()
            outputs = model(X)
            # Ensure targets and outputs types match
            if isinstance(y, torch.Tensor) and y.dtype == torch.float:
                y = y.long()
            loss = criterion(outputs, y)
            loss.backward()
            optimizer.step()
            
        return model
