#!/usr/bin/env python

#	slices wien2k cell in x direction -> each slice is a cut parallel to y-z-plane
# calling program: programname nx ny nz x_start x_end
#
# creates 2d_slices with ny times nz points
# spacing between different slices is: length of unit cell in x-direction divided by nx
# x_start and x_end are defining the x-range that is sampled, including x_start slice and x_end slice
# 	- this can be used do include mirror symmetrie along sampling axis -> you only have to sample half of the unit cell
# 	- that is important only if you want to call this script several times for the same sturcture to run on several cores
# 
# asking for 200 slides each 10 by 10 pixels: programmname 200 10 10 1 200
# be careful not to create identical slices: if you ask for 200 10 10 0 200 slice 0 and 200 are equal

# for parallel execution ask:
# programmname 200 10 10 1 50&
# programmname 200 10 10 51 100&
# programmname 200 10 10 101 150&
# programmname 200 10 10 151 200&

# Theres the possibility to shift each slice in y and z direction by defining shift vectors (in line ~180)
# if you want to sample y from 0.01:1.01 define dy=1 dy_den=100
# this can be important if you want to avoid critical potential values because the sampling exactly hits a core


import sys, os, string, array
import time

class Lapw5:
    def __init__(self, head = None):
        self.tails = ["clmsum", "clmval", "vcoul", "vtotal", "r2v"]
        if head:
            self.head = head
        else:
            self.head = os.path.basename(os.getcwd())

        self.shells = [3, 3, 3]
        self.val = 2				#Wass soll ausgegeben werden? clmsum, clmval, vcoul..
        self.in5file = head + ".in5c"

    def make_ctrl(self, origin, xend, yend, nx, ny):
        output = open(self.in5file, "w")
        output.write("%d  %d  %d  %d\n" % (origin[0], origin[1], origin[2], origin[3]))
        output.write("%d  %d  %d  %d\n" % (xend[0], xend[1], xend[2], xend[3]))
        output.write("%d  %d  %d  %d\n" % (yend[0], yend[1], yend[2], yend[3]))
        output.write("%d  %d  %d\n" % (self.shells[0], self.shells[1], self.shells[2]))
        output.write("%d   %d\n" % (nx, ny))
        output.write("RHO     \n")
        if self.val != 0:
            output.write("ATU VAL NODEBUG\n")
        else:
            output.write("ATU TOT NODEBUG\n")
        output.write("NONORTHO")
        output.close()

    def _make_def(self):
        deffile = "lapw5.def"
        os.system("x lapw5c -d")
        if self.val != 1:
            input = open(deffile, "r")
            lines = input.readlines()
            input.close()
            output = open(deffile, "w")
            for l in lines:
               l = string.replace(l, "clmval", self.tails[self.val])
               #l = string.replace(l, ".in5", ".in5c")
               output.write(l)
            output.close()

    def run(self):
        self._make_def()
        os.system("lapw5c lapw5.def")



def _string2angle(s):
    s = string.strip(s)
    if s:
        angle = string.atof(s)
        if angle == 0.0:
            angle = 90.0
        return angle
    else:
        return 90.0

def get_lattice(struct):
    input = open(struct, "r")
    line = input.readline()
    line = input.readline()
    lat_type = string.strip(line[:4])
    line = input.readline()
    line = input.readline()
    abc = map(string.atof, [line[0:10],line[10:20],line[20:30]])
    angle = map(_string2angle, [line[30:40],line[40:50],line[50:60]])
    return abc, angle

def collect_charge(head, num):
    rhofile = head + ".rho"
    rho = array.array("d")
    for i in range(num):
        suffix = ".%04d" % i
        rhof = rhofile + suffix
        input = open(rhof, "r") 
        line = input.readline()
        while 1:
            line = input.readline()
            if not line: break
            for word in string.split(line):
                rho.append(float(word))
    return rho

def erase_files(head, num):
    rhofile = head + ".rho"
    in5file = head + ".in5c"
    for i in range(num):
        suffix = ".%04d" % i
        os.unlink(rhofile + suffix)
        os.unlink(in5file + suffix)
    

if __name__ == "__main__":
    import getopt

    def print_usage():
        usage = '''usage: venus.py [-h] [-v] [-p] nx ny nz start_z end_z
        -h   print help.
        -v   verbose mode.
        -p   preserve intermediate files\n'''
        sys.stderr.write(usage)

    try:
        options, resargs = getopt.getopt(sys.argv[1:], 'hvp')
    except getopt.error, s:
        sys.stderr.write(str(s)+"\n")
        print_usage()
        sys.exit(1)
    if len(resargs) < 5:
        print_usage()
        sys.exit(1)

    verbose = 0
    preserve = 0
    for option in options:
        if option[0] == "-h":
            print_usage()
            sys.exit()
        elif option[0] == "-v":
            verbose = 1
        elif option[0] == "-p":
            preserve = 1

    nx, ny, nz = map(int, resargs[0:3])
    x_start, x_end = map(int, resargs[3:5])
    print "using lapw5c for creating slices"
    print "slicing cell in x-direction -> each slice is parallel to y-z-plane"
    print "mesh:", nx, ny, nz
    print "x_start, x_end:", x_start, x_end
    

    head = os.path.basename(os.getcwd())
    lapw5 = Lapw5(head)
    x=x_start
    while x <= x_end :
        if verbose:
            print "    --> ", x, "    ", time.strftime('%d_%m_%Y %H:%M:%S')

				# no shift	
#        origin = [x , 0*nx , 0*nx , nx]
#        xend = [x , 1*nx , 0*nx , nx]
#        yend = [x , 0*nx , 1*nx , nx] 
        
        #shift in y direction and z direction
        # if you want to shift y by 0.06 of unit-cell-vector: 
        #				0.06=6/100 -> dy=6 dy_den=100
        # if you do not want to shift set: dy=0 dy_den=1
        #
        #I'm not sure if a limit is reached for very small shifts when the total denominater becomes very large
        #there may be problems on lapw5/lapw5c because this big number may not fit into their datatype
        #(eg if lapw5 is reading the denominator from case.in5 as integer number
        
        dy=0
        dy_den=1
        dz=0
        dz_den=1
        dx=0
        dx_den=1
               
        origin = [(x*dx_den+dx*nx)*dy_den*dz_den, nx*dx_den*(0+dy)*dz_den, nx*dx_den*dy_den*(0+dz), nx*dx_den*dy_den*dz_den]
        xend =   [(x*dx_den+dx*nx)*dy_den*dz_den, nx*dx_den*(dy_den+dy)*dz_den, nx*dx_den*dy_den*(0+dz), nx*dx_den*dy_den*dz_den]
        yend =   [(x*dx_den+dx*nx)*dy_den*dz_den, nx*dx_den*(0+dy)*dz_den, nx*dx_den*dy_den*(dz_den+dz), nx*dx_den*dy_den*dz_den]               
        
        lapw5.make_ctrl(origin, xend, yend, ny, nz)
        lapw5.run()
        
        suffix = ".%04d" % x
        os.rename(head + ".in5c", head + ".in5c" + suffix)
        rhofile = head + ".rho"
        x +=1
        if os.path.exists(rhofile) and os.path.getsize(rhofile) > 0:
            os.rename(rhofile, rhofile + suffix)
        else:
            print rhofile, "does not exist or empty???"
            sys.exit(1)
				
				
    alat, angle = get_lattice(head + ".struct")
    output = open(head + ".rho3d" , "w")
    output.write("cell\n")
    output.write("%f  %f  %f\n" % (alat[1], alat[2], alat[0]))
    output.write("%f  %f  %f\n" % ( angle[1], angle[2], angle[0]))
    output.write("%d  %d  %d " % ( ny, nz,nx))
    output.write("%f  %f  %f\n" % ( alat[1], alat[2], alat[0]))

    #rho = collect_charge(head, nz/2)
  #  i = 0
  #  for z in rho:
  #      output.write("%15.8e " % z)
  #      i = i + 1
  #      if (i % 5) == 0:
  #          output.write("\n")

  #  output.close()

    #if not preserve:
    #    erase_files(head, nz-1)