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

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

    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)

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

    conf_list = collect_coords(finished_calcs, inp_vars['code'], 'screening',
                               inp_vars['special_atoms'])
    selected_confs = select_confs(conf_list, finished_calcs,
                                  inp_vars['energy_cutoff'], inp_vars['code'])
    logger.info(f"Selected {len(selected_confs)} structures to carry out the"
                f" refinement")
    run_calc('refinement', inp_vars, selected_confs)
    most_stable_conf = select_confs(conf_list, finished_calcs, 0,
                                    inp_vars['code'])[0]
    logger.info('Finished the procedures for the refinement of '
                'adsorbate-surface structures section. Most stable structure '
                f'is {most_stable_conf}')