import hydra


def get_model(cfg, dataset, num_classes=None):
    """
    Instantiate and return a model based on the given configuration and data.

    Args:
        cfg (Any): The configuration object for the model.
        data (Any): The data object containing the input features and labels.

    Returns:
        Any: The instantiated model.
    """
    if cfg.task.graph_level:
        model = hydra.utils.instantiate(cfg.train.model, out_channels=dataset.num_tasks)
    else:
        if cfg.train.model._target_ == "graphsmodel.models.LP":
            model = hydra.utils.instantiate(
                cfg.train.model, num_classes=num_classes
            )
        else:
            model = hydra.utils.instantiate(
                cfg.train.model,
                in_channels=dataset.num_features,
                out_channels=num_classes,
            )

    return model


def get_optimizer(cfg, model):
    """
    Get an optimizer instance based on the provided configuration and model parameters.

    Args:
        cfg (Config): The configuration object specifying the optimizer settings.
        model (nn.Module): The model whose parameters will be optimized.

    Returns:
        Optimizer: The optimizer instance.

    """
    optimizer = hydra.utils.instantiate(cfg, model.parameters())
    return optimizer


def weights_to_cpu(weights):
    """
    Convert the weights to CPU.

    Args:
        weights (dict): The weights dictionary.

    Returns:
        dict: The weights dictionary with the weights on the CPU.
    """
    return {layer_name: param.cpu() for layer_name, param in weights.items()}
