from math import sin, cos, pi

from ase import Atoms
from ase.build import fcc111, fcc100, add_adsorbate
from ase.db import connect
from ase.constraints import FixAtoms
from ase.lattice.cubic import FaceCenteredCubic
from ase.cluster import wulff_construction


def get_systems():
    systems = []

    cell = (5, 5, 5)
    atoms = Atoms('H2', [(0, 0, 0), (0, 0, 1.4)], cell=cell)
    atoms.center()
    systems.append((atoms, 'Hydrogen molecule'))

    atoms = FaceCenteredCubic(
        directions=[[1, -1, 0], [1, 1, 0], [0, 0, 1]],
        size=(2, 2, 2),
        symbol='Cu',
        pbc=(1, 1, 1))
    atoms.rattle(stdev=0.1, seed=42)
    systems.append((atoms, 'Shaken bulk copper'))

    a = 2.70
    c = 1.59 * a

    slab = Atoms('2Cu', [(0., 0., 0.), (1 / 3., 1 / 3., -0.5 * c)],
                 tags=(0, 1),
                 pbc=(1, 1, 0))
    slab.set_cell([(a, 0, 0),
                   (a / 2, 3**0.5 * a / 2, 0),
                   (0, 0, 1)])
    slab.center(vacuum=3, axis=2)
    mask = [a.tag == 1 for a in slab]
    slab.set_constraint(FixAtoms(mask=mask))
    systems.append((slab, 'Distorted Cu(111) surface'))

    zpos = cos(134.3 / 2.0 * pi / 180.0) * 1.197
    xpos = sin(134.3 / 2.0 * pi / 180.0) * 1.19
    co = Atoms('CO', positions=[(-xpos + 1.2, 0, -zpos),
                                (-xpos + 1.2, -1.1, -zpos)])
    slab = fcc111('Au', size=(2, 2, 2), orthogonal=True)
    add_adsorbate(slab, co, 1.5, 'bridge')
    slab.center(vacuum=6, axis=2)
    slab.set_pbc((True, True, False))
    constraint = FixAtoms(mask=[a.tag == 2 for a in slab])
    slab.set_constraint(constraint)
    systems.append((slab, 'CO on Au(111) surface'))

    atoms = Atoms(symbols='C5H12',
                  cell=[16.83752497, 12.18645905, 11.83462179],
                  positions=[[5.90380523, 5.65545388, 5.91569796],
                             [7.15617518, 6.52907738, 5.91569796],
                             [8.41815022, 5.66384716, 5.92196554],
                             [9.68108996, 6.52891016, 5.91022362],
                             [10.93006206, 5.65545388, 5.91569796],
                             [5.00000011, 6.30002353, 5.9163716],
                             [5.88571848, 5.0122839, 6.82246859],
                             [5.88625613, 5.01308931, 5.01214155],
                             [7.14329342, 7.18115393, 6.81640316],
                             [7.14551332, 7.17200869, 5.00879027],
                             [8.41609966, 5.00661165, 5.02355167],
                             [8.41971183, 5.0251482, 6.83462168],
                             [9.69568096, 7.18645894, 6.8078633],
                             [9.68914668, 7.16663649, 5.00000011],
                             [10.95518898, 5.02163182, 6.8289018],
                             [11.83752486, 6.29836826, 5.90274952],
                             [10.94464142, 5.00000011, 5.01802495]])
    systems.append((atoms, 'Pentane molecule'))

    slab = fcc100('Cu', size=(2, 2, 2), vacuum=3.5)
    add_adsorbate(slab, 'C', 1.5, 'hollow')
    mask = [a.tag > 1 for a in slab]
    constraint = FixAtoms(mask=mask)
    slab.set_constraint(constraint)
    systems.append((slab, 'C/Cu(100)'))

    surfaces = [(1, 0, 0), (1, 1, 0), (1, 1, 1)]
    esurf = [0.9151, 0.9771, 0.7953]  # Surface energies
    size = 10  # number of atoms
    atoms = wulff_construction('Al', surfaces, esurf, size, 'fcc',
                               rounding='above')
    atoms.center(vacuum=6)
    atoms.rattle(0.2)
    systems.append((atoms, 'Alumninum cluster'))
    return systems


def create_database():
    systems = get_systems()
    db = connect('systems.db', append=False)
    for atoms, description in systems:
        name = atoms.get_chemical_formula()
        db.write(atoms, description=description, name=name)

    if False:
        for atoms, description in systems:
            for seed in range(5):
                a = atoms.copy()
                a.rattle(0.1, seed=seed)
                name = a.get_chemical_formula() + '-' + str(seed)
                db.write(a, description=description, seed=seed, name=name)


if __name__ == '__main__':
    create_database()