#  Copyright (C) 2020 The authors of Py-ChemShell
#
#  This file is part of Py-ChemShell.
#
#  Py-ChemShell is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation, either version 3 of the
#  License, or (at your option) any later version.
#
#  Py-ChemShell 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 Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with Py-ChemShell.  If not, see
#  <http://www.gnu.org/licenses/>.

__author__ = 'You Lu <you.lu@stfc.ac.uk>'

def parse(filename, _regex=None, receiver=None, prefix='', **kwargs):
    '''Parse a NAMD XST file'''

    from numpy   import genfromtxt
    from os      import path
    from ...data import unitconvert

    dtype = [('step'  , 'i8'      ),
             ('A'     , 'f8', (3,)),
             ('B'     , 'f8', (3,)),
             ('C'     , 'f8', (3,)),
             ('origin', 'f8', (3,)),
             ('s'     , 'f8', (6,))
            ]

    # YL 13/08/2022: path prefix
    filename = path.join(prefix, filename)

    abuff = genfromtxt(filename,
                       dtype=dtype,
                       comments='#')

    abuff['A']      *= unitconvert.unit2au['Angstrom']
    abuff['B']      *= unitconvert.unit2au['Angstrom']
    abuff['C']      *= unitconvert.unit2au['Angstrom']
    abuff['origin'] *= unitconvert.unit2au['Angstrom']
    # cell strain(?)
    abuff['s']      *= unitconvert.unit2au['Angstrom']

    return abuff


def merge(filename0, filename1, output):
    '''Merge two XST files'''

    from numpy import hstack

    xst0 = parse(filename0)
    nrecords0 = len(xst0)
    xst1 = parse(filename1)
    nrecords1 = len(xst1)

    # add the increment
    xst1['step'] += xst0['step'][-1]

    # concatenate data
    print(f'\n Py-ChemShell XST merger: Assuming {filename0} was continued with {filename1}')
    merged = hstack((xst0, xst1[1:]))

    # write out a new file
    write(output, merged)


def write(filename, data_array):
    '''An XST writer'''

    from numpy                  import savetxt
    from numpy.lib.recfunctions import structured_to_unstructured

    with open(filename, 'w') as fp:
        fp.write('# NAMD extended system trajectory file written by Py-ChemShell\n')
        fp.write('#$LABELS step a_x a_y a_z b_x b_y b_z c_x c_y c_z o_x o_y o_z s_x s_y s_z s_u s_v s_w\n')
#        savetxt(fp, data_array.view(float).reshape(data_array.shape+(-1,)), fmt='%8d'+18*' %18.11e')
        # YL 27/07/2023: there's no better way of formatted printing structured array
        savetxt(fp, structured_to_unstructured(data_array, casting='safe'), fmt='%8d'+18*' %18.11e')





