#!/usr/bin/env python2
################################################################
#
# vasp2force:
#   convert vasp output data into potfit reference configurations
#
################################################################
#
#   Copyright 2014
#             Institute for Theoretical and Applied Physics
#             University of Stuttgart, D-70550 Stuttgart, Germany
#             http://potfit.sourceforge.net/
#
#################################################################
#
#   This file is part of potfit.
#
#   potfit is free software; you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation; either version 2 of the License, or
#   (at your option) any later version.
#
#   potfit is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with potfit; if not, see <http://www.gnu.org/licenses/>.
#
#################################################################

import argparse
import copy
import gzip
import os
import sys


def get_outcar_files(directory, recursive):
    # walk directory (recursively) and return all OUTCAR* files
    # return list of outcars' path
    sys.stderr.write(
        'Searching directory %s for OUTCAR* files ...\n' % directory)
    outcars = []
    if not recursive:
        for item in os.listdir(directory):
            if item.startswith('OUTCAR'):
                outcars.append(os.path.join(directory, item))
    else:
        for root, SubFolders, files in os.walk(directory):
            for item in files:
                if item.startswith('OUTCAR'):
                    outcars.append(os.path.join(root, item))
    if len(outcars) == 0:
        sys.stderr.write(
            'Could not find any OUTCAR files in this directory.\n')
    else:
        sys.stderr.write('Found the following files:\n')
        sys.stderr.write('  {}\n'.format(('\n  ').join(outcars)))
        return outcars
    return outcars


def uniq(seq):
    # return unique list without changing the order
    # from http://stackoverflow.com/questions/480214
    seen = set()
    seen_add = seen.add
    return [x for x in seq if x not in seen and not seen_add(x)]


def scan_outcar_file(file_handle):
    # scans an outcar file and returns a list with
    # - number of configurations
    # - atom types

    # first try TOTAL-FORCE
    configs = 0
    atom_types = []
    title = []
    potcar = []
    ipt = []
    for line in file_handle:
        if line.startswith('|'):
            continue
        if 'TOTAL-FORCE' in line:
            configs += 1
        if 'VRHFIN' in line:
            atom_types.append(
                line.split()[1].replace('=', '').replace(':', ''))
        if 'title' in line:
            title.append(line.split()[3][0:2])
        if 'POTCAR' in line:
            potcar.append(line.split()[2][0:2])
        if 'ions per type' in line:
            ipt = [int(s) for s in line.split()[4:]]

    potcar = uniq(potcar)

    if atom_types:
        return [configs, atom_types, ipt]
    elif title:
        return [configs, title, ipt]
    elif potcar:
        return [configs, potcar, ipt]
    else:
        sys.stderr.write(
            'Could not determine atom types in file %s.\n' % filename)
        sys.exit()


def process_outcar_file_v5_dev(outcars, data, numbers, types, max_types, elements=None, windex=None, fout='potfit.configs'):
    fw = open(fout, 'w')
    for i in range(len(outcars)):
        if outcars[i].endswith('.gz'):
            f = gzip.open(outcars[i], 'rb')
        else:
            f = open(outcars[i], 'r')

        # Writing current OUTCAR's information into potfit format files.
        nconfs = data[i][0]
        natoms = sum(data[i][2])  # ipt
        if windex == None:
            windex = range(nconfs)
        if windex == 'final':
            windex = [nconfs - 1]

        # reading current OUTCAR
        print("Reading %s ..." % outcars[i])
        count = -1
        line = f.readline()
        while line != '':
            line = f.readline()
            if 'Iteration' in line:
                energy = 0
                box_x = []
                box_y = []
                box_z = []
                stress = []
                atom_data = []
            # if 'energy  without' in line:
            #     # appears in each electronic-iteration steps
            #     energy = float(line.split()[6]) / natoms
            if 'free  energy   TOTEN' in line:
                energy = float (line.split()[4]) / natoms
                if count in windex:
                    fw.write("#N %s 1\n" % natoms)
                    fw.write('#C ')
                    if elements:
                        fw.write("%s " % numbers[0])
                        for j in range(1, max_types):
                            fw.write('%s\t' % numbers[j])
                    else:
                        fw.write(" %s" % data[i][1][0])
                        for j in range(1, max_types):
                            fw.write(' %s' % data[i][1][j])
                    fw.write("\n")
                    fw.write("## force file generated from file %s config %d\n" % (
                        outcars[i], count))
                    fw.write("#X %13.8f %13.8f %13.8f\n" %
                             (box_x[0], box_x[1], box_x[2]))
                    fw.write("#Y %13.8f %13.8f %13.8f\n" %
                             (box_y[0], box_y[1], box_y[2]))
                    fw.write("#Z %13.8f %13.8f %13.8f\n" %
                             (box_z[0], box_z[1], box_z[2]))
                    fw.write("#W %f\n" % (args.weight))
                    fw.write("#E %.10f\n" % (energy))
                    if stress:
                        fw.write("#S ")
                        for num in range(6):
                            fw.write('%8.7g\t' % (stress[num]))
                        fw.write('\n')
                    fw.write("#F\n")
                    fw.flush()
                    for adata in atom_data:
                        fw.write("%d %11.7g %11.7g %11.7g %11.7g %11.7g %11.7g\n" %
                                 (adata[0], adata[1], adata[2], adata[3], adata[4], adata[5], adata[6]))
            if 'VOLUME and BASIS' in line:
                for do in range(5):
                    # SKIP 5 lines
                    line = f.readline()
                box_x = [float(s)
                         for s in line.replace('-', ' -').split()[0:3]]
                line = f.readline()
                box_y = [float(s)
                         for s in line.replace('-', ' -').split()[0:3]]
                line = f.readline()
                box_z = [float(s)
                         for s in line.replace('-', ' -').split()[0:3]]
            if 'in kB' in line:
                stress = [float(s) / 1602 for s in line.split()[2:8]]
            if 'TOTAL-FORCE' in line:
                # only appears in Ion-iteration steps
                line = f.readline()  # skip 1 line
                adata = [0] * 7
                for num in range(len(data[i][2])):
                    for k in range(data[i][2][num]):
                        line = [float(s) for s in f.readline().split()[0:6]]
                        if args.c:
                            adata[0] = int(types[data[i][1][num]])
                        else:
                            adata[0] = int(num)
                        adata[1] = float(line[0])
                        adata[2] = float(line[1])
                        adata[3] = float(line[2])
                        adata[4] = float(line[3])
                        adata[5] = float(line[4])
                        adata[6] = float(line[5])
                        atom_data.append(copy.copy(adata))
                count += 1
    fw.close()
    return


def Parser():
    parser = argparse.ArgumentParser(
        description='''Converts vasp output data into potfit reference configurations.''')

    parser.add_argument('-c', type=str, required=False,
                        help='list of chemical species to use, e.g. -c Mg=0,Zn=1')
    parser.add_argument(
        '-e', type=str, required=False, help='file with single atom energies (NYI)')
    parser.add_argument('-r', '--recursive', action='store_true',
                        help='scan recursively for OUTCAR files')
    parser.add_argument('-f', '--final', action='store_true',
                        help='use only the final configuration from OUTCAR')
    parser.add_argument('-sr', '--configs_range', type=str,
                        help='range of the configurations to use')
    parser.add_argument('-w', '--weight', type=float, default=1.0,
                        help='set configuration weight for all configurations')
    parser.add_argument(
        'files', type=str, nargs='*', help='list of OUTCAR files (plain or gzipped)')

    args = parser.parse_args()
    return args

#
if __name__ == "__main__":
    # Check for sane arguments
    args = Parser()
    if args.weight < 0:
        sys.stderr.write("The weight needs to be positive!\n")
        sys.exit()

    # determine all OUTCAR files
    outcars = []
    if not args.files:
        outcars = get_outcar_files('.', args.recursive)

    for item in args.files:
        if os.path.isdir(item):
            outcars += get_outcar_files(item, args.recursive)
        else:
            outcars.append(os.path.abspath(item))

    # remove duplicate entries from the outcars table
    outcars = uniq(outcars)

    # read metadata from all outcar files
    data = []
    max_types = 1
    for item in outcars:
        if item.endswith('.gz'):
            f = gzip.open(item, 'rb')
        else:
            f = open(item, 'r')
        data.append(scan_outcar_file(f))
        f.close()
        max_types = max(max_types, len(data[-1][1]))

    # check the types string
    types = dict()
    numbers = dict()
    if args.c:
        if len(args.c.split(',')) > max_types:
            sys.stderr.write(
                "\nERROR: There are too many items in you -c string!\n")
            sys.exit()
        if len(args.c.split(',')) < max_types:
            sys.stderr.write(
                "\nERROR: There are not enough items in you -c string!\n")
            sys.exit()
        for item in args.c.split(','):
            if len(item.split('=')) != 2:
                sys.stderr.write("\nERROR: Could not read the -c string.\n")
                sys.stderr.write("Maybe a missing or extra '=' sign?\n")
                sys.exit()
            else:
                try:
                    name = str(item.split('=')[0])
                    number = int(item.split('=')[1])
                except:
                    sys.stderr.write("\nERROR: Could not read the -c string\n")
                    sys.exit()
                if number >= max_types:
                    sys.stderr.write(
                        "\nERROR: The atom type for %s is invalid!\n" % name)
                    sys.exit()
                if name in types:
                    sys.stderr.write(
                        "\nERROR: Duplicate atom type found in -c string\n")
                    sys.exit()
                if number in numbers:
                    sys.stderr.write(
                        "\nERROR: Duplicate atom number found in -c string\n")
                    sys.exit()
                types[name] = number
                numbers[number] = name

    # process all the outcar files
    sr = args.configs_range
    if sr == None:
        windex = None
    else:
        sr = sr.split()
        sr0 = int(sr[0])
        sr1 = int(sr[1])
        windex = range(sr0, sr1 + 1)

    if args.final:
        windex = 'final'

    process_outcar_file_v5_dev(
        outcars, data, numbers, types, max_types, windex=windex, fout='test.configs')
