from QuantumATK import Angstrom
from QuantumATK import BulkConfiguration
from QuantumATK import Hydrogen, Platinum, Oxygen
from QuantumATK import MoleculeConfiguration
from QuantumATK import repeat
from QuantumATK import UnitCell

# -------------------------------------------------------------
# Create the substrate configuration.
# -------------------------------------------------------------
# Set up lattice
vector_a = [3.9236, 0.0, 0.0]*Angstrom
vector_b = [0.0, 3.9236, 0.0]*Angstrom
vector_c = [0.0, 0.0, 3.9236]*Angstrom
lattice = UnitCell(vector_a, vector_b, vector_c)

# Define elements
elements = [Platinum, Platinum, Platinum, Platinum]

# Define coordinates
fractional_coordinates = [[ 0. ,  0. ,  0. ],
                          [ 0.5,  0.5,  0. ],
                          [ 0.5,  0. ,  0.5],
                          [ 0. ,  0.5,  0.5]]

# Set up configuration
bulk_configuration = BulkConfiguration(
    bravais_lattice=lattice,
    elements=elements,
    fractional_coordinates=fractional_coordinates
)

# Repeat the bulk_configuration to create the substrate.
substrate = repeat(bulk_configuration, na=3, nb=3, nc=2, stack_systems=False)

# -------------------------------------------------------------
# Create a water molecule.
# -------------------------------------------------------------
# Define elements
elements = [Hydrogen, Oxygen, Hydrogen]

# Define coordinates
cartesian_coordinates = [[ 0.474853, -0.026505, -0.631087],
                         [-0.34219 ,  0.029981, -0.147925],
                         [-0.132663, -0.003476,  0.779012]]*Angstrom

# Set up configuration
water = MoleculeConfiguration(
    elements=elements,
    cartesian_coordinates=cartesian_coordinates
)