"""
This module is designed for reading and writing of POSCAR files.
(C) Kai Sellschopp, TUHH Hamburg, Germany
"""

from __future__ import division
import os
import numpy as np


def readPOSCAR(filename):
    """
    This method reads in a POSCAR file and extracts the information about
    the lattice and atomic types, numbers and positions from it.
    
    INPUT:  filename       (String)

    OUTPUT: lc             (lattice constant, float)
            mat            (matrix of lattice vectors, numpy float array)
            atom_types     (atomic species, string array)
            atom_numbers   (atomic numbers, numpy int array)
            atom_positions (atomic positions, numpy float array)
            atom_comments  (comments for each position, list of strings)
            selDyn         (selective dynamics used?, boolean)
    """

    with open(filename, 'r') as f:
        poscar = f.read()
    
    poscar_lines=poscar.split('\n')

    # get the title
    title = poscar_lines[0]

    # get the lattice constant
    lc = float(''.join(poscar_lines[1].split()))

    # get the lattice vectors / matrix
    a = np.array(poscar_lines[2].split(), dtype='f')
    b = np.array(poscar_lines[3].split(), dtype='f')
    c = np.array(poscar_lines[4].split(), dtype='f')
    mat = np.array([a, b, c])

    # get the atom types and numbers
    atom_types = poscar_lines[5].split()
    atom_numbers = np.array(poscar_lines[6].split(), dtype='i')

    # determine whether selective dynamics was used
    key_chars = ['S', 's']
    selDyn = False
    for i in range(2):
        if poscar_lines[7+i][0] in key_chars:
            selDyn = True
            break
    
    # determine first line for atomic positions
    key_chars = ['D', 'd', 'K', 'k', 'C', 'c']
    fl = 0                  # first line variable

    for i in range(10):
        if poscar_lines[7+i][0] in key_chars:
            fl = 8+i
            break

    # get the atomic positions and comments
    N_total = np.sum(atom_numbers)
    atom_positions = np.zeros((N_total, 3))
    atom_comments = []
    for i in range(N_total):
        line = poscar_lines[fl+i].split()
        atom_positions[i] = np.array(line[0:3], dtype='f')
        atom_comments.append(' '.join(line[3:]))

    # return the information
    return lc, mat, atom_types, atom_numbers, atom_positions, atom_comments, selDyn, title


def writePOSCAR(filename, lc, mat, atom_types, atom_numbers, atom_positions, atom_comments='', selDyn=False, title='Structure'):
    """
    This method takes structural information to create a POSCAR which
    is written to a POSCAR file with name given by filename.

    INPUT: filename       (output file name, String)
           lc             (lattice constant, float)
           mat            (matrix of lattice vectors, numpy float array)
           atom_types     (atomic species, string array)
           atom_numbers   (atomic numbers, numpy int array)
           atom_positions (atomic positions, numpy float array)
           atom_comments  (comments for each position, list of strings)
           selDyn         (use selective dynamics?, boolean)

    OUTPUT: poscar        (String in the POSCAR format)
    """

    # initialize output
    poscar_lines = []

    # add header
    poscar_lines.append(title)

    # add lattice constant and vectors
    poscar_lines.append('{:f}'.format(lc))
    for vec in mat:
        poscar_lines.append('    {:.10f}    {:.10f}    {:.10f}'.format(vec[0],vec[1],vec[2]))

    # add element types and numbers
    atom_types_line = '  '
    atom_numbers_line = '  '
    for i in range(len(atom_types)):
        atom_types_line = atom_types_line + '{:5s} '.format(atom_types[i])
        atom_numbers_line = atom_numbers_line + '{:5s} '.format(repr(atom_numbers[i]))

    poscar_lines.append(atom_types_line)
    poscar_lines.append(atom_numbers_line)

    # add selective dynamics tag if chosen
    if selDyn:
        poscar_lines.append('Selective Dynamics')

    # if no comments are given, generate dummy comment ''
    if np.all(atom_comments==''):
        atom_comments = np.array('').repeat(len(atom_positions))

    # add atomic positions and comments
    poscar_lines.append('Direct')
    for i in range(len(atom_positions)):
        pos = atom_positions[i]
        com = atom_comments[i]
        if not com:
            poscar_lines.append('    {:.10f}    {:.10f}    {:.10f}'.format(pos[0],pos[1],pos[2]))
        else:
            poscar_lines.append('    {:.10f}    {:.10f}    {:.10f}    {:s}'.format(pos[0],pos[1],pos[2],com))

    # add a blank line in the end
    poscar_lines.append('')

    # write POSCAR to file
    poscar = '\n'.join(poscar_lines)
    with open(filename, 'w') as f:
        f.write(poscar)