import logging

logger = logging.getLogger('DockOnSurf')


def select_confs(orig_conf_list, calc_dirs, energy_cutoff, code):
    """From a list of atomic configurations selects the most stable ones.

    Given a list of ase.Atoms configurations and an energy cutoff, selects all
    the structures that have an energy lower than, the energy of the most stable
    conformer plus the cutoff.

    @param orig_conf_list: List of ase.Atoms objects of the conformers
    @param calc_dirs: List of the directories of the finished calculations
    @param energy_cutoff: The maximum energy above the most stable configuration
    @param code: the code used to carry out the screening of structures.
    @return: list of the the most stable configurations within the energy cutoff
    """
    from copy import deepcopy
    from modules.formats import collect_energies

    conf_list = deepcopy(orig_conf_list)
    selected_ids = []
    conf_enrgs = collect_energies(calc_dirs, code, 'screening')

    for i, conf in enumerate(conf_list):
        conf.info['energy'] = conf_enrgs[i]

    sorted_list = sorted(conf_list, key=lambda conf: conf.info['energy'])
    lowest_e = sorted_list[0].info['energy']
    return [conf for conf in sorted_list
            if conf.info['energy'] < lowest_e + energy_cutoff]


def run_refinement(inp_vars):
    """Carries out the refinement of adsorbate-slab structures after screening.

    @param inp_vars: Calculation parameters from input file.
    """
    import os
    from modules.formats import collect_coords
    from modules.calculation import run_calc, check_finished_calcs
    if not os.path.isdir("screening"):
        err = "'screening' directory not found. It is needed in order to carry "
        "out the refinement of structures to be adsorbed"
        logger.error(err)
        raise FileNotFoundError(err)

    logger.info('Carrying out procedures for the refinement of '
                'adsorbate-surface structures.')

    finished_calcs, unfinished_calcs = check_finished_calcs('screening',
                                                            inp_vars['code'])
    if len(unfinished_calcs) == 0:
        logger.info(f"Found {len(finished_calcs)} structures of "
                    f"adsorbate-surface atomic configurations.")
    else:
        logger.info(f"Found {len(finished_calcs)} structures of "
                    f"adsorbate-surface atomic configurations whose calculation"
                    f" finished normally.")
        logger.warning(f"Found {len(unfinished_calcs)} calculations more that "
                       f"did not finish normally: {unfinished_calcs}. Using "
                       f"only the {len(finished_calcs)} ones who finished "
                       f"normally")

    conf_list = collect_coords(inp_vars['code'], 'screening', 'ase',
                               inp_vars['special_atoms'])
    logger.info(f"Found {len(conf_list)} structures on screening directory.")
    selected_confs = select_confs(conf_list, finished_calcs,
                                  inp_vars['energy_cutoff'], inp_vars['code'])
    run_calc('refinement', inp_vars, selected_confs)
    logger.info('Finished the procedures for the refinement of '
                'adsorbate-surface structures section.')