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

import chemsh
from ...dl_py2f.utils  import fileutils, strutils, miscutils, objutils
from ...base.constants import MAXCONN

# YL: in mode of direct parsing of file (e.g., f = Fragment(coords='frag.punch')), receiver must be an XYZ type
def parse(filename, _regex=None, receiver=None, prefix='', **kwargs):
    '''Parse a Tcl-ChemShell punchfile and return matrix as buff'''

    from io       import StringIO
    from numpy    import array, dtype, float64, int64, loadtxt, fromfile
    from sys      import stdout
    from os       import path
    from ..       import keyword

# FIXME if there is no connectivity information, connectivity should be calculated!
    selectcases = { 'atom_charges'        :  'charges',
                    # note the column-/row-major difference between TCS and PCS!
                    'cell_vectors'        :[ 'cell.vectors' ],            # not yet supports string type 'cell.vectors' without []
                    'connectivity'        :( 'connPairs', lambda x:x-1 ), # funcname as attribute of receiver and treatment of raw data (Fortran indexing to C)
                    'coordinates'         :[ 'names', 'coords' ],
                    'point_charges'       :[ 'bqs.names', 'bqs.coords', 'bqs.charges' ],
                    'shells'              :[ 'shells.names', 'shells.coords', 'shells.charges', 'shells.coreatoms' ],
                    'title'               :  'title',
                    'charge'              :  'totalcharge',
                    'dense_integer_matrix':  '',                            # int64,
                    'dense_real_matrix'   :  '',                            # float64,
                    'matrix'              :  '',                            # type(None),
                    'matrix_title'        :  '',                            # str,
                  }
#                    'gradients'           : 'gradients',
#                    'tr_matrix'           : '',
#                    'point_group'         : 'pgroup',

    # process in particular order: connectivity depends on coordinates
    _priorities = 'coordinates', 'connectivity'


    # parse, e.g.: dictbuff = {'block': 'matrix_title', 'records': '1'}
    def _getBlock(_line):
        ''''''

        if 'block' in _line and '=' in _line:
            dictbuff = {}
            keyword.parseKeyValPairs(dictbuff, _line)
            block    = dictbuff['block'].lower()
            nrecords = int(dictbuff['records'])
            return block, nrecords
        else:
            return '', 0

    # YL 13/08/2022: path prefix
    filename = path.join(prefix, filename)

    with open(filename, 'r') as fp:

        # number of lines
        length = len(fp.read())
        fp.seek(0)

        while fp.tell() < length:

            # a drawback of using this method is that next(fp) cannot be applied in the loop body
            line = fileutils.joinContinuedLines(fp)

            # YL 04/09/2020: skip blanc and comment lines
            if line.strip().startswith('#') or not line.strip():
                continue

            # get block info
            block, nrecords = _getBlock(line)
            if chemsh.CHEMSH_DEBUG >= 5:
                print("\n Parsing punch file... block = {}, nrecords = {}".format(block, nrecords))

            # EOF
            if not block:
                break

            # empty block
            if not nrecords:
                continue


            # if unknown attr name found, return the obj itself
            attr = selectcases.get(block, '')

            # get dtype
            try:
                dt = dtype(objutils.type2dtype(objutils.getAttr(receiver, attr)))
                try:
                    # try to name the fields
                    dt.names = [ strutils.getRightmost(key) for key in attr ]
                except:
                    # pass if not structured
                    pass
            # if not structured, it is safe to use float64 as default as in the NumPy npyio.py source code
            except:
                dt = float64

            # split the whole file into blocks
            # do not dump to file
            with StringIO('') as pun:
                pun.write(line)
                for i in range(nrecords):
                    pun.write(fp.readline())

                # load data into a structured array `databuff`
                try:
                    pun.seek(0)
                    databuff = loadtxt(pun, dtype=dt, skiprows=1)
                except ValueError as emessage:
                    if chemsh.CHEMSH_DEBUG >= 1:
                        print("\n >>> WARNING: failed to load {} records of data buffer as type for block \"{}\":\n".format(nrecords, block), dt)
                    # YL 15/02/2019: added debugging message for it could be difficult for users to find out mistakes in punch
                    if chemsh.CHEMSH_DEBUG >= 3:
                        print("\n {}\n".format(emessage))
                        stdout.flush()
                        # if intruding content found in the supposed length of data block
                        if 'could not convert' in str(emessage):
                            # YL 29/11/2022: the error messaage of NumPy has changed
                            try:
                                intrusive = str(emessage).split(':')[1].strip()
                            except:
                                intrusive = str(emessage).strip().split('\'')[1].split('\n')[0]
                            pun.seek(0)
                            lines = pun.readlines()[1:]
                            for i, l in enumerate(lines):
                                if intrusive in l.split():
                                    print(" record {} of block \"{}\" in file \"{}\":\n {}\n".format(i+1, block, filename, line))
                    # YL 29/03/2018: changed "pass" to "continue" to prevent proceding with wrong databuff
                    continue

                # if block unmatches _regex
                # this `continue` command MUST NOT be brought up to before StringIO dumping which contains `readline()`
                if block in [ 'cell_vectors' ]:
                    pass
                if _regex and block != _regex:
                    continue

                try:
                    objutils.setAttr(receiver, attr, databuff)

                    # special treatment due to the column-/row-major difference between TCS and PCS!
                    if block in [ 'cell_vectors' ]:
                        handle = objutils.getAttr(receiver, attr)
                        handle[0][:] = handle[0].T

                # e.g., may fail to loadtxt as int from block 'matrix_title' and it means the block does not match
                # or `dt` cannot be defined: try to apply a function specified in `selectcases` for special treatment in this case
                except AttributeError:
                    # get function reference from the first part of tuple `attr`
                    func    = getattr(receiver, attr[0])
                    # preprocess the raw data according to the second part of tuple `attr`
                    newdata = attr[1](databuff)
                    # run function
                    func(newdata)

                except:
                    print(" >>> NOTE: ChemShell does not know how to assign data block", '"'+block+'"', "in file", filename, "for", type(receiver))

                    continue

    return receiver


# TODO Check and handle objects other than fragments
# TODO Add missing fragment blocks
#@miscutils.time
def write(filename, frag, **kwargs):
    '''Write out a fragment in punch file format.
    Format specifications are taken from Tcl-ChemShell.'''

    from numpy import count_nonzero

    with open(filename, 'w') as fp:

        # Header
        fp.write("block = fragment records = 0\n")
        fp.write("block = title records = 1\n")
        fp.write(frag.title.rstrip() + "\n")

        # Atom coordinates
        fp.write("block = coordinates records = %d\n" % frag.natoms)
        #print(frag.names)

        # YL 15/01/2018: frag.coords2str() is not faster than explicit loops of writing
        strbuff = frag.coords2str(atoms='names', coords=14, charges=None, notations={'c':'e'})
        fp.write(strbuff)

        # WONTDO Gradients
        # TODO Fractional coordinates
        # TODO Space group
        # AJL, Aug2017: (Total) charge
        fp.write("block = charge records = 1\n")
        fp.write(str(frag.totalcharge) + "\n")

        # Atom charges
        # AJL, Aug2017: Stumbled across this on the to-do list and fixed it
        # TODO Need a charges flag in the fragment object
        #hasCharges = False
        #tiny = 1.0e-10
        #for i in range(frag.natoms):
        #    if abs(frag.charges[i]) > tiny:
        #        hasCharges = True
        #        break
        #if hasCharges:

        if frag.hasCharges():
            fp.write("block = atom_charges records = %d\n" % frag.natoms)
            for i in range(frag.natoms):
                fp.write("%20.10f\n" % frag.charges[i])

        # TODO Fractional atom charges

        # Point charges
        if frag.nbqs > 0:
            fp.write("block = point_charges records = %d\n" % frag.nbqs)
            for i in range(frag.nbqs):
                fp.write("%8s %20.14e %20.14e %20.14e %12.6f \n" % 
                         (frag.bqs.names[i].decode(),
                         frag.bqs.coords[i][0], frag.bqs.coords[i][1], frag.bqs.coords[i][2],
                         frag.bqs.charges[i]))

        # TODO Fractional point charges
        # WONTDO Externals

        # Shells
        if frag.nshells > 0:
            # get wrapped shell data once only
            # TODO delete once fragment object has been refactored for performance
            shells = frag.shells
            fp.write("block = shells records = %d\n" % frag.nshells)
            for i in range(frag.nshells):
                fp.write("%s  %20.14e %20.14e %20.14e %16.8e %d \n" %
                         (shells.names[i].decode(),
                          shells.coords[i][0], shells.coords[i][1], shells.coords[i][2],
                          shells.charges[i], shells.coreatoms[i]))

        # TODO Fractional shells
        # TODO Mapping between primitive cells and clusters and fractional list (xtal_map)

        # Cell vectors
        if frag.cell.natoms:
            fp.write("block = cell_vectors records = %d\n" % frag.cell.ndimensions)
            for i in range(frag.cell.ndimensions):
                fp.write("%20.14e %20.14e %20.14e\n" % (frag.cell.vectors[i,0], frag.cell.vectors[i,1], frag.cell.vectors[i,2]))

        # connectivity
        # YL 19/02/2020: corrected nconns
        nconns = int(count_nonzero(frag.conn+1)/2)
        ndigits = len(str(nconns))
        fmt = '{:>'+'{}'.format(ndigits)+'d}  {:>'+'{}'.format(ndigits)+'d}\n'
        if nconns:
            fp.write("block = connectivity records = %d\n"%nconns)
            for i in range(frag.natoms):
                for j in range(MAXCONN):
                    if frag.conn[i][j] > i:
                        # starting from 1 rather than 0
                        fp.write(fmt.format(i+1, frag.conn[i][j]+1))

        # WONTDO Residue (group) to which atoms belong

