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

from ..qm        import _QM
from ...parallel import parallel
from ...dl_py2f  import utils

class Gaussian(_QM):
    '''Interface to Gaussian'''

    _kwargs = {
                'check'    : '_gaussian.chk',
                'input'    : '_gaussian.inp',
                'path'     : 'g16',
                'output'   : '_gaussian.log',  # hardcoded as ".log" by Gaussian
                'threshold': 1.0E-8,
              }

    _internals = {}


    def getInpBuff(self):
        ''''''

        from math        import log10
        from numpy       import unique

#        import platform
#        import socket
        from sys import stdout

        # archive old files
        utils.fileutils.archiveFile([self.input, self.output])

        strbuff = ''

        # shared-memory multiprocessor parallel execution
        if parallel.nprocs_per_wg > 1:
#            strbuff += '%nprocshared={}\n'.format(parallel.nprocs_per_wg)
            strbuff += '{}'.format(self.cpu_indices)

        # cluster/network parallelism
#        hostnames = unique(parallel.workgroup_hostnames)
#        if hostnames.size > 1:
#            lindaworkers = ','.join(hostnames)
#            strbuff += '%lindaworkers={}\n'.format(lindaworkers)

        selectcases = { 'hf'    :'hf',
                        'tdhf'  :'hf',
                        'td-hf' :'hf',
                        'dft'   : self.functional,
                        'tddft' : self.functional,
                        'td-dft': self.functional,
                      }

        gradients = { True: 'force',
                    }

        scftypes = { 'rhf' :     selectcases[self.method.lower()],
                     'uhf' :'u' +selectcases[self.method.lower()],
                     'rohf':'ro'+selectcases[self.method.lower()],
                   }

        threshold = 'scf=(conver={}, maxcycle={})'.format(-int(log10(self.threshold)), self.maxiter)

        # time-dependent methods
        if self._time_dependent:
            if self.restart:
                td_restart = ', read, restart'
            else:
                td_restart = ''
            td = 'td(nstates={}, {}, root={}, conver={}{}, maxcycle=100) '.format(self.nstates,
                                                                                  self.spin.lower(),
                                                                                  self.target_state,
                                                                                  -int(log10(self.threshold)),
                                                                                  td_restart)
        else:
            td = ''

        # uhf/rhf/rohf
        try:
            method = scftypes.get(self.scftype) or selectcases[self.method.lower()]
        except KeyError:
            print(" >>> ChemShell ERROR: unknown method \"{}\"".format(self.method.lower()))
            raise

        # check file
        strbuff += '%chk={}\n'.format(self.check)

        # method and basis set
        strbuff += '#p {}{}/{}\n{}\n{}\n'.format(td, method, self.basis, gradients.get(self._gradients, ''), threshold)

        # integration grid
        selectcases = { 'ultrafine' :'ultrafine',
                        'ultra fine':'ultrafine',
                        'ultra-fine':'ultrafine',
                        'ultra_fine':'ultrafine',
                        'medium'    :'ultrafine',
                        'superfine' :'superfine',
                        'super fine':'superfine',
                        'super-fine':'superfine',
                        'super_fine':'superfine',
                        'high'      :'superfine',
                      }

        strbuff += 'integral(grid={})\n'.format(selectcases.get(self.grid.strip().lower(), 'ultrafine'))

        # restart
        # NB: do NOT use chkbasis!
        if self.restart:
            strbuff += 'guess=read charge=read\n'

        # title line
        strbuff += '\n{}\n'.format(self.title)

        # charge and multiplicity
        strbuff += '\n{} {}\n'.format(self.charge, self.mult)

        # coordinates
        strbuff += self.frag.coords2str(charges=None, unit='angstrom')

        # a blank line
        strbuff += '\n\n'

        return strbuff


    @property
    def _outkeys(self):
        '''Define {task:keyword} for searching in output file'''

        # each keyword should contain an array of integers specifying:
        #     [ row_0, # starting row
        #       row_n, # ending row
        #       row_i, # row increment
        #       col_0, # starting col
        #       col_n, # ending col
        #       col_i, # col increment
        #       skip , # skip the previous (skip-1) match,
        #       coeff, # unit conversion coefficient
        #     ]
        # referencing to the line where the key is matched
        _outkeys = {
                    'energy'     :{'SCF Done:'                  : [0,0,1,4,4,1,0,1.0]},
                    'e_root'     :{'This state for optimization': [1,1,1,4,4,1,0,1.0]},
                    # apply -1.0 because Gaussian prints "forces" rather than "gradients"
                    'gradients'  :{'Forces \\(Hartrees/Bohr\\)' : [3,3+self.frag.natoms-1,1,2,4,1,0,-1.0]},
                    'ecp_grads'  :{'ENERGY GRADIENTS'           : [4+self.frag.natoms,4+self.frag.natoms+self.necps-1,1,5,7,1,1,1.0]},
                   }

        # multistate calculation
        for i in range(self.nstates):
            # typically: Excited State   1:      Singlet-A1     x.xxxx eV  yyy.yy nm  f=z.zzzz  <S**2>=0.000
            # internally we start from 0 but the Gaussian labels start from 1
            _outkeys.update({'excitation {}'.format(i):{'Excited State   {}'.format(i+1): [0,0,1,3,8,1,0,1.0]}})

        return _outkeys


    def parseMainOutput(self):
        '''Parse main output file self.output'''

        from ...io import file

        # energy
        self._result.energy = file.getArrayFromFile(self.output, 'out', self._result.energy.shape, regex=self._outkeys['energy'])

        # excited state calculations
        for i in range(self.nstates):
            abuff = file.getArrayFromFile( self.output,
                                               'out',
                                                self._result.energy.shape,
                                                regex=self._outkeys['excitation {}'.format(i)],
                                                number_only=False,
                                               )
            self._result.excitations[i] = abuff[1], abuff[5].split('=')[1], abuff[3], abuff[0].split('-')[1].upper(), abuff[0].split('-')[0].lower()

        # in excited state calculation we use E0 for ground state energy and result.energy for the state of interest (root state)
        if self._excited_state:
            # get the energy of the state of interest
            E_root = file.getArrayFromFile(self.output, 'out', self._result.energy.shape, regex=self._outkeys['e_root'])
            # swap result.energy and result.E0 
            self._result.E0     = self._result.energy[0]
            self._result.energy = E_root

        # gradients
        try:
            if self._gradients:
                self._result.gradients = file.getArrayFromFile(self.output, 'out', self._result.gradients.shape, regex=self._outkeys['gradients'])
#                if self.necps and self.separate_ecps:
#                    self.ecps.gradients = file.getArrayFromFile(self.output, 'out', self.ecps.gradients.shape, regex=self._outkeys['ecp_grads'])
        except:
            print("\n >>> ChemShell NWChem: failed parsing gradients from output:\n")
            raise


    # legends: number - world rank; [] - workgroup; row - nodes
    #
    # scenario I: 5 workgroups over 2 nodes containing 10 cores each
    #     node0: [0   1   2   3][ 4   5   6   7][ 8   9
    #     node1: 10  11][12  13  14  15][16  17  18  19]
    # Gaussian keywords should be
    #     wg0: %cpu:0-3\n%lindaworkers=node0
    #     wg1: %cpu:4-7\n%lindaworkers=node0
    #     wg2: %cpu:8-9\n%lindaworkers=node0,node1:2
    #     wg3: %cpu:2-5\n%lindaworkers=node1
    #     wg4: %cpu:6-9\n%lindaworkers=node1
    #
    # scenario II: 2 workgroups over 5 nodes containing 6 cores each
    #     node0: [0   1   2   3   4   5
    #     node1:  6   7   8   9  10  11
    #     node2: 12  13  14][15  16  17
    #     node3: 18  19  20  21  22  23
    #     node4: 24  25  26  27  28  29]
    # Gaussian keywords should be
    #     wg0: %cpu:0-5\n%lindaworkers=node0,node1,node2:3
    #     wg1: %cpu:3-5\n%lindaworkers=node2,node3:6,node4:6
    @property
    def cpu_indices(self):
        ''''''

        from numpy       import array, unique
        from collections import Counter
        from sys import stdout

        world_hostnames     = parallel.world_hostnames
        workgroup_hostnames = parallel.workgroup_hostnames

# uncomment to test
#        # scenario I test
#        world_hostnames = ['node0']*10+['node1']*10
#        if parallel.workgroup in [0,1]:
#            workgroup_hostnames = ['node0']*4
#        if parallel.workgroup in [3,4]:
#            workgroup_hostnames = ['node1']*4
#        if parallel.workgroup in [2]:
#            workgroup_hostnames = ['node0']*2+['node1']*2
#        # scenario II test
#        world_hostnames = ['node0']*6+['node1']*6+['node2']*6+['node3']*6+['node4']*6
#        if parallel.workgroup in [0]:
#            workgroup_hostnames = ['node0']*6+['node1']*6+['node2']*3
#        if parallel.workgroup in [1]:
#            workgroup_hostnames = ['node2']*3+['node3']*6+['node4']*6

        hostnames, nnodes = unique(workgroup_hostnames, return_counts=True)
        nhostnames = hostnames.size
        stdout.flush()

        counter = Counter(workgroup_hostnames)

        # if the current workgroup is contained by a single node
        if nhostnames == 1:
            hostname = hostnames[0]
            offset = world_hostnames.index(hostname)
            strbuff = '%cpu={}\n'.format(','.join(map(str, parallel.world_ranks-offset)))
            strbuff += '%lindaworkers={}\n'.format(hostname)
            return strbuff

        # if the current workgroup spans more than 1 node
        elif nhostnames > 1:
            # starting point of the current workgroup
            offset = world_hostnames.index(hostnames[0])
            # CPU indices of only the first node of the current workgroup
            cpus = parallel.world_ranks[:counter[hostnames[0]]]-offset
#            # to avoid %cpu=0\n%lindaworkers=node1,node2:x which is illegal in Gaussian
#            if cpus.size == 1:
#                cpus = [ cpus[0], cpus[0]+1 ]
            strbuff = '%cpu={}\n'.format(','.join(map(str, cpus)))
            # add the first node of the current workgroup
            l_hostnames = [hostnames[0]]
            # add the rest nodes
            for hostname in hostnames[1:]:
                l_hostnames += ['{}:{}'.format(hostname, counter[hostname])]
            lindaworkers = ','.join(l_hostnames)
            strbuff += '%lindaworkers={}\n'.format(lindaworkers)
            return strbuff


