#  Copyright (C) 2017 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/>.

# YL 16/03/2021: we can no longer directly use `from .. import CHEMSH_DEBUG` which is not set at initialisation
import chemsh
from ..dl_py2f.utils import fileutils, miscutils

@miscutils.time
def relaxShells(qmmmObj, dryrun=False):
    '''Iteratively relax shells until converged'''

    from sys     import stdout
    from numpy   import copy   as np_copy
    from numpy   import square as np_square
    from numpy   import sum    as np_sum
    from numpy   import count_nonzero

    _parts = { qmmmObj.qm:qmmmObj.mm,
               qmmmObj.mm:qmmmObj.qm }

    # calculate MM charges first if not yet available
    if not count_nonzero(qmmmObj.mm.frag.charges):
        print("\n Calculating MM charges...")
        qmmmObj.mm.run(task='charges', _energy=False, _gradients=False, _shells=False, dryrun=dryrun)

    if qmmmObj.shl_maxcycles:
        print("\n Self-Consistent Relaxation of Shell Positions")
        print(" "+"="*70)
#        print(" "+">"*70)

    def _printLeavingBanner():
#        print(" "+"<"*70)
        print(" End of Shell Relaxation")
        print(" "+"="*70+"\n")

    archivedir = '_shells_relaxation'

    if chemsh.CHEMSH_DEBUG >= 3:
#        fileutils.archiveDir(qmmmObj._qm._underscored_name)
#        fileutils.archiveDir(qmmmObj._mm._underscored_name)
        fileutils.archiveDir(archivedir)

    fmt = '.%0'+str(len(str(qmmmObj.shl_maxcycles)))+'d'

    for cycle in range(qmmmObj.shl_maxcycles):

        if chemsh.CHEMSH_DEBUG >= 4:
            print("\n >>> Original shells (a.u.):\n")
            print(qmmmObj.frag.shells.coords2str(atoms='names'))

        # calculate QM to get gradients (note: `qm` rather than `_qm` to make everything up-to-date)
        # restart is False if cycle 0; or is True if cycle > 0 or forced to restart
        restart = cycle > 0 or qmmmObj.qm.restart
        # YL 25/04/2020: to match Tcl-ChemShell we reduce the threshold of SCF convergence for shells relaxation
        #                FIXME "scf_threshold" is now a GAMESS-UK keyword needing harmonising with other interfaces
        ierror = qmmmObj.qm.run(_gradients='shells', restart=restart, dryrun=dryrun, scf_threshold=5)

        # YL 24/04/2020: turn on restart after a successful run for shell relaxation
        qmmmObj.qm.restart = True
        stdout.flush()

        # archive input and output files when successfully run
        if chemsh.CHEMSH_DEBUG >= 3:
            fileutils.archiveFile([qmmmObj._qm.input,
                                   qmmmObj._qm.output,
                                   qmmmObj._qm.punch]+
                                   qmmmObj._qm._extra_outputs,
                                   dir=archivedir,
                                   suffix=fmt%cycle)

        # save old shells positions
        abuff = np_copy(qmmmObj.frag.shells.coords)

        # run MM to update shell positions (note: `mm` rather than `_mm` to make everything up-to-date)
        ierror = qmmmObj.mm.run(_shells=True, dryrun=dryrun)
        stdout.flush()

        # archive input and output files when successfully run
        if chemsh.CHEMSH_DEBUG >= 3:
            from os import sep
            fileutils.archiveFile([qmmmObj._mm.input,
                                   qmmmObj._mm.output]+
                                   qmmmObj._mm._extra_outputs,
                                   dir=archivedir,
                                   suffix=fmt%cycle)
            qmmmObj.frag.save(archivedir+sep+'_shells_relaxation.pun'+fmt%cycle, format='punch')

        # calculate RMS
        abuff -= qmmmObj.frag.shells.coords

        if chemsh.CHEMSH_DEBUG >= 3:
            print("\n >>> Relaxed shells (a.u.):\n")
            print(qmmmObj.frag.shells.coords2str(atoms='names'))

        RMS = (np_sum(np_square(abuff))/qmmmObj.frag.nshells)**0.5

        print('\n Shells relaxation cycle {}'.format(cycle+1))
        print(' RMS      : {:.9f} a.u.'.format(RMS))
        print(' Tolerance: {:.9f} a.u.'.format(qmmmObj.shl_tolerance))
        stdout.flush()

        if RMS <= qmmmObj.shl_tolerance:
            print("\n Shells relaxation has converged.\n")
            _printLeavingBanner()

            # "restart" defaults to True after successful shells relaxation
            # TODO: allow to disable restart in any case
            qmmmObj.restart = True

            return 0

    # a warning if relaxation doesn't converge
    if qmmmObj.shl_maxcycles:
        print("\n Maximum number of cycles reached, however shells relaxation did NOT converge!\n")
        _printLeavingBanner()


