#  Copyright (C) 2019 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 01/05/2019

'''Superposition of fragment molecular orbitals
See: http://www.nwchem-sw.org/index.php/Release66:Hartree-Fock_Theory_for_Molecules#Superposition_of_fragment_molecular_orbitals
Step A: get a 
'''

from sys        import modules
from ...dl_py2f import utils

movecsModule = modules[__name__]


def sfmoType(NWChemObj):
    '''Return True if broken down into single atoms'''

    if type(NWChemObj.fragmovecs) is str:
        if NWChemObj.fragmovecs.lower() in [ 'atoms', 'individual' ]:
            return 'atoms'
        elif NWChemObj.fragmovecs.lower() in [ 'sections' ]:
            return 'sections'
        elif NWChemObj.fragmovecs.lower() in [ 'species', 'elements' ]:
            return 'species'
    else:
        try:
            # get a list containing the user-defined sections, such as [ [5,7,...], [77,78,...] ]
            setattr(movecsModule, 'sections', [utils.iterutils.getFlattened(list(s)) for s in iter(NWChemObj.fragmovecs)])
            # a list of atoms not included in the user-defined list of list(s)
            notincluded = set(range(NWChemObj.natoms)) - set(utils.iterutils.getFlattened(list(movecsModule.sections)))
            movecsModule.sections += [ [i] for i in notincluded ]
            setattr(NWChemObj, '_fragmovecs', NWChemObj.fragmovecs)
            NWChemObj.fragmovecs = 'sections'
            return 'sections'
        except TypeError:
            return 'atoms'


def getReorderedAtoms(NWChemObj):
    ''''''

    from numpy import argsort, unique

    # recommended: break down into every individual atom
    if sfmoType(NWChemObj) == 'atoms' or 'species':

        return NWChemObj.frag.coords2str(atoms='names', charges=None, coords=14)

    elif sfmoType(NWChemObj) == 'sections':

        strbuff = ''
        for section in movecsModule.sections:
            strbuff += NWChemObj.frag.coords2str(select=section, atoms='names', charges=None, coords=14)
        return strbuff

    # grouped into species
    else:

        fragObj = NWChemObj.frag
    
        indices = argsort(fragObj.names)
    
        setattr(movecsModule, 'indices', indices)
    
        sets, set_starts = unique(fragObj.names[indices], return_index=True)
    
        setattr(movecsModule, 'sets'      , sets)
        setattr(movecsModule, 'set_starts', set_starts)
    
        return NWChemObj.frag.coords2str(select=indices, atoms='names', charges=None, coords=14)


def getFragments(NWChemObj):
    '''Split into fragments according to fragObj.names'''

    fragObj = NWChemObj.frag

    strbuff = ''

    # recommended: break down into every individual atom
    if sfmoType(NWChemObj) == 'atoms':

        # geometry blocks containing a single atom each
        for i in range(fragObj.natoms):
            frag_id = '{:03d}{}'.format(i, fragObj.names[i].decode())
            strbuff += 'geometry _{}\n'.format(frag_id)
            strbuff += '    symmetry c1\n'
            strbuff += fragObj.coords2str(select=[i], atoms='names', charges=None, coords=14)
            strbuff += 'end\n'

    #
    elif sfmoType(NWChemObj) == 'sections':
        # geometry blocks containing a section each
        for isec, section in enumerate(movecsModule.sections):
            frag_id = '{:03d}'.format(isec)
            strbuff += 'geometry _{}\n'.format(frag_id)
            strbuff += '    symmetry c1\n'
            strbuff += fragObj.coords2str(select=section, atoms='names', charges=None, coords=14)
            strbuff += 'end\n'

    # grouped into species
    elif sfmoType(NWChemObj) == 'species':

        from numpy import argsort, unique

#        indices = argsort(fragObj.names)
#        setattr(movecsModule, 'indices', indices)
        sets, set_starts = unique(fragObj.names, return_index=True)
        setattr(movecsModule, 'sets'      , sets)
        setattr(movecsModule, 'set_starts', set_starts)
    
        # fragments
        for i, name in enumerate(sets):
            strbuff += 'geometry _{} print units bohr nocenter noautoz\n'.format(name.decode())
            strbuff += '    symmetry c1\n'
            strbuff += '    {:6} {:20.14f} {:20.14f} {:20.14f}\n'.format(name.decode(), 0.0, 0.0, 0.0)
            strbuff += 'end\n'
    else:
        strbuff += '\n'
    
    return strbuff
    

def getSCFs(NWChemObj):
    '''SCF tasks of each fragment'''

    from ...base import forcefield
    from ...data import atomicdata

    fragObj = NWChemObj.frag
    strbuff = 'scf\n    thresh {:E}\nend\n\n'.format(NWChemObj.tol_pre_hf)

    # recommended: break down into every individual atom
    if sfmoType(NWChemObj) == 'atoms':

        for i in range(fragObj.natoms):

            formal_charge = NWChemObj.formal_charges.get(fragObj.names[i].decode(), atomicdata.formal_charges[fragObj.znums[i]])

            frag_id = '{:03d}'.format(i)

            strbuff += 'charge {}\n'.format(formal_charge)
            strbuff += 'set geometry _{}{}\n'.format(frag_id, fragObj.names[i].decode())
            strbuff += 'scf\n'
            strbuff += '    maxiter {}\n'.format(NWChemObj.maxiter)
            if NWChemObj.direct:
                strbuff += '    direct\n'
            if NWChemObj.mult == 1:
                strbuff += '    singlet\n'
                strbuff += '    {}\n'.format(NWChemObj.scftype)
            elif NWChemObj.mult == 2:
                strbuff += '    doublet\n'
                strbuff += '    {}\n'.format(NWChemObj.scftype)
            elif NWChemObj.mult == 3:
                strbuff += '    triplet\n'
                strbuff += '    {}\n'.format(NWChemObj.scftype)
            strbuff += '    vectors input atomic output {}\n'.format(frag_id)
            strbuff += 'end\n'
            strbuff += 'task scf energy\n\n'

    elif sfmoType(NWChemObj) == 'sections':

        for isec, section in enumerate(movecsModule.sections):

            formal_charge = 0
            for iatom in section:
                formal_charge += NWChemObj.formal_charges.get(fragObj.names[iatom].decode(), atomicdata.formal_charges[fragObj.znums[iatom]])

            frag_id = '{:03d}'.format(isec)

            strbuff += 'charge {}\n'.format(formal_charge)
            strbuff += 'set geometry _{}\n'.format(frag_id)
            strbuff += 'scf\n'
            strbuff += '    maxiter {}\n'.format(NWChemObj.maxiter)
            if NWChemObj.mult == 1:
                strbuff += '    singlet\n'
                strbuff += '    {}\n'.format(NWChemObj.scftype)
            elif NWChemObj.mult == 2:
                strbuff += '    doublet\n'
                strbuff += '    {}\n'.format(NWChemObj.scftype)
            elif NWChemObj.mult == 3:
                strbuff += '    triplet\n'
                strbuff += '    {}\n'.format(NWChemObj.scftype)
            strbuff += '    vectors input atomic output {}\n'.format(frag_id)
            strbuff += 'end\n'
            strbuff += 'task scf energy\n\n'

    elif sfmoType(NWChemObj) == 'species':

        for i in range(movecsModule.sets.size):

            pos_in_all = movecsModule.set_starts[i]

            name = fragObj.names[pos_in_all].decode()
            znum = forcefield.names2znums([name])[0]

            formal_charge = NWChemObj.formal_charges.get(name, atomicdata.formal_charges[fragObj.znums[pos_in_all]])

            frag_id = '{:03d}'.format(i)

            strbuff += 'charge {}\n'.format(int(formal_charge))
            strbuff += 'set geometry _{}\n'.format(name)
            strbuff += 'scf\n'
            strbuff += '    maxiter {}\n'.format(NWChemObj.maxiter)
            if NWChemObj.direct:
                strbuff += '    direct\n'
            
            mult = (znum - int(formal_charge))%2 + 1
            if mult == 1:
                strbuff += '    singlet\n'
                strbuff += '    {}\n'.format(NWChemObj.scftype)
            elif mult == 2:
                strbuff += '    doublet\n'
                strbuff += '    {}\n'.format(NWChemObj.scftype)
            elif mult == 3:
                strbuff += '    triplet\n'
                strbuff += '    {}\n'.format(NWChemObj.scftype)
            strbuff += '    vectors input atomic output {}\n'.format(name)
            strbuff += 'end\n'
            strbuff += 'task scf energy\n\n'


    # grouped into species
    else:

        for i, frag in enumerate(movecsModule.sets):
            
            strbuff += 'charge {}\n'.format(charges[i])
            strbuff += 'set geometry {}\n'.format(frag.decode())
            strbuff += 'scf\n'
            strbuff += '    vectors input atomic output {}.mo\n'.format(frag.decode())
            strbuff += 'end\n'
            strbuff += 'task scf energy\n\n'

    strbuff += 'unset scf:*\n\n'
    
    return strbuff


def getPreSCF(NWChemObj):
    '''A full SCF combining all fragments'''

    fragObj = NWChemObj.frag

    # a full list of all fragment MOs files
    # NB: NWChem reads in only a limited length
    strbuff = '    vectors input fragment'

    # recommended: break down into every individual atom
    if sfmoType(NWChemObj) == 'atoms':
        lmovecs = ''
        for i in range(fragObj.natoms):
            frag_id = '{:03d}'.format(i)
            lmovecs += ' {}'.format(frag_id)
        # NWChem cannot read in too many characters (< ~998) from this list of files
        # then we support < 249 atoms
        if len(lmovecs) > 998:
            print("\n >>> ChemShell/NWChem: cannot calculate so many fragment MOs")
            exit(999)
        strbuff += lmovecs

    elif sfmoType(NWChemObj) == 'sections':
        lmovecs = ''
        for isec, section in enumerate(movecsModule.sections):
            frag_id = '{:03d}'.format(isec)
            lmovecs += ' {}'.format(frag_id)
        # NWChem cannot read in too many characters (< ~998) from this list of files
        # then we support < 249 atoms
        if len(lmovecs) > 998:
            print("\n >>> ChemShell/NWChem: cannot calculate so many fragment MOs")
            exit(999)
        strbuff += lmovecs

    elif sfmoType(NWChemObj) == 'species':
        lmovecs = ''
        for i in range(fragObj.natoms):
            lmovecs += ' {}'.format(fragObj.names[i].decode())
        # NWChem cannot read in too many characters (< ~998) from this list of files
        # then we support < 249 atoms
        if len(lmovecs) > 998:
            print("\n >>> ChemShell/NWChem: cannot calculate so many fragment MOs")
            exit(999)
        strbuff += lmovecs

    # grouped into species
    else:
        for i, frag in enumerate(movecsModule.sets):
            strbuff += ' {}.mo'.format(frag.decode())

#    strbuff += ' \\\n'
    strbuff += ' output _{}_prescf.movecs\n'.format(fragObj.formula)

    return strbuff









