#!/usr/bin/env python

__author__ = "Jin Cao"
__copyright__ = "Copyright 2016, qfmda.com"
__version__ = "0.1"
__maintainer__ = "Jin Cao"
__email__ = "jincao2013@outlook.com"
__date__ = "Dec 2, 2016"

import re
from math import cos
from math import sin
from math import sqrt
import copy


def POSCAR_to_crystal_data(file_name):
# This funtion translate POSCAR file to a python list named crystal_data.
# caution_1: ensure a file named 'POSCAR' ,not 'COUNTCAR'
# caution_2: ensure c is the normal diraction of 2D_sheet,moreover,it should 
#           rolls around 'a' axil of 2D_sheet under the rule of right hand.
    file_name = file_name    
    POSCAR = []
    crystal_data = []
    
    file_path  = './'+file_name
    with open(file_path, 'r') as f:
        for line in f.readlines():
            # print line.strip()
            POSCAR.append(line.strip())
    
    # line_1: New structure
    crystal_data.append(POSCAR[0])
    # line_2: 1.0
    crystal_data.append(POSCAR[1])
    # line_3,4,5: lattice constant
    for i in [2,3,4]:
        a = re.findall(r"\d+\.?\d*",POSCAR[i])
        b = []
        b.append(float(a[0]))
        b.append(float(a[1]))
        b.append(float(a[2]))
        crystal_data.append(b)
    # line_6: atom labels
    atom_labels = re.findall(r"\w+",POSCAR[5])
    crystal_data.append(atom_labels)
    # line_7: atom_numbers    
    atom_numbers_str = re.findall(r"\d+",POSCAR[6])
    atom_numbers = []
    for i in range(len(atom_numbers_str)):
        atom_numbers.append(int(atom_numbers_str[i]))
    crystal_data.append(atom_numbers)
    # line_8: Direct or Cartesian
    crystal_data.append(POSCAR[7])
    
    range_POSCAR=range(len(POSCAR))
    
    for i in [0,1,2,3,4,5,6,7]:
        del range_POSCAR[0]
    # lines_9+: positons of atom
    for i in range_POSCAR:
        a = re.findall(r"\d+\.?\d*",POSCAR[i])
        b = []
        b.append(float(a[0]))
        b.append(float(a[1]))
        b.append(float(a[2]))
        crystal_data.append(b)
    # print crystal_data
    return crystal_data



def crystal_data_to_POSCAR(crystal_data,output_file_name):
# This funtion used to creat POSCAR from crystal_data.
# The only paramate of input of this funtion is crystal_data that is 
# automatically generated by POSCAR_to_crystal_data or make_NT funtions.
    output_file_name = output_file_name
    output_file_path = './' + output_file_name   
    
    with open(output_file_path, 'w') as g:
        # line_1: New structure
        # + line_2: 1.0
        for i in [0,1]:
            g.write(str(crystal_data[i]))
            g.write('\n')
        # line_3,4,5: lattice constant
        for i in [2,3,4]:
            g.write('\t')
            g.write(str(crystal_data[i][0]))
            g.write('\t\t\t')
            g.write(str(crystal_data[i][1]))
            g.write('\t\t\t')
            g.write(str(crystal_data[i][2]))
            g.write('\t\t\t')
            g.write('\n')
        # line_6: atom labels
        for i in range(len(crystal_data[5])):
            g.write('\t')
            g.write(str(crystal_data[5][i]))
        g.write('\n')
        # line_7: atom_numbers 
        for i in range(len(crystal_data[6])):
            g.write('\t')
            g.write(str(crystal_data[6][i]))
        g.write('\n')
        # line_8: Direct or Cartesian
        g.write(str(crystal_data[7]))
        g.write('\n')
            
        range_crystal_data=range(len(crystal_data))
        
        for i in [0,1,2,3,4,5,6,7]:
            del range_crystal_data[0]
        # lines_9+: positons of atom
        for i in range_crystal_data:
            g.write('\t')
            g.write(str(crystal_data[i][0]))
            g.write('\t\t')
            g.write(str(crystal_data[i][1]))
            g.write('\t\t')
            g.write(str(crystal_data[i][2]))
            g.write('\t\t')
            g.write('\n')


def Cartesian_to_Direct(crystal_data):
    # input
    crystal_data = crystal_data
    
    # prepare process
    crystal_data_Direct = crystal_data
    
    range_crystal_data=range(len(crystal_data_Direct))
    
    for i in [0,1,2,3,4,5,6,7]:
        del range_crystal_data[0]
    
    a = crystal_data_Direct[2][0]
    b = crystal_data_Direct[3][1]
    c = crystal_data_Direct[4][2]
    # main
    if crystal_data_Direct[7] == 'Cartesian':
        crystal_data_Direct[7] = 'Direct'
        for i in range_crystal_data:
            crystal_data_Direct[i][0] = crystal_data_Direct[i][0]/a
            crystal_data_Direct[i][1] = crystal_data_Direct[i][1]/b
            crystal_data_Direct[i][2] = crystal_data_Direct[i][2]/c
        return crystal_data_Direct
    if crystal_data_Direct[7] == 'Direct':
        return crystal_data_Direct

def Direct_to_Cartesian(crystal_data):
    # input
    crystal_data = crystal_data
    
    # prepare process
    crystal_data_Cartesian = crystal_data
    
    range_crystal_data=range(len(crystal_data_Cartesian))
    
    for i in [0,1,2,3,4,5,6,7]:
        del range_crystal_data[0]
    
    a = crystal_data_Cartesian[2][0]
    b = crystal_data_Cartesian[3][1]
    c = crystal_data_Cartesian[4][2]
    # main
    if crystal_data_Cartesian[7] == 'Direct':
        crystal_data_Cartesian[7] = 'Cartesian'
        for i in range_crystal_data:
            crystal_data_Cartesian[i][0] = crystal_data_Cartesian[i][0]*a
            crystal_data_Cartesian[i][1] = crystal_data_Cartesian[i][1]*b
            crystal_data_Cartesian[i][2] = crystal_data_Cartesian[i][2]*c
        return crystal_data_Cartesian
    if crystal_data_Cartesian[7] == 'Cartesian':
        return crystal_data_Cartesian



def build_supercell(A,B,C,_crystal_data):
    # input
    A = int(A)
    B = int(B)
    C = int(C)

    # preper process
    if _crystal_data[7] == 'Cartesian':
        crystal_data = copy.deepcopy(_crystal_data)
        crystal_data_supercell = _crystal_data
    if _crystal_data[7] == 'Direct':
        crystal_data_supercell = Direct_to_Cartesian(_crystal_data) 
        crystal_data = copy.deepcopy(Direct_to_Cartesian(_crystal_data))
    # main
    lattice_a = sqrt((crystal_data[2][0])**2+(crystal_data[2][1])**2+(crystal_data[2][2])**2)
    lattice_b = sqrt((crystal_data[3][0])**2+(crystal_data[3][1])**2+(crystal_data[3][2])**2)
    lattice_c = sqrt((crystal_data[4][0])**2+(crystal_data[4][1])**2+(crystal_data[4][2])**2)


    
    # build lattice constant
    for i in [2,3,4]:
        crystal_data_supercell[i][0] = A*crystal_data_supercell[i][0]
        crystal_data_supercell[i][1] = B*crystal_data_supercell[i][1]
        crystal_data_supercell[i][2] = C*crystal_data_supercell[i][2]
    # build atom labels and numbers
    for i in range(len(crystal_data_supercell[6])):
        crystal_data_supercell[6][i] = A*B*C*crystal_data_supercell[6][i]
    # build atoms
    count = 9
    insert_point = 9
    for n in range(len(crystal_data[6])):
        insert_point = insert_point + crystal_data[6][n]
        for m in range(crystal_data[6][n]):
            for i in range(A):
                for j in range(B):
                    for k in range(C):
                        if (i==0)and(j==0)and(k==0):
                            continue
                        new_atom = []
                        new_atom.append(i*lattice_a+crystal_data[count+m-1][0])
                        new_atom.append(j*lattice_b+crystal_data[count+m-1][1])
                        new_atom.append(k*lattice_c+crystal_data[count+m-1][2])
                        crystal_data_supercell.insert(insert_point-1,new_atom)
                        insert_point = insert_point + 1
        count = count + crystal_data[6][n]
    
    return crystal_data_supercell


def make_NT(crystal_data,vacuum):
    # input
    vacuum = vacuum
    crystal_data = crystal_data
    
    # sub program

    ## prepare process
    if crystal_data[7] == 'Cartesian':
        crystal_data = crystal_data
    if crystal_data[7] == 'Direct':
        crystal_data = Direct_to_Cartesian(crystal_data)

    crystal_data_NT = crystal_data
    
    range_crystal_data=range(len(crystal_data))
    
    for i in [0,1,2,3,4,5,6,7]:
        del range_crystal_data[0]
    
    ## find Z max, Z min
    z_list = []
    for i in range_crystal_data:
        z_list.append(float(crystal_data[i][2]))
    
    z_max = max(z_list)
    z_min = min(z_list)
    
    ## paramater of NT
    L = crystal_data[3][1]      # equal to b of 2D_sheet
    R = L/(2*3.1415926)
    a = 2*R + vacuum + (z_max - z_min)*2
    b = 2*R + vacuum + (z_max - z_min)*2
    c = crystal_data[2][0]
    
    ## write into crystal_data_NT
    crystal_data_NT[2][0] = a
    crystal_data_NT[2][1] = 0
    crystal_data_NT[2][2] = 0
    
    crystal_data_NT[3][0] = 0
    crystal_data_NT[3][1] = b
    crystal_data_NT[3][2] = 0
    
    crystal_data_NT[4][0] = 0
    crystal_data_NT[4][1] = 0
    crystal_data_NT[4][2] = c
    
    for i in range_crystal_data:
        # data of sheet
        x_sheet = crystal_data[i][0]
        y_sheet = crystal_data[i][1]
        z_sheet = crystal_data[i][2]
        # translate to NT
        R_this_layer = R+((z_max + z_min)/2 - z_sheet)
        x_NT = R_this_layer*cos(y_sheet/R) + a/2
        y_NT = R_this_layer*sin(y_sheet/R) + b/2
        z_NT = x_sheet
        # write to crystal_data
        crystal_data_NT[i][0] = x_NT
        crystal_data_NT[i][1] = y_NT
        crystal_data_NT[i][2] = z_NT
        
    return crystal_data_NT



'''
print '\n\n\n'
print '########################################################################'
print '################ Welcome to making Nanotube(NT) program ################'
print '########################################################################'
print '\n'
print 'This program generate POSCAR_NT from input file POSCAR with '
print 'a only paramater -vacuum-.To generate a resonable result you expected,'
print 'ensuring the follwing items.'
print '##Items:'
print "1. ensuring a input file named 'POSCAR' ,not a such like 'POSCAR.vasp' or" 
print "   'COUNTCAR'."
print "2. ensuring 'c' is the normal diraction of 2D_sheet,moreover,it should" 
print "   rolls around 'a' axil of 2D_sheet under the rule of right hand."
print "3. Cartesian is recommanded than Direct"
print '4. atoms should formate a 2D_sheet,but not via periodic bountary '
print '   condition, i.e. atoms should not cross the bountary, if you get a'
print '   result with a strangely large R, it may cause from this.'

vacuum = input('vacuum = ')
A = input('A = ')
B = input('B = ')
#A = 1
C = 1
# vacuum = 15
crystal_data = POSCAR_to_crystal_data('POSCAR')
crystal_data_supercell =build_supercell(A,B,C,crystal_data)
#crystal_data_to_POSCAR(crystal_data_supercell)
crystal_data_NT = make_NT(crystal_data_supercell,vacuum)
crystal_data_to_POSCAR(crystal_data_NT,'POSCAR_NT')
'''


# the following lines of code is for test
'''
crystal_data = POSCAR_to_crystal_data()
crystal_data_Cartesian = Direct_to_Cartesian(crystal_data)
# crystal_data_Direct = Cartesian_to_Direct(crystal_data)
crystal_data_to_POSCAR(crystal_data_Cartesian)
'''


## test build_supercell
'''
crystal_data = POSCAR_to_crystal_data()
crystal_data_super = build_supercell(2,2,2,crystal_data)
crystal_data_to_POSCAR(crystal_data_super)
'''





