#!/usr/bin/python3

from numpy import *
import string
import int_values
import abstracts
import sys

'''
program pos2xyzStatic.py
	- converts POSCAR into xyz file format

	- example of usage:

	python3 pos2xyzStatic.py POSCAR.1 0 0.2 0 > anim.xyz

	in this case, fractional coordinates from POSCAR are shifted by the vector 
	(0.0, 0.2, 0.0) before converting it to Cartesian coordinates. The result is 
	written in file anim.xyz that can be visualized, e.g, by molden, or jmol
'''

def same_cell(x, y):
    r = 0.0
    for i in range(len(x)):
        for j in range(len(x[0])):
            while (x[i][j]-y[i][j]) > 0.5:
                x[i][j] -= 1
            while (x[i][j]-y[i][j]) <= -0.5:
                x[i][j] += 1
    return x


fname = sys.argv[1]
shift = zeros(3, float)
# print len(sys.argv)
# print xaf
if len(sys.argv) > 4:
    shift[0] = float(sys.argv[2])
    shift[1] = float(sys.argv[3])
    shift[2] = float(sys.argv[4])

mult = ones(3, int)

if len(sys.argv) > 7:
    mult[0] = int(sys.argv[5])
    mult[1] = int(sys.argv[6])
    mult[2] = int(sys.argv[7])

mult0 = mult[0]*mult[1]*mult[2]

xdat = abstracts.Xdatcarreader()
xdat.readFile(fname, 0)
numofatoms = xdat.numofatoms
# katoms=xdat.atoms
types = xdat.types
ntypes = xdat.ntypes

flags = xdat.tags_atoms

# print len(katoms),"species detected"
#flags=raw_input('give me atomic flags:\n')
# flags=flags.split()

allflags = []
for i in range(len(flags)):
    allflags = allflags+types[i]*[flags[i]]


# user-defined shift
for i in range(len(xdat.data)):
    xdat.data[i] += shift
    for j in range(numofatoms):
        for k in range(3):
            while xdat.data[i][j][k] >= 1.:
                xdat.data[i][j][k] -= 1.
            while xdat.data[i][j][k] < 0.0:
                xdat.data[i][j][k] += 1.

# put in the same cell
for i in range(len(xdat.data)):
    if i > 0:
        xdat.data[i] = same_cell(array(xdat.data[i]), array(xdat.data[i-1]))


lattmat = xdat.lattmat[0]

lattmat[0] *= mult[0]
lattmat[1] *= mult[1]
lattmat[2] *= mult[2]

for i in range(ntypes):
    types[i] *= mult0

for i in range(len(xdat.data)):
    # if i>0:
    #  xdat.data[i]=same_cell(array(xdat.data[i]),array(xdat.data[i-1]))
    xnew = zeros((mult0*numofatoms, 3), float)
    count = 0
    for j in range(numofatoms):
        x0 = xdat.data[i][j]
        for t1 in range(mult[0]):
            for t2 in range(mult[1]):
                for t3 in range(mult[2]):
                    xnew[count][0] = (x0[0]+1.*t1)/mult[0]
                    xnew[count][1] = (x0[1]+1.*t2)/mult[1]
                    xnew[count][2] = (x0[2]+1.*t3)/mult[2]
                    count += 1
    coords_c = dot(xnew, lattmat)
    print(numofatoms*mult0)
    print("")
    for j in range(numofatoms*mult0):
        indx = int(j/mult0)
        print(allflags[indx], coords_c[j, 0], coords_c[j, 1], coords_c[j, 2])
