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

__author__ = 'you.lu@stfc.ac.uk'

from . import _Tasks

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

    _kwargs = {
                'coord' :[],
                'theory': None,
                'range' : 0.1,
#                'result':[],
                'step'  : 0.01,
              }

    def run(self, taskfarm=True, **kwargs):
        ''''''


        selectcases = {
                        2: self.scanBond,
                        3: self.scanAngle,
                        4: self.scanDihedral,
                      }

        return selectcases.get(len(self.coord), self.raiseError)(taskfarm=taskfarm, **kwargs)


    def scanBond(self, **kwargs):
        ''''''


    def scanAngle(self, **kwargs):
        ''''''


    def scanDihedral(self, taskfarm=True, **kwargs):
        '''I-J-K-L'''

        from numpy      import arange, array_split, asarray, cross, dot, deg2rad, cos, sin, copy
        from ..parallel import parallel

        # entry to sync
        entry = 'energies'

        # handles
        theory = self.theory
        frag   = theory.frag
        iatom, jatom, katom, latom = self.coord[:]

        # original dihedral (in rad)
        rad_0 = deg2rad(frag.dihedral(iatom, jatom, katom, latom))

        # range (in rad)
        rad = arange(rad_0*(1.0-self.range), rad_0*(1.0+self.range), rad_0*self.step)

        # save original positions of atoms k and l
        coords_k0 = copy(frag.coords[katom])
        coords_l0 = copy(frag.coords[latom])

        nfreedoms = len(rad)
        getattr(self.result, entry).resize(nfreedoms)
#        entryRef.resize(nfreedoms)
        energies = [] 

        # task-farming
        # workgroup ID and number of workgroups
        if taskfarm:
            workgroup   = parallel.workgroup
            nworkgroups = parallel.nworkgroups
        else:
            workgroup   = 0
            nworkgroups = 1

        if taskfarm:
            # array entry to sync
            self.result.title = entry

            # pass `result` to Fortran
            parallel.createPyObject(self.result,
                                    nfreedoms,
                                    0,
                                    0)

        one_dim_indices = arange(nfreedoms)
        indices_of_wg = array_split(one_dim_indices, nworkgroups)[workgroup]

        for theta in rad_0-rad[indices_of_wg]:

            # unit vector k & v, respectively
            uvec_jk = frag.unitvec(jatom, katom)
            uvec_kl = frag.unitvec(katom, latom)
            # Rodrigues' rotation formula: v' = v*cos(theta) + (k x v)*sin(theta) + k*(k dot v)*(1-cos(theta))
            uvec_kl_rot = uvec_kl*cos(theta) + cross(uvec_jk, uvec_kl)*sin(theta) + uvec_jk*dot(uvec_jk, uvec_kl)*(1-cos(theta))
            # get vector from unit vector
            vec_kl_rot  = uvec_kl_rot*frag.dist(katom, latom)
            # get displaced position of atom l
            frag.coords[latom] = coords_k0 - vec_kl_rot

            theory.run(**kwargs)

            energies.append(self.result.energy[0])
            # make sure atom l is restored
            frag.coords[latom] = coords_l0


        getattr(self.result, entry)[indices_of_wg] = energies[:]
#            self.result.append(theory._result.energy)

        # sync all workgroups
        if taskfarm:
            parallel.syncPyObject(entry=entry)

        return self
        #print("\n >>> ChemShell 1D-scan energies:", energies)



    def raiseError(self, *args):
        ''''''

