import shutil
import subprocess
import xml.etree.ElementTree as ET
import z2pack

qedir = '/opt/josephson/quantum-espresso/5.4/'
wandir = '/home/olmos.asar/programas/wannier90-1.2/'

mpirun = 'srun -n 16 '
pwcmd = mpirun + qedir + 'pw.x '
pw2wancmd = mpirun + qedir + 'pw2wannier90.x '
wancmd = wandir + 'wannier90.x '

material_name = 'C'
num_wcc =8
z2cmd = (wancmd + ' ' + material_name + ' -pp;' +
         pwcmd + '< ' + material_name + '.nscf.in >& pw.log;' +
         pw2wancmd + '< ' + material_name + '.pw2wan.in >& pw2wan.log;')

if not os.path.exists('./scf'):
   os.makedirs('./scf')
   shutil.copyfile('input/' + material_name + '.scf.in', 'scf/' + material_name +'.scf.in')
   print(pwcmd + '< ' + material_name + '.scf.in >& scf.out')
   subprocess.call(pwcmd + '< ' + material_name + '.scf.in > scf.out', shell=True, cwd='./scf')

system = z2pack.fp.System(input_files=
                          ['input/'+material_name+'.nscf.in',
                           'input/'+material_name+'.pw2wan.in',
                           'input/'+material_name+'.win'],
                          kpt_fct=
                          [z2pack.fp.kpoint.qe,
                           z2pack.fp.kpoint.wannier90],
                          kpt_path=
                          [material_name+".nscf.in",
                           material_name+".win"],
                          command=z2cmd,
                          executable='/bin/bash',
                          mmn_path=material_name+'.mmn'
)

result = z2pack.surface.run(
    system=system,
    surface=lambda t1, t2: [t1, t2, 0],
    num_lines=201,
    pos_tol=0.001,
    iterator=range(40, 60, 5),
    gap_tol=0.01,
    move_tol=0.01
)

print('Z2=')
print(z2pack.invariant.z2(result))