# Copyright (C) 2015 Atsushi Togo
# All rights reserved.
#
# This file is part of phonopy.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in
#   the documentation and/or other materials provided with the
#   distribution.
#
# * Neither the name of the phonopy project nor the names of its
#   contributors may be used to endorse or promote products derived
#   from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

import sys
import numpy as np

from phonopy.file_IO import collect_forces
from phonopy.interface.vasp import (get_scaled_positions_lines,
                                    sort_positions_by_symbols,
                                    check_forces,
                                    get_drift_forces)
from phonopy.units import Bohr
from phonopy.structure.atoms import PhonopyAtoms as Atoms
from phonopy.structure.atoms import symbol_map

def parse_set_of_forces(num_atoms, forces_filenames, verbose=True):
    hook = 'Forces :'
    is_parsed = True
    force_sets = []

    for i, filename in enumerate(forces_filenames):
        if verbose:
            sys.stdout.write("%d. " % (i + 1))
        f = open(filename)
        elk_forces = collect_forces(f,
                                    num_atoms,
                                    hook,
                                    [3, 4, 5],
                                    word='total force')
        if check_forces(elk_forces, num_atoms, filename, verbose=verbose):
            drift_force = get_drift_forces(elk_forces,
                                           filename=filename,
                                           verbose=verbose)
            force_sets.append(np.array(elk_forces) - drift_force)
        else:
            is_parsed = False

    if is_parsed:
        return force_sets
    else:
        return []

def read_elk(filename):
    elk_in = ElkIn(open(filename).readlines())
    tags = elk_in.get_variables()
    avec = [tags['scale'][i] * np.array(tags['avec'][i]) for i in range(3)]
    spfnames = tags['atoms']['spfnames']
    symbols = [x.split('.')[0] for x in spfnames]
    numbers = []
    for s in symbols:
        if s in symbols:
            numbers.append(symbol_map[s])
        else:
            numbers.append(0)
            
    for i, n in enumerate(numbers):
        if n == 0:
            for j in range(1, 119):
                if not (j in numbers):
                    numbers[i] = j
                    break
    pos_all = []
    num_all = []
    for num, pos in zip(numbers, tags['atoms']['positions']):
        pos_all += pos
        num_all += [num] * len(pos)

    return Atoms(numbers=num_all,
                 cell=avec,
                 scaled_positions=pos_all), spfnames

def write_elk(filename, cell, sp_filenames):
    f = open(filename, 'w')
    f.write(get_elk_structure(cell, sp_filenames))

def write_supercells_with_displacements(supercell,
                                        cells_with_displacements,
                                        sp_filenames):
    write_elk("supercell.in", supercell, sp_filenames)
    for i, cell in enumerate(cells_with_displacements):
        write_elk("supercell-%03d.in" % (i + 1), cell, sp_filenames)

def get_elk_structure(cell, sp_filenames=None):
    lattice = cell.get_cell()
    (num_atoms,
     symbols,
     scaled_positions,
     sort_list) = sort_positions_by_symbols(cell.get_chemical_symbols(),
                                            cell.get_scaled_positions())

    if sp_filenames is None:
        spfnames = [s + ".in" for s in symbols]
    else:
        spfnames = sp_filenames

    lines = ""
    lines += "avec\n"
    lines += ((" %21.16f" * 3 + "\n") * 3) % tuple(lattice.ravel())
    lines += "atoms\n"
    n_pos = 0
    lines += " %d\n" % len(num_atoms)
    for i, (n, s) in enumerate(zip(num_atoms, spfnames)):
        lines += " \'%s\'\n" % s
        lines += " %d\n" % n
        lines += get_scaled_positions_lines(scaled_positions[n_pos:(n_pos + n)])
        if i < len(num_atoms) - 1:
            lines += "\n"
        n_pos += n
        
    return lines

class ElkIn(object):
    def __init__(self, lines):
        self._set_methods = {'atoms':  self._set_atoms,
                             'avec':   self._set_avec,
                             'scale':  self._set_scale,
                             'scale1': self._set_scale1,
                             'scale2': self._set_scale2,
                             'scale3': self._set_scale3}
        self._tags = {'atoms': None,
                      'avec':  None,
                      'scale': [1.0, 1.0, 1.0]}
        self._lines = iter(lines)
        self._collect()

    def get_variables(self):
        return self._tags

    def _collect(self):
        elements = {}
        while True:
            try:
                line_str = self._lines.next().strip()
            except StopIteration:
                break

            if len(line_str) == 0:
                continue
            if line_str[0] == '!':
                continue

            elems = line_str.split()
            if elems[0] in self._set_methods:
                self._set_methods[elems[0]]()

    def _set_atoms(self):
        nspecies = int(self._lines.next().split()[0])
        spfnames = []
        positions = []
        for i in range(nspecies):
            spfnames.append(self._lines.next().split()[0].strip('\''))
            natoms = int(self._lines.next().split()[0])
            pos_sp = []
            for j in range(natoms):
                pos_sp.append(
                    [float(x) for x in self._lines.next().split()[:3]])
            positions.append(pos_sp)
            
        self._tags['atoms'] = {'spfnames':  spfnames,
                               'positions': positions}
            
    def _set_avec(self):
        avec = []
        for i in range(3):
            avec.append([float(x) for x in self._lines.next().split()[:3]])
        self._tags['avec'] = avec

    def _set_scale(self):
        scale = float(self._lines.next().split()[0])
        for i in range(3):
            self._tags['scale'][i] = scale

    def _set_scale1(self):
        self._tags['scale'][0] = float(self._lines.next().split()[0])

    def _set_scale2(self):
        self._tags['scale'][1] = float(self._lines.next().split()[0])

    def _set_scale3(self):
        self._tags['scale'][2] = float(self._lines.next().split()[0])
        
if __name__ == '__main__':
    import sys
    from phonopy.structure.symmetry import Symmetry
    cell, sp_filenames = read_elk(sys.argv[1])
    symmetry = Symmetry(cell)
    print("# %s" % symmetry.get_international_table())
    print(get_elk_structure(cell, sp_filenames=sp_filenames))
