#!/usr/bin/env python

import numpy as np

from ase import Atoms, Atom
#from ase.visualize import view
#from ase.all import *
from gpaw import GPAW, Mixer
from ase.lattice.surface import fcc110
from gpaw.transport.jstm import dump_hs, dump_lead_hs
from ase.io import read

# surface calculation
srf = read('POSCAR', format='vasp') 
#srf = srf.repeat(rep=(3,1,1))
srf.pbc = (1, 1, 0)
#view(srf)

# choose a grid spacing ~0.2 but divisible by 8
h = 0.2
div = 8
cell_c = np.sum(srf.get_cell()**2, axis=1)**0.5
N_c = div * np.round(cell_c / (div * h))

calc = GPAW(gpts=N_c,
            mixer=Mixer(0.1, 5, weight=50),
            width=0.1, 
            mode='lcao',
            basis='dzp',
            txt='dumphs.txt',
            usesymm=False,
            #kpts=(3,3,1),
            convergence={'energy': 0.05, # eV
                         'density': 5.0e-3,
                         'eigenstates': 1.0e-6,
                         'bands': 'occupied'})

srf.set_calculator(calc)
srf.get_potential_energy()
calc.write('srf')

# Dump the overlap matrix and the Hamiltonian matrix to the local directory. 
# Here the keyword 'cvl' refers  to the number of basis functions in the 
# in the convergence layer, i.e. for the present system four atomic layers are used.

dump_hs(calc, 'srf', region='surface', cvl=2*6*15) 
                                                 
# tip calculation 
a = 0.75 # lattice constant
tip = Atoms('H12', pbc=(1, 1, 0), cell=[5, 5, 12 * a + 7])
tip.positions[:,2] = [i * a for i in range(12)]
tip.positions[:] += (tip.cell / 2.0)[0, :] + (tip.cell / 2.0)[1, :]
tip.translate([0, 0, 6])
#view(tip)

# choose a grid spacing ~0.2 but divisible by 8
cell_c = np.sum(tip.get_cell()**2, axis=1)**0.5
N_c = div * np.round(cell_c / (div * h))
calc.set(gpts=N_c)

tip.set_calculator(calc)
tip.get_potential_energy()
calc.write('tip')
dump_hs(calc, 'tip', region='tip', cvl=5) # dump overlap and hamiltonian matrix
calc.set(kpts=(1, 1, 7)) # for the lead calculations we use kpoints in the z-direction

# surface principal layer calculation
#srf_p = fcc110('Cu', size=(2, 2, 4)) 
srf_p = fcc110('Cu', size=(3, 2, 4), a=3.60431112)
#srf_p=read('/home/test_y/CuO_9layers/CuPL/POSCAR',format='vasp')
srf_p.pbc = (1, 1, 1)

# choose a grid spacing ~0.2 but divisible by 8
cell_c = np.sum(srf_p.get_cell()**2, axis=1)**0.5
N_c = div * np.round(cell_c / (div * h))
calc.set(gpts=N_c)

srf_p.set_calculator(calc)
srf_p.get_potential_energy()
dump_lead_hs(calc, 'srf_p') # dump overlap and hamiltonian matrix

# tip principal layer calculation
tip_p = Atoms('H4', pbc=(1,1,1), cell=[5, 5, 4*a])
tip_p.positions[:,2] = [i * a for i in range(4)]
tip_p.positions[:] += (tip_p.cell / 2.0)[0, :] + (tip_p.cell / 2.0)[1, :]
#view(tip_p)

# choose a grid spacing ~0.2 but divisible by 8
cell_c = np.sum(tip_p.get_cell()**2, axis=1)**0.5
N_c = div * np.round(cell_c / (div * h))
calc.set(gpts=N_c)

tip_p.set_calculator(calc)
tip_p.get_potential_energy()
dump_lead_hs(calc, 'tip_p') 
# dump overlap and hamiltonian matri