from ATK.TwoProbe import *

# Li chain lattice constant
a = 2.90
# Construct the electrode unit cell
unit_cell = [ [3*a, 0.0, 0.0 ],
              [0.0, 3*a, 0.0 ],
              [0.0, 0.0, 4*a ] ] * Angstrom

# Define the electrode
electrode_Li = PeriodicAtomConfiguration(
    super_cell_vectors=unit_cell,
    elements=4*[Lithium],
    fractional_coordinates=[(0.5, 0.5, float(i)/4.0) for i in range(0,4)]
    )

# Setup the two-probe scattering region

# Distances between Li-H and H-H (found from relaxing a Li4H2 cluster)
dist_HH = 0.804
dist_LiH = 2.366

# The atoms in the central region
elements = 3*[Lithium] + 2*[Hydrogen] + 3*[Lithium]

positions = [
    (0.0, 0.0, 0*a),
    (0.0, 0.0, 1*a),
    (0.0, 0.0, 2*a),
    (0.0, 0.0, 2*a + dist_LiH),
    (0.0, 0.0, 2*a + dist_LiH + dist_HH),
    (0.0, 0.0, 2*a + dist_LiH + dist_HH + dist_LiH + 0*a),
    (0.0, 0.0, 2*a + dist_LiH + dist_HH + dist_LiH + 1*a),
    (0.0, 0.0, 2*a + dist_LiH + dist_HH + dist_LiH + 2*a)
    ] * Angstrom
    
# Combine electrode and scattering region
# into a two-probe system
two_probe = TwoProbeConfiguration(
    electrodes = (electrode_Li,electrode_Li),
    scattering_region_elements              = elements,
    scattering_region_cartesian_coordinates = positions
    )

# Export the two-probe system to VNL file.
vnl_file = VNLFile("lih2li.vnl")
vnl_file.addToSample(two_probe, "lih2li")