#!/usr/bin/python3.6

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

'''
program pos2xyzAnim_for.py
	- converts file in XDATCAR format into the xyz formal that could be read by most of 
	visualization programs

	- example of usage:

	pos2xyzAnim_for.py m/XDATCAR 10 0 0.2 0 > anim.xyz

	in this case, every 10th configuration is selected and fractional coordinates are
	shifted by the vector (0.0, 0.2, 0.0)
'''

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


fname1 = sys.argv[1]

pace = 1

shift = zeros(3, float)
# print len(sys.argv)
# print xaf

if len(sys.argv) > 2:
    pace = int(sys.argv[2])


if len(sys.argv) > 5:
    shift[0] = float(sys.argv[3])
    shift[1] = float(sys.argv[4])
    shift[2] = float(sys.argv[5])

mult = ones(3, int)

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

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

xdat = abstracts.Xdatcarreader()
xdat.readFile(fname1)
numofatoms = xdat.numofatoms
# katoms=xdat.atoms
types = xdat.types
ntypes = xdat.ntypes

flags = xdat.tags_atoms


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

#c select every pace-th structure
coordsall = []

for i in range(0,len(xdat.data),pace):
    coordsall.append(xdat.data[i])

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

# put in the same cell
for i in range(len(coordsall)):
    if i > 0:
        coordsall[i] = same_cell(array(coordsall[i]), array(coordsall[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(coordsall)):
    #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 = coordsall[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])
