#!/usr/bin/env python
# =======================================
#    STM of Al(fcc)100(N x M unit cell)
# =======================================
from Dacapo import Dacapo
from ASE import Atom,ListOfAtoms
from Numeric import sqrt,array
from ASE.Visualization.VMD import VMD
import os

a0     = 4.05         # cubic fcc lattice constant
N      = 2            # repetition along x
M      = 2            # repetition along y
layers = 2            # slab layers
electronsperatom = 3 
vaclay = 5            # interlayer dist = a0/2


# ---------------------------------------------------
# ------- Setup the structure/calculation -----------
# ---------------------------------------------------

atoms   = ListOfAtoms([])
for n in range(layers):
    for i in range(N):
        for j in range(M):
            scaledpos = [(i+(n%2)/2.)/sqrt(2.),(j+(n%2)/2.)/sqrt(2.),-n/2.]
            atoms.append(Atom('Al', a0*array(scaledpos)))
                                       
                                       
unitcell = [[N/sqrt(2.), 0.0,        0.0],
            [0.0,        M/sqrt(2.), 0.0],
            [0.0,        0.0,        (vaclay+layers)/2.]]

atoms.SetUnitCell(a0*array(unitcell),fix=True)

# VMD(atoms)

# check if we have already calculated the nc file
if not os.path.isfile('Al100.nc'):

    calc = Dacapo(planewavecutoff = 150, 
              nbands = 10 + len(atoms)*electronsperatom/2, 
              kpts=(4,4,1),
              xc = 'LDA',
              usesymm=True,
              out = 'Al100.nc',txtout = 'Al100.txt') 

    atoms.SetCalculator(calc) 

else: 
    atoms = Dacapo.ReadAtoms('Al100.nc')

energy = atoms.GetPotentialEnergy()

# =======================================
#    STM of Al(fcc)100(N x M unit cell)
# =======================================
# --------------------------------------------------
# ------------ STM image generation part -----------
# --------------------------------------------------

from Dacapo.ElectronicStates import ElectronicStates
from Dacapo import Dacapo
from ASE.Utilities.STMTool import STMTool
from ASE.Utilities.STMLineScan import STMLineScan

atoms = Dacapo.ReadAtoms('Al100.nc')
from ASE.Visualization.VMD import VMD
# VMD(atoms)

loe = ElectronicStates(filename='Al100.nc')
stm = STMTool(loe,contourvalue=1.0e-4, channel="s",
             normalaxis=2, smoothfactor=0.1)    # available channels: s,px,py,p



from pylab import *
import matplotlib

# select a linescan
linescan = STMLineScan(stm,fftnumber=12,axis=1)

# contour surface plot 
extent = stm.GetExtentOfContourSurface()
im2 = imshow(stm.GetContourSurface().Array,
             interpolation='bicubic',
             origin='lower',cmap=cm.jet, 
             extent = extent) 
colorbar()

# add position of line scan to contour plot
linex,liney  = linescan.GetLine()
plot(linex,liney,linewidth=2,color='black')
axis(extent)

savefig('contourplot')

# line scan plot
figure()

values = linescan.GetArray()
plot(values)
xlabel('Distance along surface')
ylabel('Height above surface')
title('Linescan at FFT number 12 along the y-axis') 
grid(True)

savefig('linescan')