#!/usr/bin/env python

import Numeric
from ASE import Atom, ListOfAtoms
from Dacapo import Dacapo
from ASE.Utilities.Wannier import Wannier,CompositeWannierCells

# Read the atoms for the DFT calculation of the scattering region 
#(1 principal layer on each side of S must be included) 
atoms1 = Dacapo.ReadAtoms("out_scatter_kpt-0.375x-0.375.nc",save_memory=True)
# Read the atoms for the DFT calculation of the lead
atoms2 = Dacapo.ReadAtoms("out_lead.nc",save_memory=True)

calc1 = atoms1.GetCalculator()
calc2 = atoms2.GetCalculator()

# Initialize a Wannier object for the Wannier functions in S
wannier1 = Wannier(numberofwannier=432,calculator=calc1,occupationenergy=3.0)
wannier1.ReadZIBlochMatrix('zibloch_scatter.pickle')
wannier1.ReadRotation('Wannier_scatter')

# Initialize a Wannier object for the Wannier functions in the lead
wannier2 = Wannier(numberofwannier=18,calculator=calc2,occupationenergy=3.0)
wannier2.ReadZIBlochMatrix('zibloch_lead.pickle')
wannier2.ReadRotation('Wannier_lead')

# Setup cells describing 1 principal layer and the scattering region
cell1=CompositeWannierCells.WannierCell(wannier2,atoms2,repeat=[1,3,3])
cell2=CompositeWannierCells.WannierCell(wannier1,atoms1,repeat=[1,1,1])
cell3=CompositeWannierCells.WannierCell(wannier2,atoms2,repeat=[1,3,3])

# Glue the three regions together and specify the k-point
container=CompositeWannierCells.CompositeWannierCell([cell1,cell2,cell3],kpoint=[-0.375,-0.375],executable='/home/niflheim/lhansen/thul/dacapo/dacapo.run')
# Construct the Hamiltonian
container.WriteHSMatrixToNetCDFFile('HS.nc')