from copy import deepcopy


def create_SAC_algorithm(variant, *args, **kwargs):
    from .sac import SAC

    algorithm = SAC(*args, **kwargs)

    return algorithm


def create_SQL_algorithm(variant, *args, **kwargs):
    from .sql import SQL

    algorithm = SQL(*args, **kwargs)

    return algorithm

def create_MVE_algorithm(variant, *args, **kwargs):
    from .mve_sac import MVESAC

    algorithm = MVESAC(*args, **kwargs)

    return algorithm

def create_MultitaskCQL_algorithm(variant, *args, **kwargs):
    from cuds.algorithms.multi_task_cql import MultitaskCQL

    algorithm = MultitaskCQL(*args, **kwargs)

    return algorithm

def create_BC_algorithm(variant, *args, **kwargs):
    from cuds.algorithms.bc import BC

    algorithm = BC(*args, **kwargs)

    return algorithm


ALGORITHM_CLASSES = {
    'SAC': create_SAC_algorithm,
    'SQL': create_SQL_algorithm,
    'MultitaskCQL': create_MultitaskCQL_algorithm,
    'BC': create_BC_algorithm,
}


def get_algorithm_from_variant(variant,
                               *args,
                               **kwargs):
    algorithm_params = variant['algorithm_params']
    algorithm_type = algorithm_params['type']
    algorithm_kwargs = deepcopy(algorithm_params['kwargs'])
    if 'learn_bc_policy' or 'use_old_policy' in algorithm_kwargs:
        algorithm_kwargs['policy_variant'] = variant
    if 'use_fqe' in algorithm_kwargs:
        algorithm_kwargs['Q_variant'] = variant
    algorithm = ALGORITHM_CLASSES[algorithm_type](
        variant, *args, **algorithm_kwargs, **kwargs)

    return algorithm
