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 
# (dist_HH + 2*dist_LiH = 3.0 Angstrom)
dist_HH = 0.2
dist_LiH = 1.5 - 0.5 * dist_HH

# 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_conf = TwoProbeConfiguration(
    electrodes = (electrode_Li,electrode_Li),
    scattering_region_elements = elements,
    scattering_region_cartesian_coordinates = positions
    )

# Export the two-probe configuration to a VNL file
vnl_file = VNLFile("lih2li-ini.vnl")
vnl_file.addToSample(two_probe_conf, "lih2li-ini")