
#  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/>.

import chemsh
from numpy            import int32, float32, float64
from sys              import stdout
from ...dl_py2f.utils import fileutils, miscutils

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

# YL 20/06/2024: fixed bugs

# YL 28/07/2023: NAMD format (tuple for dtype and number of data)
field2dtype = {'magic'    :( int32  ,1), 
               'header'   :('S4'    ,1),
               'nframes'  :( int32  ,1),
               'step0'    :( int32  ,1),
               'freq_dcd' :( int32  ,1),
               'nsteps'   :( int32  ,1),
               '5 dummies':( int32  ,5),
               'timestep' :( float32,1),
               'unitcell' :( int32  ,1),
               '8 dummies':( int32  ,8),
               '4 magics' :( int32  ,4),
               'title1'   :('S80'   ,1),
               'title2'   :('S80'   ,1),
               '2 dummies':( int32  ,2),
               'natoms'   :( int32  ,1),
               '1 dummy'  :( int32  ,1),
#               'nconsts1' :( int32  ,1),
#               'consts'   :( float64,1),
#               'nconsts2' :( int32  ,1),
              }


# NAMD DCD is binary in single-precision (32-bit)
# see: NAMD/src/dcdlib.C
def parse(filename, _regex=None, receiver={}, nrecords_only=False, prefix='', **kwargs):
    '''Parse a DCD trajectory file'''

    from math  import ceil
    from numpy import fromfile, int32, float32, vstack
    from os    import path

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

    with open(filename, 'rb') as fp:

        # hardcoded by NAMD: 84
        magic = fromfile(fp, dtype=int32, count=1)[0]
        receiver.update({'magic':magic})
        if chemsh.CHEMSH_DEBUG > 4:
            print("\n ChemShell DCD parser: magic number =", magic)
    
        header = fromfile(fp, dtype='S4', count=1)[0]
        receiver.update({'header':header})
        if chemsh.CHEMSH_DEBUG > 4:
            print("\n ChemShell DCD parser: header =", header)
    
        # number of frames
        nframes = fromfile(fp, dtype=int32, count=1)[0]
        receiver.update({'nframes':nframes})

        # starting timestep
        step0 = fromfile(fp, dtype=int32, count=1)[0]
        receiver.update({'step0':step0})

        # frequency of writing DCD
        freq_dcd = fromfile(fp, dtype=int32, count=1)[0]
        receiver.update({'freq_dcd':freq_dcd})

        # number of total timesteps
        nsteps = fromfile(fp, dtype=int32, count=1)[0]
        receiver.update({'nsteps':nsteps})
    
        # five zeros
        dummies = fromfile(fp, dtype=int32, count=5)
        receiver.update({'5 dummies':dummies})
    
        # timestep in fs converted with an NAMD factor
        timestep = fromfile(fp, dtype=float32, count=1)[0]*float32(48.88821)
        receiver.update({'timestep':timestep})
        if chemsh.CHEMSH_DEBUG > 4:
            print("\n ChemShell DCD parser: timestep =", timestep)

        # TODO
        unitcell = fromfile(fp, dtype=int32, count=1)[0]
        receiver.update({'unitcell':unitcell})
    
        # eight zeros
        dummies = fromfile(fp, dtype=int32, count=8)
        receiver.update({'8 dummies':dummies})
    
        # hardcoded by NAMD: 24, 84, 164, 2
        magic = fromfile(fp, dtype=int32, count=4)
        receiver.update({'4 magics':magic})
        if chemsh.CHEMSH_DEBUG > 4:
            print("\n ChemShell DCD parser: magic numbers =", magic)
    
        title1 = fromfile(fp, dtype='S80', count=1)
        receiver.update({'title1':title1})
    
        title2 = fromfile(fp, dtype='S80', count=1)
        receiver.update({'title2':title2})
    
        # hardcoded by NAMD: 164, 4
        dummies = fromfile(fp, dtype=int32, count=2)
        receiver.update({'2 dummies':dummies})

        # natoms
        natoms = fromfile(fp, dtype=int32, count=1)[0]
        receiver.update({'natoms':natoms})
        if chemsh.CHEMSH_DEBUG > 4:
            print("\n ChemShell DCD parser: natoms =", natoms)

        # hardcoded by NAMD: 4
        dummy = fromfile(fp, dtype=int32, count=1)
        receiver.update({'1 dummy':dummy})
    
        nrecords = ceil(nsteps/freq_dcd)

        # YL 25/05/2022: get nrecords only
        # YL FIXME: here may be nframes??
        if nrecords_only:
            return nrecords

        receiver.update({'nrecords':nrecords})
        if chemsh.CHEMSH_DEBUG > 4:
            print("\n ChemShell DCD parser: nrecords =", nrecords)
    
        for i in range(nframes):

            start = fp.tell()

            nconsts1 = fromfile(fp, dtype=int32, count=1)[0]
            # original order: A, gamma, B, beta, alpha, C
            consts = fromfile(fp, dtype='float64', count=6)[[0,2,5,4,3,1]]
            nconsts2 = fromfile(fp, dtype=int32, count=1)[0]
            if nconsts1 == nconsts2 == 48:
                if chemsh.CHEMSH_DEBUG > 4:
                    print("\n ChemShell DCD parser: number of lattice constants =", int(nconsts1/8))
                    print("\n ChemShell DCD parser: lattice constants (A,B,C,alpha,beta,gamma) =\n", consts)
                receiver.update({'cell_nconsts_{}'.format(i):nconsts1})
                receiver.update({'cell_constants_{}'.format(i):consts})
            # wind back if no lattice block
            else:
                fp.seek(start)

            # NB: here is natoms*4
            nx1 = fromfile(fp, dtype=int32, count=1)[0]

            # NAMD: float32
            x  = fromfile(fp, dtype='float32', count=natoms)
            # NB: here is natoms*4
            nx2 = fromfile(fp, dtype=int32, count=1)[0]
            if nx2 != nx1:
                print("\n ChemShell DCD parser: parsing has failed! nx2 =", nx2)
            if chemsh.CHEMSH_DEBUG > 4:
                print("\n ChemShell DCD parser: nx1 =", nx2)
            # NB: here is natoms*4
            ny1 = fromfile(fp, dtype=int32, count=1)[0]
            if chemsh.CHEMSH_DEBUG > 4:
                print("\n ChemShell DCD parser: ny1 =", ny1)
            y  = fromfile(fp, dtype='float32', count=natoms)
            # NB: here is natoms*4
            ny2 = fromfile(fp, dtype=int32, count=1)[0]
            if chemsh.CHEMSH_DEBUG > 4:
                print("\n ChemShell DCD parser: ny2 =", ny2)
            # NB: here is natoms*4
            nz1 = fromfile(fp, dtype=int32, count=1)[0]
            if chemsh.CHEMSH_DEBUG > 4:
                print("\n ChemShell DCD parser: nz1 =", nz1)
            z  = fromfile(fp, dtype='float32', count=natoms)
            # NB: here is natoms*4
            nz2 = fromfile(fp, dtype=int32, count=1)[0]
            if chemsh.CHEMSH_DEBUG > 4:
                print("\n ChemShell DCD parser: nz2 =", nz2)
            receiver.update({i:vstack([x,y,z]).T})
    
    return 0


def write(filename, data_dict, byteorder='big', mute=False):
    '''A DCD writer'''

    from numpy      import int32, float32
    from ...dl_py2f import utils

    selectcases = {int32  :4,
                   float32:4,
                   float64:8,
                  }

    with open(filename, 'wb') as fp:
        for k, (dtype, n) in field2dtype.items():
            if k == 'timestep':
                (data_dict[k]/float32(48.88821)).tofile(fp)
                continue
            data_dict[k].tofile(fp)
        indices = [ i for i in data_dict.keys() if utils.miscutils.isInteger(i) ]
        indices.sort()
        natoms = data_dict['natoms']
        # YL 20/06/2024: we must cast the type (though NumPy 2.0 may have changed the behaviour)
        natoms4 = int32(natoms*4)
        for i in indices:
            # YL 20/06/2024: non-periodic
            if f'cell_nconsts_{i}' in data_dict:
                # nconsts of cell
                data_dict[f'cell_nconsts_{i}'].tofile(fp)
            if f'cell_constants_{i}' in data_dict:
                # constants of cell (reordered from A,B,C,alpha,beta,gamma to A,gamma,B,beta,alpha,C)
                data_dict[f'cell_constants_{i}'][[0,5,1,4,3,2]].tofile(fp)
            if f'cell_nconsts_{i}' in data_dict:
                # nconsts of cell
                data_dict[f'cell_nconsts_{i}'].tofile(fp)
            # natoms*4 of X
            natoms4.tofile(fp)
            # X coordinates
            data_dict[i][:,0].tofile(fp)
            # natoms*4 of X
            natoms4.tofile(fp)
            # natoms*4 of Y
            natoms4.tofile(fp)
            # Y coordinates
            data_dict[i][:,1].tofile(fp)
            # natoms*4 of Y
            natoms4.tofile(fp)
            # natoms*4 of Z
            natoms4.tofile(fp)
            # Z coordinates
            data_dict[i][:,2].tofile(fp)
            # natoms*4 of Z
            natoms4.tofile(fp)

    if not mute:
        print(f' >>> \"{filename}\" saved.')


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

    import re
    from os.path  import join
    from numpy    import array, int32
    from datetime import date
    from ...io    import file, format

    dcd1 = {}
    parse(filename1, 'dcd', receiver=dcd1)
    dcd0 = {}
    parse(filename0, 'dcd', receiver=dcd0)
    nframes0 = dcd0['nframes']
    stdout.flush()

    # check and update the 1st DCD
    for k, v in dcd0.items():
        if 'magic' in str(k):
            pass
        if 'dummies' in str(k):
            pass
        if k == 'nframes':
            format.printwrap(f'Merging {filename1} of {dcd1[k]} records into {filename0} of {dcd0[k]} records...')
            dcd0.update({k:dcd1[k]+dcd0[k]})
        if k == 'nsteps':
            dcd0.update({k:dcd1[k]+dcd0[k]})
        if k == 'timestep':
            if dcd1[k] != dcd0[k]:
                print(f"\n >>> ChemShell DCD merger WARNING: {filename0} and {filename1} have different timesteps!")
        if k == 'title1':
            dcd0.update({k:array([f'REMARKS FILENAME: {output} MERGED BY Py-ChemShell'[:80]], dtype=field2dtype['title1'][0])})
        if k == 'title2':
            dcd0.update({k:array([f'REMARKS DATE: {date.today().strftime("%d/%m/%y")}'[:80]], dtype=field2dtype['title2'][0])})
        if k == 'natoms':
            if dcd1[k] != dcd0[k]:
                print(f"\n >>> ChemShell DCD merger ERROR: {filename0} and {filename1} have different nsteps!")
                stdout.flush()
                exit(999)

    # attach data from the 2nd DCD to the first one
    for k, v in dcd1.items():
        if 'cell_nconsts_' in str(k):
            index = int(re.sub('[^0-9]', '', k)) + int(nframes0)
            dcd0.update({'cell_nconsts_'+str(index):v})
        if 'cell_constants_' in str(k):
            index = int(re.sub('[^0-9]', '', k)) + int(nframes0)
            dcd0.update({'cell_constants_'+str(index):v})
        if type(k) is int:
            dcd0.update({k+int(nframes0):v})

    write(output, dcd0)


def takeSnapshots(input, fragObj, nsnapshots=10, pbc_wrap=[], pbc_wrap_centre=-1, box_parameters=[], mute=False, irecord=None, recentre=True):
    '''Take nsnapshots snapshots out of the DCD file'''

    from numpy    import linspace
    from ...data  import unitconvert
    from ...io    import format

    result = dict()

    # filename as input
    if type(input) is str:
        dbuff = dict()
        parse(input, receiver=dbuff)
    # dict data as input
    elif type(input) is dict:
        dbuff = input

    nframes = dbuff.get('nframes', 0)

    freq     = dbuff['freq_dcd']
    nsteps   = dbuff['nsteps']
    timestep = dbuff['timestep']

    if nframes:

        # a dummy fragment object to hold the data
        dummyFrag = fragObj.merge(reconnect=False)

        # return only the requested snapshot
        if miscutils.isInteger(irecord):
            print( '\n '+'-'*70)
            print(f' Requesting snapshot of record {irecord}...')
            print( ' '+'-'*70)
            stdout.flush()
            dummyFrag.coords = dbuff[irecord]*unitconvert.Angstrom2Bohr
            # apply PBC wrap
            if len(box_parameters):
                dummyFrag.pbc_wrap(box_parameters=box_parameters[irecord], select=pbc_wrap, centre=pbc_wrap_centre)
            return dummyFrag

        print( '\n '+'-'*70)
        print(f' Taking snapshots out of {input}...')
        print( ' '+'-'*70)
        stdout.flush()

        # nsnapshots*2 cannot be larger than nframes
        nsnapshots = min(int(nframes/2), nsnapshots)
        len_index = len(str(nsnapshots))

        # use only structures of the late trajectory
        points = linspace(0, nframes-1, nsnapshots*2, dtype=int)[nsnapshots:]

        # archive old files
        fileutils.archiveDir('_snapshots')

        # for printing
        time_snapshots = []

        # loop over snapshots
        for isnapshot, irecord in enumerate(points):
            strbuff = f'\n Snapshot {isnapshot}'
            print(strbuff)
            print(' '+'-'*(len(strbuff)-2))
            stdout.flush()
            time_snapshots += ['{:.2f}'.format(irecord*freq*timestep/1000)]
            dummyFrag.coords = dbuff[irecord]*unitconvert.Angstrom2Bohr

            # apply PBC wrap
            if len(box_parameters):
                dummyFrag.cell.vectors[:,0] = box_parameters[irecord]['A']
                dummyFrag.cell.vectors[:,1] = box_parameters[irecord]['B']
                dummyFrag.cell.vectors[:,2] = box_parameters[irecord]['C']
                dummyFrag.pbc_wrap(box_parameters=box_parameters[irecord], select=pbc_wrap, centre=pbc_wrap_centre)

            # YL 30/10/2023: recentre
            if recentre:
                dummyFrag.recentre()
            filename = '_snapshot_{:0'+str(len_index)+'d}.pdb'
            dummyFrag.save(filename.format(isnapshot), mute=mute)
            fileutils.archiveFile(filename.format(isnapshot), dir='_snapshots', suffix='')
            # use the occ column to save the charges
            filename = '_snapshot_{:0'+str(len_index)+'d}.pqr'
            dummyFrag.save(filename.format(isnapshot), mute=mute)
            fileutils.archiveFile(filename.format(isnapshot), dir='_snapshots', suffix='')

        print( '\n '+'-'*70)
        print( ' Summary of snapshots taken from DCD')
        print( ' '+'-'*70)
        print(f' Num. of frames   : {nframes}')
        print(f' MD length        : {nsteps*timestep/1000:.2f} ps')
        print( ' Directory        : _snapshots')
        print(f' Num. of snapshots: {nsnapshots}')
        print( ' Simulation times (in ps) of snapshots:')
        format.printwrap(' '.join(time_snapshots), prepend_blanc_line=False)
        print( ' '+'-'*70)

    else:
        print("\n ERROR: no record found in trajectory file {}\n".format(filename))
        exit(999)


