from ase.build import bulk
from ase.calculators.espresso import Espresso
from multiprocessing import Pool
from asetest import datafiles


# Ugly pseudopotentials boilerplate
pseudo_dir = datafiles.paths['espresso'][0]
pseudopotentials = {}
for path in pseudo_dir.glob('*.UPF'):
    fname = path.name
    # Names are e.g. si_lda_v1.uspp.F.UPF
    symbol = fname.split('_', 1)[0].capitalize()
    pseudopotentials[symbol] = fname


def mkcalc(sym):
    # Tons of boilerplate here.
    # Only two things are important:
    #  * We set the command to something with mpiexec
    #  * We set the directory to avoid clashes with other processes

    input_data = {'system': {'occupations': 'smearing',
                             'smearing': 'fermi-dirac',
                             'degauss': 0.02}}

    return Espresso(
        pseudo_dir=str(pseudo_dir),  # FIXME doesn't like Path objects
        ecutwfc=90,  # Ry
        pseudopotentials=pseudopotentials,
        command='mpiexec -n 2 pw.x -in PREFIX.pwi > PREFIX.pwo',
        directory=f'calc-{sym}',
        input_data=input_data,
    )


def calc(sym):
    atoms = bulk(sym)
    atoms.calc = mkcalc(sym)
    atoms.get_potential_energy()
    return atoms


symbols = ['Cu', 'Ag', 'Au', 'Pt',
           'Pd', 'Ti', 'Al', 'C',
           'Si', 'Fe', 'Ni']


with Pool(4) as pool:
    out_images = pool.map(calc, symbols)
    for img in out_images:
        print(img)