#  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@stfc.ukri.org'

import chemsh
from math         import ceil
from .            import _Tasks
from ..interfaces import namd
from ..objects    import resultmd

# TODO: to implement
dl_poly = None
#from .. import dl_poly

# we provide users some synonyms
_synons = {
            'dl_poly': dl_poly,
            'dl-poly': dl_poly,
            'dlpoly' : dl_poly,
            'dlp'    : dl_poly,
            'namd'   : namd.NAMD,
          }

class MD(_Tasks):
    '''Molecular Dynamics'''

    _kwargs = {
                'boundary'              :'periodic',
                'chk'                   :'',
                'dcd'                   :'',
                'driver'                :'dl_poly',
                'ensemble'              :'NPT',
                'ff'                    :'charmm',
                'fix'                   :[],
                'freq_chk'              :-1,
                'freq_dcd'              :-1,
                'freq_xst'              :-1,
                'freq_out_energy'       :-1,
                'freq_out_pressure'     :-1,
                'langevin'              : True,
                'langevin_temperature'  : None,
                'langevin_damping'      : 1,
                'langevin_H'            : False,
                'langevin_piston'       : False,
                'langevin_piston_target':   1.01325,        # bar
                'langevin_piston_period': 100.0,
                'langevin_piston_decay' :  50.0,
                'langevin_piston_temp'  : 293.15,
                'minimise'              : 100,
                'nsnapshots'            : 10,
                'nsteps'                : 100,
                'pbc_wrap'              :[],                # list of atom indices to perform PBC wrapping
                'result'                : resultmd.ResultMD(),
                'rigid'                 :'all',
                'seed'                  : 2020,
                'spheric_r0'            :-1.0,
                'spheric_k'             : 10.0,
                'spheric_exponent'      : 2,
                'temperature'           : 293.15,
                'theory'                : None,
                'timestep'              : 1.0,               # femtosecond
              }

    _internals  = { '_default_freq_chk': True,
                    '_default_freq_dcd': True,
                    '_default_freq_xst': True,
                    '_default_freq_out_energy': True,
                    '_default_freq_out_pressure': True,
                  }


    @property
    def nsteps(self):
        '''Number of MD steps'''

        return self._nsteps


    @nsteps.setter
    def nsteps(self, val):
        '''Setter of nsteps'''

        self._nsteps = val


    @property
    def freq_chk(self):
        '''Freqency of writing CHK file'''

        if self._freq_chk == -1 or self._default_freq_chk:
            if self.nsteps > 100:
                self._freq_chk = ceil(self.nsteps/100)
            else:
                self._freq_chk = 1

        # YL 08/08/2023: freq can't be greater than nsteps
        self._freq_chk = min(self._freq_chk, self.nsteps)

        return self._freq_chk


    @freq_chk.setter
    def freq_chk(self, val):
        '''Setter of freq_chk'''

        if val != -1:
            self._default_freq_chk = False

        self._freq_chk = val


    @property
    def freq_dcd(self):
        '''Freqency of writing CHK file'''

        if self._freq_dcd == -1 or self._default_freq_dcd:
            if self.nsteps > 100:
                self._freq_dcd = ceil(self.nsteps/100)
            else:
                self._freq_dcd = 1

        # YL 08/08/2023: freq can't be greater than nsteps
        self._freq_dcd = min(self._freq_dcd, self.nsteps)

        return self._freq_dcd


    @freq_dcd.setter
    def freq_dcd(self, val):
        '''Setter of freq_dcd'''

        if val != -1:
            self._default_freq_dcd = False

        self._freq_dcd = val


    @property
    def freq_xst(self):
        '''Freqency of writing CHK file'''

        if self._freq_xst == -1 or self._default_freq_xst:
            if self.nsteps > 100:
                self._freq_xst = ceil(self.nsteps/100)
            else:
                self._freq_xst = 1

        # YL 08/08/2023: freq can't be greater than nsteps
        self._freq_xst = min(self._freq_xst, self.nsteps)

        return self._freq_xst


    @freq_xst.setter
    def freq_xst(self, val):
        '''Setter of freq_xst'''

        if val != -1:
            self._default_freq_xst = False

        self._freq_xst = val


    @property
    def freq_out_energy(self):
        '''Freqency of writing CHK file'''

        if self._freq_out_energy == -1 or self._default_freq_out_energy:
            if self.nsteps > 100:
                self._freq_out_energy = ceil(self.nsteps/100)
            else:
                self._freq_out_energy = 1

        # YL 08/08/2023: freq can't be greater than nsteps
        self._freq_out_energy = min(self._freq_out_energy, self.nsteps)

        return self._freq_out_energy


    @freq_out_energy.setter
    def freq_out_energy(self, val):
        '''Setter of freq_out_energy'''

        if val != -1:
            self._default_freq_out_energy = False

        self._freq_out_energy = val


    @property
    def freq_out_pressure(self):
        '''Freqency of writing CHK file'''

        if self._freq_out_pressure == -1 or self._default_freq_out_pressure:
            if self.nsteps > 100:
                self._freq_out_pressure = ceil(self.nsteps/100)
            else:
                self._freq_out_pressure = 1

        # YL 08/08/2023: freq can't be greater than nsteps
        self._freq_out_pressure = min(self._freq_out_pressure, self.nsteps)

        return self._freq_out_pressure


    @freq_out_pressure.setter
    def freq_out_pressure(self, val):
        '''Setter of freq_out_pressure'''

        if val != -1:
            self._default_freq_out_pressure = False

        self._freq_out_pressure = val


    @property
    def is_periodic(self):
        '''Is periodic or not'''

        from numpy import any, char

        if issubclass(type(self.boundary), str):
            return not any(char.startswith(['sphere','spherical','droplet'], self.boundary.strip().lower()))
        else:
            if self.boundary:
                return True
            else:
                return False


    def run(self, info_filename='_profile_md.npz', **kwargs):
        ''''''

        from numpy   import arange, savez_compressed

        # Langevin
        if not self.langevin_temperature:
            self.langevin_temperature = self.temperature

        if not self.chk:
            self.chk = '_'+self.theory.__class__.__name__.lower()+'.chk'
        if not self.dcd:
            self.dcd = '_'+self.theory.__class__.__name__.lower()+'.dcd'

        # if path can be imported as module call run.runLib(), otherwise run.runExec() (see <dict> run.run().runTypes)
        driver = _synons.get(self.driver.lower(), dl_poly)

        # YL 27/02/2021: special treatments for read-only options
        #                TODO these can be made automatic in the initialisation of <class Options>
        self.options.is_periodic = self.is_periodic

        # not all MD pacakge can be used as a general-purpose driver
        if issubclass(type(self.theory), driver):
            if self.spheric_r0 == -1.0:
                self.spheric_r0 = self.theory.frag.radius
            ierror = self.theory.run(MDOptions=self.options, **kwargs)
        else:
            ierror = driver.run(self.theory, **kwargs)

        print(f"\n Number of data points saved to file {info_filename}:")
        print( " Energy      : {:>5d}".format(self.result.energies.size))
        print( " Potential   : {:>5d}".format(self.result.potentials.size))
        print( " Kinetic     : {:>5d}".format(self.result.kinetics.size))
        print( " Temperature : {:>5d}".format(self.result.temperatures.size))
        print( " Minimisation: {:>5d}".format(self.result.opt_energies.size))
        print( "\n Run command to visualise the data (matplotlib required):")
        print(f" chemsh --plot {info_filename} --units [...]\n\n")

        # dump the result arrays to a compressed file in NumPy .npz format
        # unit: always a.u. and fs
        savez_compressed(info_filename,
                         energies       = self.result.energies,
                         potentials     = self.result.potentials,
                         kinetics       = self.result.kinetics,
                         temperatures   = self.result.temperatures,
                         opt_energies   = self.result.opt_energies,
                         times          = arange(0.0, self.timestep*len(self.result.energies)*self.freq_out_energy, self.timestep*self.freq_out_energy),
                        )

        if not ierror and self.nsnapshots and not kwargs.get('dryrun', False):
            self.takeSnapshots()

        return self


    def takeSnapshots(self, pbc_wrap_centre=-1):
        '''Take nsnapshots snapshots out of the trajectory file'''

        from sys  import stdout
        from ..io import file, files
        from ..dl_py2f.utils import miscutils

        # take snapshots with PBC wrapping
        try:
            # box parameters from saved data file (e.g., .xst)
            # the shape (nrecords,) is dummy
            box_parameters  = files.xst.parse(self.theory.xst)
            nbox_parameters = box_parameters.shape[0]
            # try something else (e.g., .xsc) if the data don't look right
            nrecords = files.dcd.parse(self.dcd, receiver={}, nrecords_only=True)
            if nrecords+1 != nbox_parameters:
                print("\n ChemShell MD WARNING: numbers of records in {} and {} do not match!\n".format(self.dcd, self.theory.xst))
                box_parameters  = files.xst.parse(self.theory.xsc)
                # not in use
                nbox_parameters = 1
            files.dcd.takeSnapshots(self.dcd, self.theory.frag, nsnapshots=self.nsnapshots, box_parameters=box_parameters, pbc_wrap=self.pbc_wrap, pbc_wrap_centre=pbc_wrap_centre)

        # take snapshots without PBC wrapping
        except:
            print('\n Taking snapshots without PBC wrapping...')
            stdout.flush()
            # skip failing snapshots unless at high debug level
            if chemsh.CHEMSH_DEBUG > 4:
                raise
            files.dcd.takeSnapshots(self.dcd, self.theory.frag, nsnapshots=self.nsnapshots)


