from ATK.TwoProbe import *
from numpy import *

# Define constants for the fcc lattice
a     = 4.078250
d     = a/sqrt(2)
d_sep = a/sqrt(3)

# Define vector generators for 3x3 layers
u  = array([1.0,             0.0, 0.0])*d
v  = array([0.5,     sqrt(3)/2.0, 0.0])*d
du = array([0.5, 1.0/2.0/sqrt(3), 0.0])*d

# Setup the A layer
n = 3
A = reshape(zeros(3*n*n), (n*n,3))*1.0
for j in range(n):
    for i in range(n):
        A[i+n*j] = i*u + j*v

# Setup the B and C layers
B1 = A + 1.0*du
B2 = A + 1.0*du
C  = A + 2.0*du

# Setup DTB molecule
dCC  = 1.397 
dCH  = 1.084 
dCS  = 1.754
dDTB = 2*(dCC + dCS)

n = 6
DTB = reshape(zeros(2*3*n), (2*n,3))*1.0
for i in range(n):
    t = i*pi/3.0
    DTB[i] = array([0.0, sin(t), cos(t)])*dCC
    if (i%3 == 0):
        DTB[i+n] = array([0.0, sin(t), cos(t)])*(dCC + dCS)
    else:
        DTB[i+n] = array([0.0, sin(t), cos(t)])*(dCC + dCH)

# Translate left S atom to z = 0
DTB[:,2] += (dCC + dCS)

# Define AB-DTB-BC central region
dAuDTB = 1.7

B1[:,2]  += d_sep
DTB[:,2] += d_sep + dAuDTB
B2[:,2]  += d_sep + dAuDTB + dDTB + dAuDTB
C[:,2]   += d_sep + dAuDTB + dDTB + dAuDTB + d_sep
DTB      += 6.0*du

# Central region elements
scattering_elements  = 18*[Gold]
scattering_elements += 6*[Carbon] + [Sulfur] + 2*[Hydrogen] + [Sulfur] + 2*[Hydrogen]
scattering_elements += 18*[Gold]

# Central region positions
scattering_positions  = concatenate((A,B1,DTB,B2,C)).tolist()
scattering_positions *= Angstrom

# Electrode elements and positions
electrode_elements  = 3*[Gold]
electrode_positions = [
    [0.0,      0.0,        0.0],
    [0.5*d, d*0.5/sqrt(3), 1*d_sep],
    [1.0*d, d*1.0/sqrt(3), 2*d_sep]
    ]*Angstrom

# Electrode unit cell
electrode_unitcell = [
    [d,               0.0, 0.0],
    [0.5*d, sqrt(3)/2.0*d, 0.0],
    [0.0,             0.0, 3.0*d_sep]
    ]*Angstrom

# Electrode
electrode = PeriodicAtomConfiguration(
    electrode_unitcell,
    electrode_elements,
    electrode_positions
    )

# Two-probe configuration
two_probe = TwoProbeConfiguration(
    (electrode,electrode),
    scattering_elements,
    scattering_positions,
    electrode_repetitions=[[3,3],[3,3]],
    equivalent_atoms = ([0,4],[2,43])
    )

# Export the two-probe system to an VNL file.
vnl_file=VNLFile('au-dtb-au.vnl')
vnl_file.addToSample(two_probe, 'au-dtb-au')