#!/usr/bin/env python
#
# displace.py
#
# Simple script to generate input files of given displacement patterns.
# Currently, VASP, Quantum-ESPRESSO, and xTAPP are supported.
#
# Copyright (c) 2014, 2015, 2016 Terumasa Tadano
#
# This file is distributed under the terms of the MIT license.
# Please see the file 'LICENCE.txt' in the root directory
# or http://opensource.org/licenses/mit-license.php for information.
#

"""
Input file generator for displaced configurations.
"""

import optparse
import numpy as np

usage = "usage: %prog [options] file.pattern_HARMONIC file.pattern_ANHARM3 ... \n \
      file.pattern_* can be generated by 'alm' with MODE = suggest."
parser = optparse.OptionParser(usage=usage)
parser.add_option('--mag',
                  help="Magnitude of displacement in units of \
                        Angstrom (default: 0.02)")

parser.add_option('--prefix',
                  help="Prefix of the files to be created. ")

parser.add_option('--QE',
                  metavar='orig.pw.in',
                  help="Quantum-ESPRESSO input file with equilibrium atomic positions (default: None)")
parser.add_option('--VASP',
                  metavar='orig.POSCAR',
                  help="VASP POSCAR file with equilibrium atomic \
                        positions (default: None)")
parser.add_option('--xTAPP',
                  metavar='orig.cg',
                  help="xTAPP CG file with equilibrium atomic \
                        positions (default: None)")

# Functions for VASP


def read_POSCAR(file_in):

    file_pos = open(file_in, 'r')

    file_pos.readline()
    a = float(file_pos.readline().rstrip())
    lavec = np.zeros((3, 3))

    for i in range(3):
        arr = file_pos.readline().rstrip().split()
        if len(arr) != 3:
            print "Could not read POSCAR properly"
            exit(1)

        for j in range(3):
            lavec[i, j] = a * float(arr[j])

    lavec = lavec.transpose()
    invlavec = np.linalg.inv(lavec)

    elements = file_pos.readline().rstrip().split()

    if elements[0].isdigit():
        nat_elem = [int(tmp) for tmp in elements]
        elements = []

    else:
        nat_elem = [int(tmp) for tmp in file_pos.readline().rstrip().split()]

    nat = np.sum(nat_elem)
    basis = file_pos.readline().rstrip()
    x = np.zeros((nat, 3))

    for i in range(nat):
        arr = file_pos.readline().rstrip().split()
        for j in range(3):
            x[i][j] = float(arr[j])

    if basis == "Direct" or basis == "direct" or basis == "D" or basis == "d":
        xf = x
    else:
        xf = np.dot(x, invlavec)

    file_pos.close()

    return lavec, invlavec, elements, nat_elem, xf


def write_POSCAR(prefix, counter, header, nzerofills,
                 lavec, elems, nat, disp, coord):

    filename = prefix + str(counter).zfill(nzerofills) + ".POSCAR"
    f = open(filename, 'w')
    f.write("%s\n" % header)
    f.write("%s\n" % "1.0")

    for i in range(3):
        f.write("%20.14f %20.14f %20.14f\n" % (lavec[0][i],
                                               lavec[1][i],
                                               lavec[2][i]))

    for i in range(len(elems)):
        f.write("%s " % elems[i])
    if len(elems) > 0:
        f.write("\n")

    for i in range(len(nat)):
        f.write("%d " % nat[i])
    f.write("\n")

    f.write("Direct\n")

    for i in range(len(disp)):
        for j in range(3):
            f.write("%20.16f" % (coord[i][j] + disp[i][j]))
        f.write("\n")
    f.close()


# Functions for QE

def get_namelist(file_in, namelist_tag):

    list_out = []
    flag_add = False

    with open(file_in) as openfileobject:
        for line in openfileobject:
            line_upper = line.upper()
            if namelist_tag in line_upper:
                flag_add = True
                list_out.append(line)
            elif line.strip() == "/":
                flag_add = False
            elif flag_add:
                list_out.append(line)

    if len(list_out) == 0:
        print "%s field not found" % namelist_tag
        exit(1)

    list_out.append("/\n")
    return list_out


def gen_lattice_vector(ibrav, celldm, list_CELL_PARAMETERS):
    """.

    Computer lattice vector in units of Angstrom for given ibrav and celldm.
    Doc/INPUT_PW.txt was used as a reference.
    """
    import math

    Bohr_to_angstrom = 0.5291772108

    lavec = np.zeros((3, 3))

    if ibrav == 0:

        if list_CELL_PARAMETERS is None:
            print "CELL_PARAMETERS must be given when ibrav = 0."
            exit(1)

        else:
            for i in range(3):
                lavec[i][:] = [float(entry) for entry in
                               list_CELL_PARAMETERS[i + 1].rstrip().split()]

            lavec = np.array(lavec)

            mode = list_CELL_PARAMETERS[0].rstrip().split()

            if len(mode) == 1:
                print "Error : Please specify either alat, bohr, or angstrom for CELL_PARAMETERS"
                exit(1)
            else:
                mode_str = mode[1].lower()

            if "alat" in mode_str:

                if not celldm[0]:
                    print "celldm(1) must be given when 'alat' is used for CELL_PARAMETERS"
                    exit(1)

                for i in range(3):
                    for j in range(3):
                        lavec[i][j] *= celldm[0]

            elif "angstrom" in mode_str:

                for i in range(3):
                    for j in range(3):
                        lavec[i][j] /= Bohr_to_angstrom

            elif "bohr" not in mode_str:

                print "Error : Invalid option for CELL_PARAMETERS: %s" % mode[1]
                exit(1)

    elif ibrav == 1:

        if not celldm[0]:
            print "celldm(1) must be given when ibrav = 1."
            exit(1)

        else:
            a = celldm[0]
            lavec = np.array([[a, 0.0, 0.0],
                              [0.0, a, 0.0],
                              [0.0, 0.0, a]])

    elif ibrav == 2:

        if not celldm[0]:
            print "celldm(1) must be given when ibrav = 2."
            exit(1)

        else:
            a = celldm[0] / 2.0
            lavec = np.array([[-a, 0.0, a],
                              [0.0, a, a],
                              [-a, a, 0.0]])

    elif ibrav == 3:

        if not celldm[0]:
            print "celldm(1) must be given when ibrav = 3."
            exit(1)

        else:
            a = celldm[0] / 2.0
            lavec = np.array([[a, a, a],
                              [-a, a, a],
                              [-a, -a, a]])

    elif ibrav == 4:

        if not celldm[0] or not celldm[2]:
            print "celldm(1) and celldm(3) must be given when ibrav = 4."
            exit(1)

        else:
            a = celldm[0]
            c = celldm[0] * celldm[2]
            lavec = np.array([[a, 0.0, 0.0],
                              [-0.5 * a, math.sqrt(3.) / 2.0 * a, 0.0],
                              [0.0, 0.0, c]])

    elif ibrav == 5 or ibrav == -5:

        if not celldm[0] or not celldm[3]:
            print "celldm(1) and celldm(4) must be given when ibrav = 5, -5."
            exit(1)

        else:
            a = celldm[0]
            cosalpha = celldm[3]
            tx = a * math.sqrt((1.0 - cosalpha) / 2.)
            ty = a * math.sqrt((1.0 - cosalpha) / 6.)
            tz = a * math.sqrt((1.0 + 2.0 * cosalpha) / 3.)

            if ibrav == 5:
                lavec = np.array([[tx, -ty, tz],
                                  [0.0, 2.0 * ty, tz],
                                  [-tx, -ty, tz]])

            else:
                a_prime = a / math.sqrt(3.0)
                u = tz - 2.0 * math.sqrt(2.0) * ty
                v = tz + math.sqrt(2.0) * ty

                u *= a_prime
                v *= a_prime

                lavec = np.array([[u, v, v],
                                  [v, u, v],
                                  [v, v, u]])

    elif ibrav == 6:

        if not celldm[0] or not celldm[2]:
            print "celldm(1) and celldm(3) must be given when ibrav = 6."
            exit(1)

        else:
            a = celldm[0]
            c = celldm[0] * celldm[2]
            lavec = np.array([[a, 0.0, 0.0],
                              [0.0, a, 0.0],
                              [0.0, 0.0, c]])

    elif ibrav == 7:

        if not celldm[0] or not celldm[2]:
            print "celldm(1) and celldm(3) must be given when ibrav = 7."
            exit(1)

        else:
            a = celldm[0]
            c = celldm[0] * celldm[2]
            lavec = np.array([[a / 2.0, -a / 2.0, c / 2.0],
                              [a / 2.0,  a / 2.0, c / 2.0],
                              [-a / 2.0, -a / 2.0, c / 2.0]])

    elif ibrav == 8:

        if not celldm[0] or not celldm[1] or not celldm[2]:
            print "celldm(1), celldm(2), and celldm(3) must be given\
             when ibrav = 8."
            exit(1)

        else:
            a = celldm[0]
            b = celldm[0] * celldm[1]
            c = celldm[0] * celldm[2]

            lavec = np.array([[a, 0.0, 0.0],
                              [0.0, b, 0.0],
                              [0.0, 0.0, c]])

    elif ibrav == 9 or ibrav == -9:

        if not celldm[0] or not celldm[1] or not celldm[2]:
            print "celldm(1), celldm(2), and celldm(3) must be given\
             when ibrav = 9 or -9."
            exit(1)

        else:
            a = celldm[0]
            b = celldm[0] * celldm[1]
            c = celldm[0] * celldm[2]

            if ibrav == 9:
                lavec = np.array([[a / 2., b / 2., 0.0],
                                  [-a / 2., b / 2., 0.0],
                                  [0.0, 0.0, c]])
            else:
                lavec = np.array([[a / 2., -b / 2., 0.0],
                                  [a / 2., b / 2., 0.0],
                                  [0.0, 0.0, c]])

    elif ibrav == 10:

        if not celldm[0] or not celldm[1] or not celldm[2]:
            print "celldm(1), celldm(2), and celldm(3) must be given\
             when ibrav = 10."
            exit(1)

        else:
            a = celldm[0] / 2.0
            b = celldm[0] * celldm[1] / 2.0
            c = celldm[0] * celldm[2] / 2.0
            lavec = np.array([[a, 0.0, c],
                              [a, b, 0.0],
                              [0.0, b, c]])

    elif ibrav == 11:

        if not celldm[0] or not celldm[1] or not celldm[2]:
            print "celldm(1), celldm(2), and celldm(3) must be given\
             when ibrav = 11."
            exit(1)

        else:
            a = celldm[0] / 2.0
            b = celldm[0] * celldm[1] / 2.0
            c = celldm[0] * celldm[2] / 2.0
            lavec = np.array([[a, b, c],
                              [-a, b, c],
                              [-a, -b, c]])

    elif ibrav == 12:

        if not celldm[0] or not celldm[1] or not celldm[2] or \
           not celldm[3]:
            print "celldm(1), celldm(2), celldm(3), and celldm(4)\
             must be given when ibrav = 12."
            exit(1)

        else:
            a = celldm[0]
            b = celldm[0] * celldm[1]
            c = celldm[0] * celldm[2]
            gamma = math.acos(celldm[3])
            lavec = np.array([[a, 0.0, 0.0],
                              [b * math.cos(gamma), b * math.sin(gamma), 0.0],
                              [0.0, 0.0, c]])

    elif ibrav == -12:

        if not celldm[0] or not celldm[1] or not celldm[2] or \
           not celldm[4]:
            print "celldm(1), celldm(2), celldm(3), and celldm(5)\
             must be given when ibrav = -12."
            exit(1)

        else:
            a = celldm[0]
            b = celldm[0] * celldm[1]
            c = celldm[0] * celldm[2]
            beta = math.acos(celldm[4])
            lavec = np.array([[a, 0.0, 0.0],
                              [0.0, b, 0.0],
                              [c * math.cos(beta), 0.0, c * math.sin(beta)]])

    elif ibrav == 13:

        if not celldm[0] or not celldm[1] or not celldm[2] or\
           not celldm[3]:
            print "celldm(1), celldm(2), celldm(3), and celldm(4)\
             must be given when ibrav = 13."
            exit(1)

        else:
            a = celldm[0]
            b = celldm[0] * celldm[1]
            c = celldm[0] * celldm[2]
            gamma = math.acos(celldm[3])
            lavec = np.array([[a / 2.0, 0.0, -c / 2.0],
                              [b * math.cos(gamma), b * math.sin(gamma), 0.0],
                              [a / 2.0, 0.0, c / 2.0]])

    elif ibrav == 14:

        if not celldm[0] or not celldm[1] or not celldm[2] or \
           not celldm[3] or not celldm[4] or not celldm[5]:
            print "All celldm must be given when ibrav = 14."
            exit(1)

        else:
            a = celldm[0]
            b = celldm[0] * celldm[1]
            c = celldm[0] * celldm[2]
            alpha = math.acos(celldm[3])
            beta = math.acos(celldm[4])
            gamma = math.acos(celldm[5])

            lavec = np.array([[a, 0.0, 0.0],
                              [b * math.cos(gamma), b * math.sin(gamma), 0.0],
                              [c * math.cos(beta),
                               c * (math.cos(alpha) - math.cos(beta) *
                                    math.cos(gamma)) / math.sin(gamma),
                               c * math.sqrt(1.0 + 2.0 * math.cos(alpha) * math.cos(beta) * math.cos(gamma)
                                             - math.cos(alpha) ** 2 - math.cos(beta) ** 2 - math.cos(gamma) ** 2) / math.sin(gamma)]])

    else:

        print "Invalid ibrav = %s" % ibrav
        exit(1)

    # Transpose for later use
    lavec = lavec.transpose()

    # Convert to Angstrom unit
    for i in range(3):
        for j in range(3):
            lavec[i][j] *= Bohr_to_angstrom

    return lavec


def get_system_info(list_in):

    list_mod = []

    for obj in list_in:
        obj_split = obj.rstrip().split(',')
        for subobj in obj_split:
            if subobj:
                index = subobj.find('=')
                if index > 0:
                    subobj = subobj[:index] + " = " + subobj[index + 1:]
                list_mod.append(subobj)

    str_input = ""

    for entry in list_mod:
        str_input += entry + " "

    entrylist = str_input.split()

    celldm = [[] for i in range(6)]

    for i in range(len(entrylist)):

        if "ibrav" in entrylist[i]:
            ibrav = int(entrylist[i + 2])

        if "nat" in entrylist[i]:
            nat = int(entrylist[i + 2])

        if "ntyp" in entrylist[i]:
            ntyp = int(entrylist[i + 2])

        if "celldm(1)" in entrylist[i]:
            celldm[0] = float(entrylist[i + 2])

        if "celldm(2)" in entrylist[i]:
            celldm[1] = float(entrylist[i + 2])

        if "celldm(3)" in entrylist[i]:
            celldm[2] = float(entrylist[i + 2])

        if "celldm(4)" in entrylist[i]:
            celldm[3] = float(entrylist[i + 2])

        if "celldm(5)" in entrylist[i]:
            celldm[4] = float(entrylist[i + 2])

        if "celldm(6)" in entrylist[i]:
            celldm[5] = float(entrylist[i + 2])

    return ibrav, celldm, nat, ntyp


def get_options(option_tag, taglists, file_in):

    list_out = []
    flag_add = False

    with open(file_in) as openfileobject:
        for line in openfileobject:

            if option_tag in line:
                flag_add = True
                list_out.append(line)
            elif len(line.split()) > 0 and line.split()[0] in taglists:
                flag_add = False
            elif flag_add:
                list_out.append(line)

    return list_out


def get_fractional_coordinate(aa, N, list_in, a_Bohr):

    Bohr_to_angstrom = 0.5291772108

    list_tmp = list_in[0].rstrip().split()

    if len(list_tmp) == 1:
        print "Error : Please specify either alat, bohr, angstrom, or crystal for ATOMIC_POSITIONS"
        exit(1)
    else:
        mode_str = list_tmp[1].lower()

    if "crystal_sg" in mode_str:
        print "Error : Sorry. 'crystal_sg' is not supported in this script. Please use another option."
        exit(1)

    xtmp = np.zeros((N, 3))
    kd = []

    for i in range(N):
        list_tmp = list_in[i + 1].rstrip().split()
        kd.append(list_tmp[0])
        xtmp[i][:] = [float(j) for j in list_tmp[1:4]]

    aa_inv = np.linalg.inv(aa)

    if "alat" in mode_str:
        a_angstrom = a_Bohr * Bohr_to_angstrom

        for i in range(3):
            for j in range(3):
                aa_inv[i][j] *= a_angstrom

        for i in range(N):
            xtmp[i][:] = np.dot(xtmp[i][:], aa_inv.transpose())

    elif "bohr" in mode_str:

        for i in range(3):
            for j in range(3):
                aa_inv[i][j] *= Bohr_to_angstrom

        for i in range(N):
            xtmp[i][:] = np.dot(xtmp[i][:], aa_inv.transpose())

    elif "angstrom" in mode_str:

        for i in range(N):
            xtmp[i][:] = np.dot(xtmp[i][:], aa_inv.transpose())

    elif "crystal" not in mode_str:
        print "Error : Invalid option for ATOMIC_POSITIONS: %s" % mode_str
        exit(1)

    return kd, xtmp


def read_original_QE(file_in):

    # Parse fortran namelists
    list_CONTROL = get_namelist(file_in, "&CONTROL")
    list_SYSTEM = get_namelist(file_in, "&SYSTEM")
    list_ELECTRONS = get_namelist(file_in, "&ELECTRONS")

    # Parse general options
    tags = ["ATOMIC_SPECIES", "ATOMIC_POSITIONS", "K_POINTS",
            "CELL_PARAMETERS", "OCCUPATIONS", "CONSTRAINTS", "ATOMIC_FORCES"]

    list_ATOMIC_SPECIES = get_options("ATOMIC_SPECIES", tags, file_in)
    list_ATOMIC_POSITIONS = get_options("ATOMIC_POSITIONS", tags, file_in)
    list_K_POINTS = get_options("K_POINTS", tags, file_in)
    list_CELL_PARAMETERS = get_options("CELL_PARAMETERS", tags, file_in)
    list_OCCUPATIONS = get_options("OCCUPATIONS", tags, file_in)

    # Get ibrav, celldm, nat, and ntyp
    # and then calculate the lattice vector
    ibrav, celldm, nat, ntyp = get_system_info(list_SYSTEM)
    lavec = gen_lattice_vector(ibrav, celldm, list_CELL_PARAMETERS)
    lavec_inv = np.linalg.inv(lavec)

    # Get fractional coordinate
    kd_symbol, x_frac = get_fractional_coordinate(lavec,
                                                  nat,
                                                  list_ATOMIC_POSITIONS,
                                                  celldm[0])
    list_namelist_merged = []
    list_namelist_merged.extend(list_CONTROL)
    list_namelist_merged.extend(list_SYSTEM)
    list_namelist_merged.extend(list_ELECTRONS)

    return list_namelist_merged, list_ATOMIC_SPECIES, list_K_POINTS, \
        list_CELL_PARAMETERS, list_OCCUPATIONS, \
        nat, lavec, kd_symbol, x_frac, lavec_inv


def generate_QE_input(prefix, suffix, counter, nzerofills, list_namelist,
                      list_ATOMS, list_KP, list_CELL, list_OCCU,
                      nat, kd_symbol, x, u):

    filename = prefix + str(counter).zfill(nzerofills) + "." + suffix
    f = open(filename, 'w')

    for entry in list_namelist:
        f.write(entry)

    for entry in list_ATOMS:
        f.write(entry)

    f.write("ATOMIC_POSITIONS crystal\n")
    for i in range(nat):
        f.write("%s %20.16f %20.16f %20.16f\n" % (kd_symbol[i],
                                                  x[i][0] + u[i, 0],
                                                  x[i][1] + u[i, 1],
                                                  x[i][2] + u[i, 2]))

    for entry in list_KP:
        f.write(entry)
    for entry in list_CELL:
        f.write(entry)
    for entry in list_OCCU:
        f.write(entry)

    f.write("\n")
    f.close()


# Functions for xTAPP

def read_tappinput(file_in):

    list_tappinput = []
    flag_add = False

    with open(file_in) as openfileobject:
        for line in openfileobject:
            if "main" in line and "data" in line:
                flag_add = True
                list_tappinput.append(line)
            elif "#" in line:
                flag_add = False
            elif flag_add:
                list_tappinput.append(line)

    if len(list_tappinput) == 0:
        print "main data entry not found"
        exit(1)

    list_tappinput_new = []

    for obj in list_tappinput:
        obj_split = obj.rstrip().split(',')
        for subobj in obj_split:
            if subobj:
                list_tappinput_new.append(subobj)

    str_input = ""

    for entry in list_tappinput_new:
        str_input += entry + " "

    entrylist = str_input.split()
    lavec_list = []

    a = 0.0
    nkd = 0
    nat = 0

    # get lattice_factor
    for i in range(len(entrylist)):
        if "lattice_factor" in entrylist[i]:
            a = float(entrylist[i + 2])

        if "lattice_list" in entrylist[i]:
            for j in range(9):
                lavec_list.append(entrylist[i + j + 2])

        if "number_element" in entrylist[i]:
            nkd = int(entrylist[i + 2])

        if "number_atom" in entrylist[i]:
            nat = int(entrylist[i + 2])

    if a == 0.0:
        print "Couldn't read lattice_factor"
        exit(1)
    if nkd == 0:
        print "Couldn't read number_element"
        exit(1)
    if nat == 0:
        print "Couldn't read number_atom"
        exit(1)
    if len(lavec_list) != 9:
        print "Couldn't read lattice_list"
        exit(1)

    lavec = np.zeros((3, 3))

    Bohr_to_angstrom = 0.5291772108
    a *= Bohr_to_angstrom

    for i in range(3):
        for j in range(3):
            lavec[j][i] = a * float(lavec_list[3 * i + j])

    return lavec, nat, nkd, list_tappinput


def read_kpdata(file_in):

    list_kpoint = []
    flag_add = False

    with open(file_in) as openfileobject:
        for line in openfileobject:
            if "k-points" in line.rstrip():
                flag_add = True
                list_kpoint.append(line)
            elif "#" in line.strip():
                flag_add = False
            elif flag_add:
                list_kpoint.append(line)

    if len(list_kpoint) == 0:
        print "k-points data entry not found"
        exit(1)

    return list_kpoint


def read_structure_optimize(file_in):

    list_opt = []
    flag_add = False

    with open(file_in) as openfileobject:
        for line in openfileobject:
            if "struct_opt" in line.rstrip():
                flag_add = True
                list_opt.append(line)
            elif "#" in line.strip():
                flag_add = False
            elif flag_add:
                list_opt.append(line)

    if len(list_opt) == 0:
        print "struct_opt entry not found"
        exit(1)

    list_opt2 = []
    flag_add = False

    with open(file_in) as openfileobject:
        for line in openfileobject:
            if "str_opt_constr" in line.rstrip():
                flag_add = True
                list_opt2.append(line)
            elif "#" in line.strip():
                flag_add = False
            elif flag_add:
                list_opt2.append(line)

    if len(list_opt2) == 0:
        print "str_opt_constr entry not found"
        exit(1)

    return list_opt, list_opt2


def read_atomdata(file_in, nat_in, nkd_in):

    list_atom = []
    flag_add = False

    with open(file_in) as openfileobject:
        for line in openfileobject:
            if "atom" in line and "data" in line:
                flag_add = True
                list_atom.append(line)
            elif "#" in line.strip():
                flag_add = False
            elif flag_add:
                list_atom.append(line)

    if len(list_atom) == 0:
        print "atom data entry not found"
        exit(1)

    x_out = np.zeros((nat_in, 3), dtype=float)
    kd_out = np.zeros(nat_in, dtype=int)

    for i in range(nat_in):
        list_tmp = list_atom[i + nkd_in + 1].rstrip().split()
        kd_out[i] = int(list_tmp[0])
        for j in range(3):
            x_out[i][j] = float(list_tmp[j + 1])

    return x_out, kd_out, list_atom


def read_CG(file_in):

    lavec, nat, nkd, str_tappinput = read_tappinput(file_in)
    str_kpoint = read_kpdata(file_in)
    str_struct_opt, str_opt_constr = read_structure_optimize(file_in)
    x, kd, str_atom = read_atomdata(file_in, nat, nkd)

    str_header = ""

    for entry in str_tappinput:
        str_header += entry
    for entry in str_kpoint:
        str_header += entry
    for entry in str_struct_opt:
        str_header += entry
    for entry in str_opt_constr:
        str_header += entry
    for i in range(nkd + 1):
        str_header += str_atom[i]

    lavec = np.matrix(lavec)
    lavec_inv = np.array(lavec.I)

    return str_header, nat, nkd, lavec, lavec_inv, x, kd


def gen_CG(prefix, suffix, counter, nzerofills, str_header,
           nat, kd, x, u, nsym, symop, denom_tran, has_inv):

    filename = prefix + str(counter).zfill(nzerofills) + "." + suffix
    f = open(filename, 'w')
    f.write("%s" % str_header)

    for i in range(nat):
        f.write("%i %20.16f %20.16f %20.16f\n" % (kd[i], x[i][0] + u[i, 0],
                                                  x[i][1] + u[i, 1],
                                                  x[i][2] + u[i, 2]))

    f.write("# symmetry data\n")
    f.write("&symmetry\n")
    f.write("  number_sym_op = %i\n" % nsym)
    f.write("  has_inversion = %i\n" % has_inv)
    f.write("  denom_trans = %i\n" % denom_tran)
    f.write("/\n")

    mat_tmp = np.zeros((3, 3), dtype=int)

    for elems in symop:
        for i in range(3):
            for j in range(3):
                mat_tmp[i][j] = elems[3 * i + j]

        mat_inv = np.matrix(mat_tmp).I

        for i in range(3):
            for j in range(3):
                f.write("%4i" % mat_inv[i, j])

        f.write("   ")
        for i in range(3):
            f.write("%4i" % elems[9 + i])

        f.write("\n")

    f.write("\n")
    f.close()


# Other functions

def parse_displacement_patterns(files_in):

    pattern = []

    for file in files_in:
        pattern_tmp = []

        f = open(file, 'r')
        tmp, basis = f.readline().rstrip().split(':')
        if basis == 'F':
            print "Warning: DBASIS must be 'C'"
            exit(1)

        while True:
            line = f.readline()

            if not line:
                break

            line_split_by_colon = line.rstrip().split(':')
            is_entry = len(line_split_by_colon) == 2

            if is_entry:
                pattern_set = []
                natom_move = int(line_split_by_colon[1])
                for i in range(natom_move):
                    disp = []
                    line = f.readline()
                    line_split = line.rstrip().split()
                    disp.append(int(line_split[0]))
                    for j in range(3):
                        disp.append(float(line_split[j + 1]))

                    pattern_set.append(disp)
                pattern_tmp.append(pattern_set)

        print "File %s containts %i displacement patterns" \
              % (file, len(pattern_tmp))

        for entry in pattern_tmp:
            if entry not in pattern:
                pattern.append(entry)

        f.close()

        print
    print "Number of unique displacement patterns = ", len(pattern)

    return pattern


def char_xyz(entry):

    if entry % 3 == 0:
        return 'x'
    elif entry % 3 == 1:
        return 'y'
    elif entry % 3 == 2:
        return 'z'


def gen_displacement(counter_in, pattern, disp_mag, nat, invlavec):

    poscar_header = "Disp. Num. %i" % counter_in
    poscar_header += " ( %f Angstrom" % disp_mag

    disp = np.zeros((nat, 3))

    for displace in pattern:
        atom = displace[0] - 1

        poscar_header += ", %i : " % displace[0]

        str_direction = ""

        for i in range(3):
            if abs(displace[i + 1]) > 1.0e-10:
                if displace[i + 1] > 0.0:
                    str_direction += "+" + char_xyz(i)
                else:
                    str_direction += "-" + char_xyz(i)

            disp[atom][i] += displace[i + 1] * disp_mag

        poscar_header += str_direction

    poscar_header += ")"

    for i in range(nat):
        disp[i] = np.dot(disp[i], invlavec.T)

    return poscar_header, disp


def get_number_of_zerofill(npattern):

    nzero = 1

    while True:
        npattern /= 10

        if npattern == 0:
            break

        nzero += 1

    return nzero


if __name__ == '__main__':

    options, args = parser.parse_args()
    file_pattern = args[0:]

    print "*****************************************************************"
    print "             displace.py -- Input file generator                 "
    print "*****************************************************************"
    print

    if len(file_pattern) == 0:
        print "Usage: displace.py [options] file1.pattern_HARMONIC\
 file2.pattern_ANHARM3 ..."
        print "file.pattern_* can be generated by 'alm' with MODE = suggest."
        print
        print "For details of available options, \
 please type\n$ python displace.py -h"
        exit(1)

    if options.VASP is None and options.QE is None and options.xTAPP is None:
        print "Error : Either --VASP, --QE, or --xTAPP option must be given."
        exit(1)

    elif options.VASP and options.QE or options.VASP and options.xTAPP or \
            options.QE and options.xTAPP:
        print "Error : --VASP, --QE, and --xTAPP cannot be given simultaneously."
        exit(1)

    elif options.VASP:
        code = "VASP"
        print "--VASP option is given: Generate POSCAR files for VASP"
        print

    elif options.QE:
        code = "QE"
        print "--QE option is given: Generate input files for Quantum-ESPRESSO."
        print

    elif options.xTAPP:
        code = "xTAPP"
        print "--xTAPP option is given: Generate input files for xTAPP."
        print

    # Assign the magnitude of displacements
    if options.mag is None:
        options.mag = "0.02"
        disp_length = 0.02
        print "--mag option not given. Substituted by the default (0.02 Angstrom)"
        print

    else:
        disp_length = float(options.mag)

    if options.prefix is None:
        prefix = "disp"
        print "--prefix option not given. Substituted by the default (\"disp\"). "
        print
    else:
        prefix = options.prefix

    print "-----------------------------------------------------------------"
    print

    if code == "VASP":
        str_outfiles = "%s{counter}.POSCAR" % prefix
        file_original = options.VASP

    elif code == "QE":
        str_outfiles = "%s{counter}.pw.in" % prefix
        file_original = options.QE
        suffix = "pw.in"

    elif code == "xTAPP":
        str_outfiles = "%s{counter}.cg" % prefix
        file_original = options.xTAPP

    # Read the original file
    if code == "VASP":
        aa, aa_inv, elems, nats, x_frac = read_POSCAR(file_original)
        nat = np.sum(nats)

    elif code == "QE":
        list_namelist, list_ATOMIC_SPECIES, \
            list_K_POINTS, list_CELL_PARAMETERS, list_OCCUPATIONS, \
            nat, lavec, kd_symbol, x_frac, aa_inv = read_original_QE(
                file_original)

    elif code == "xTAPP":
        str_header, nat, nkd, aa, aa_inv, x_frac, kd = read_CG(file_original)
        suffix = "cg"

    print "Original file                  : %s" % file_original
    print "Output file format             : %s" % str_outfiles
    print "Magnitude of displacements     : %s Angstrom" % disp_length
    print "Number of atoms                : %i" % nat
    print

    disp_pattern = parse_displacement_patterns(args[:])
    nzerofills = get_number_of_zerofill(len(disp_pattern))
    counter = 0

    for pattern in disp_pattern:
        counter += 1
        header, disp = gen_displacement(counter, pattern, disp_length,
                                        nat, aa_inv)

        if code == "VASP":
            write_POSCAR(prefix, counter, header, nzerofills,
                         aa, elems, nats, disp, x_frac)

        elif code == "QE":
            generate_QE_input(prefix, suffix, counter, nzerofills, list_namelist,
                              list_ATOMIC_SPECIES, list_K_POINTS,
                              list_CELL_PARAMETERS, list_OCCUPATIONS,
                              nat, kd_symbol, x_frac, disp)
        elif code == "xTAPP":
            nsym = 1
            symop = []
            symop.append([1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0])
            denom_tran = 1
            has_inv = 0

            gen_CG(prefix, suffix, counter, nzerofills, str_header, nat, kd,
                   x_frac, disp, nsym, symop, denom_tran, has_inv)

    print
    print "All input files are created."
