import os
import logging

logger = logging.getLogger('DockOnSurf')


def create_bak_calc_dir(run_type):
    """Checks if calculations directory already exists, backs it up if so and
    creates an empty one.
    @param run_type: Type of calculation. 'isolated', 'screening' or
    'refinement'
    """
    dir_name = run_type
    bak_num = 0
    while dir_name in os.listdir("."):
        bak_num += 1
        dir_name = dir_name.split(".")[0] + f".bak{bak_num}"
    if bak_num > 0:
        os.rename(run_type, dir_name)
        logger.warning(f"'{run_type}' directory already present. Moved former "
                       f"directory to {dir_name}")
    os.mkdir(run_type)


def prep_cp2k(inp_file, run_type, atms_list):
    """Prepares the directories to run isolated calculation with CP2K.

    @param inp_file: CP2K Input file to run the calculations with.
    @param run_type: Type of calculation. 'isolated', 'screening' or
        'refinement'
    @param atms_list: list of ase.Atoms objects to run the calculation of.
    @return: None
    """
    from shutil import copy
    import ase.io
    from pycp2k import CP2K
    cp2k = CP2K()
    cp2k.parse(inp_file)
    force_eval = cp2k.CP2K_INPUT.FORCE_EVAL_list[0]
    coord_file = force_eval.SUBSYS.TOPOLOGY.Coord_file_name

    # Creating and setting up directories for every atoms configuration.
    for i, conf in enumerate(atms_list):
        os.mkdir(f'{run_type}/conf_{i}')
        copy(inp_file, f'{run_type}/conf_{i}/')
        ase.io.write(f'{run_type}/conf_{i}/{coord_file}', conf)


def run_calc(run_type, inp_vars, atms_list):
    """Directs the calculation run according to the provided arguments.

    @param run_type: Type of calculation. 'isolated', 'screening' or
    'refinement'
    @param inp_vars: Calculation parameters from input file.
    @param atms_list: List of ase.Atoms objects containing the sets of atoms
    aimed to run the calculations of.
    """
    run_types = ['isolated', 'screening', 'refinement']
    run_type_err = f"'run_type' should be one of the following: {run_types}"
    if not isinstance(run_type, str) or run_type.lower() not in run_types:
        logger.error(run_type_err)
        raise ValueError(run_type_err)

    logger.info(f"Running {run_type} calculation with {inp_vars['code']} on "
                f"{inp_vars['batch_q_sys']}")
    create_bak_calc_dir(run_type)
    if run_type == 'isolated':
        if inp_vars['code'] == 'cp2k':
            prep_cp2k(inp_vars['isol_inp_file'], run_type, atms_list)