#  Copyright (C) 2021 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.ukri.org>'


from sys             import stdout
from ..dl_py2f.utils import fileutils, miscutils

@miscutils.time
def analyseMD(fragObj,
              dcd='',
#              xsc='',
#              xst='',
              nsnapshots=10,
              select=[],
              align=True,
              units=['a.u.', 'ps'],
              info_filename='_profile_analysis.npz',
              timestep:float=0.0,
              nsteps:int=0,
              directory='_analysis',
              **kwargs):
    ''''''

    from .       import getUnits
    from ..io    import banner, files
    from ..base  import geometry
    from numpy   import arange, empty, linspace, mean, nan, savez_compressed, sqrt

    lunit, tunit, lfactor, tfactor = getUnits(units)

    # simulation time
    total_time = timestep*nsteps

    banner.printModuleBanner('ChemShell MD Analysis Toolkit')

    # read the DCD data to a dictionary
    dict_dcd = dict()
    files.dcd.parse(dcd, receiver=dict_dcd)
    nrecords = dict_dcd.get('nrecords', 0)
    if nrecords:
        print(" {} records found in {}\n".format(nrecords, dcd))
        # archive old files
        fileutils.archiveDir(directory)
    else:
        return

    points = linspace(0, nrecords-1, num=nsnapshots+1, dtype=int)[1:]

# YL 13/08/2022: decided not to introduce PBC wrapping to MD analysis for now (still can do it when taking snapshots)
#    # YL 13/08/2022: extract the box parameters from the XST file
#    box_parameters = []
#    if xst:
#        box_parameters  = files.xst.parse(xst)
#        nbox_parameters = box_parameters.shape[0]
#        if nrecords+1 != nbox_parameters:
#            print("\n WARNING: numbers of records in {} and {} do not match!\n".format(dcd, xst))
#            # try something else (e.g., .xsc) if the data don't look right
#            if xsc:
#                box_parameters  = files.xst.parse(xsc)
#                nbox_parameters = box_parameters.shape[0]
#                if nrecords+1 != nbox_parameters:
#                    print("\n WARNING: numbers of records in {} and {} do not match!\n".format(dcd, xsc))
#                    nbox_parameters = 1

    # get the first (starting) structure from the DCD
    snapshot0 = files.dcd.takeSnapshots(dict_dcd, fragObj, mute=True, irecord=0)
    dummyFrag0 = snapshot0.getSelected(select)
    # translate to centroid
    dummyFrag0.coords[:] -= dummyFrag0.centroid

    # write out the structure
    len_index = len(str(nsnapshots))
    if align:
        filename_pdb = '_snapshot_{:0'+str(len_index)+'d}_aligned.pdb'
        filename_pqr = '_snapshot_{:0'+str(len_index)+'d}_aligned.pqr'
        sbuff =  ' Aligned snapshots:\n'
    else:
        filename_pdb = '_snapshot_{:0'+str(len_index)+'d}_unaligned.pdb'
        filename_pqr = '_snapshot_{:0'+str(len_index)+'d}_unaligned.pqr'
        sbuff =  ' Unaligned snapshots:\n'
    dummyFrag0.save(filename_pdb.format(0), mute=True)
    dummyFrag0.save(filename_pqr.format(0), mute=True)
    fileutils.archiveFile(filename_pdb.format(0), dir=directory, suffix='')
    fileutils.archiveFile(filename_pqr.format(0), dir=directory, suffix='')
    if total_time:
        sbuff += f' Snapshot    Time ({tunit})    DCD record    RMSD ({lunit})\n'
        sbuff += ' {:>5d}   {:>12.1f}       {:>5d}    {:>12.4f}\n'.format(0, 0.0, 0, 0.0)
    else:
        sbuff += f' Snapshot    DCD record    RMSD ({lunit})\n'
        sbuff += ' {:>5d}         {:>5d}    {:>12.4f}\n'.format(0, 0, 0.0)

    if not len(select):
        select = ...
        nselecteds = snapshot0.natoms
    else:
        nselecteds = len(select)

    print(" RMSD analysis: {} atoms selected\n".format(nselecteds))
    stdout.flush()

    array_rmsds    = empty((len(points)+1,))
    array_rmsds[:] = nan
    array_rmsds[0] = 0.0
    array_times    = [0.0]

    # get the remaining structures one by one so that we keep at most two structures in memory at the same time
    for isnapshot, ipoint in enumerate(points):

        # get the ith structure from the DCD
        snapshoti = files.dcd.takeSnapshots(dict_dcd, fragObj, mute=True, irecord=ipoint)
        dummyFrag = snapshoti.getSelected(select)

        # align
        if align:
            # get the optimal rotation matrix of the Kabsch algorithm
            R = geometry.getKabschMatrix(snapshot0.coords[select], dummyFrag.coords)
            # rotate with the matrix
            dummyFrag.coords = dummyFrag.coords @ R
            # translate to centroid
            dummyFrag.coords[:] -= dummyFrag.centroid

        # calculate the RMSD
        # NB: to avoid misunderstanding, we always save the data in converted units
        rmsd = sqrt(((dummyFrag.coords-dummyFrag0.coords)**2).sum()/nselecteds)*lfactor
        array_rmsds[isnapshot+1] = rmsd

        if total_time:
            # NB: to avoid misunderstanding, we always save the data in converted units
            time = total_time/len(points)*(isnapshot+1)*tfactor
            array_times.append(time)
            sbuff += ' {:>5d}   {:>12.1f}       {:>5d}    {:>12.4f}\n'.format(isnapshot+1, time, ipoint, rmsd)
        else:
            # dummy time series
            array_times.append(0.0)
            sbuff += ' {:>5d}         {:>5d}    {:>12.4f}\n'.format(isnapshot+1, ipoint, rmsd)
        # write out the structure
        dummyFrag.save(filename_pdb.format(isnapshot+1), mute=True)
        dummyFrag.save(filename_pqr.format(isnapshot+1), mute=True)
        fileutils.archiveFile(filename_pdb.format(isnapshot+1), dir=directory, suffix='')
        fileutils.archiveFile(filename_pqr.format(isnapshot+1), dir=directory, suffix='')

        savez_compressed(info_filename,
                         rmsds   = array_rmsds,
                         times   = array_times,
                         lunit   = lunit,
                         tunit   = tunit,
#                         lfactor = lfactor,
#                         tfactor = tfactor,
                        )
        fileutils.archiveFile(info_filename, dir=directory, suffix='')

    print(sbuff)
    print("\n {} snapshots taken out from the trajectory and saved to directory:\n {}".format(nsnapshots, directory))
    print("\n Run command to visualise the data (matplotlib required):")
    print(f" chemsh --plot {directory}/{info_filename}")


