"""Module for building atomic structures"""

from ASE import ListOfAtoms, Atom


def BCC100(symbol, a, layers, h):
    """Build a bcc(100) surface

    symbol: chmical symbol ('H', 'Li', ...)
    a     : lattice constant
    layers: number of layers
    h     : height of unit cell"""

    # Distance between layers:
    z = a / 2

    assert h > layers * z, 'unit cell too small for ' + str(layers) + ' layers'
    
    # Start with an empty ListOfAtoms object:
    atoms = ListOfAtoms([], periodic=True)
    
    # Fill in the atoms using scaled coordinates:
    for n in range(layers):
        position = [0.5 * (n % 2),
                    0.5 * (n % 2),
                    -n * z / h]
        atoms.append(Atom(symbol, position))
        
    # The unit cell is orthorhombic:
    atoms.SetUnitCell((a, a, h))
    # The unit cell could also have been set like this:
    # atoms.SetUnitCell([(a, 0, 0),
    #                    (0, a, 0),
    #                    (0, 0, h)])
    return atoms


if __name__ == '__main__':
    bcc = BCC100('Al', 3.25, 4, 20.0)
    from ASE.Visualization.RasMol import RasMol
    p = RasMol(bcc, (4, 4, 2))
    raw_input('Press [enter] to finish. ')