from stereograph.classification import edge_conv_classifier
from stereograph.energy import energy_regressor
from stereograph.direction import direction_regressor

def load_model(config, checkpoint_path=None):
    """
    Load a model based on the provided configuration and optional checkpoint.
    Parameters
    ----------
    config : dict
        Configuration dictionary containing model and training parameters.
        Expected keys:
        - "network": dict with keys "model", "hidden_dim", "learning_rate", "dropout_rate", "num_layers".
        - "data": dict with key "input_features".
    checkpoint_path : str, optional
        Path to the checkpoint file to load the model from. If None, a new model instance is created.
    Returns
    -------
    model : object
        An instance of the model class specified in the configuration, either newly created or loaded from the checkpoint.
    """

    
    model_class = globals()[config["network"]["model"]]

    if checkpoint_path is None:
        model = model_class(
            input_dim=len(config["data"]["input_features"]),
            hidden_dim=config["network"]["hidden_dim"],
            learning_rate=config["network"]["learning_rate"],
            dropout_rate=config["network"]["dropout_rate"],
            num_layers=config["network"]["num_layers"],
        )
    else:
        model = model_class.load_from_checkpoint(
            checkpoint_path,
            input_dim=len(config["data"]["input_features"]),
            hidden_dim=config["network"]["hidden_dim"],
            learning_rate=config["network"]["learning_rate"],
            dropout_rate=config["network"]["dropout_rate"],
            num_layers=config["network"]["num_layers"],
        )
    return model
