from pathlib import Path
from ase.db import connect
from ase.constraints import ExpCellFilter
from ase.optimize import BFGS
from ase.build import bulk
from ase.dft.bandgap import bandgap
from gpaw import GPAW, PW


if Path('database.db').is_file():
    Path('database.db').unlink()

structures = ['Si', 'Ge', 'C']
db = connect('database.db')

for f in structures:
    db.write(bulk(f))

for row in db.select():
    atoms = row.toatoms()
    calc = GPAW(mode=PW(400),
                kpts=(4, 4, 4),
                txt=f'{row.formula}-gpaw.txt', xc='LDA')
    atoms.calc = calc
    atoms.get_stress()
    filter = ExpCellFilter(atoms)
    opt = BFGS(filter)
    opt.run(fmax=0.05)
    db.write(atoms=atoms, relaxed=True)


for row in db.select(relaxed=True):
    atoms = row.toatoms()
    calc = GPAW(mode=PW(400),
                kpts=(4, 4, 4),
                txt=f'{row.formula}-gpaw.txt', xc='LDA')
    atoms.calc = calc
    atoms.get_potential_energy()
    bg, _, _ = bandgap(calc=atoms.calc)
    db.update(row.id, bandgap=bg)