#  Copyright (C) 2022 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'
# Jul 2022

import sys
import chemsh

__all__ = 'bse'


class BSE():
    '''BSE type'''


    def __init__(self, name, *args, **kwargs):
        ''''''

        self.name = name


    def run(self, *args, frag=None, strict=True, fmt='', default='3-21g', **kwargs):
        '''Run BSE as a Python module'''
    
        import sys
        from numpy    import argmax, char, unique
        from warnings import catch_warnings, resetwarnings, simplefilter
        from os.path  import abspath, dirname, join
        from difflib  import SequenceMatcher
        from ...data  import atomicdata, references
        from ...io    import banner, format
    
        print(" >>> Using basis set from the Basis Set Exchange (BSE) library.\n")
    
        try:
            from .src.build.lib import basis_set_exchange as bse
            # BSE will search in sys.path for the dat directory
            sys.path.append(join(abspath(dirname(__file__)), 'src', 'build', 'lib', 'basis_set_exchange'))
        except ModuleNotFoundError as e:
            format.printwrap(">>> ChemShell/BSE ERROR: Missing dependencies")
            print(f' {e}')
        except:
            format.printwrap(">>> ChemShell/BSE ERROR: BSE library not available in the current installation. Rebuild ChemShell with: ./setup --bse\n")
#            exit(112)

        # automatically determine the format by the name of the QM code    
        formats = list(bse.writers.write._writer_map.keys())
        fmt = formats[argmax([ SequenceMatcher(None, fmt, f).ratio() for f in formats ])]

        strbuff = '\n'
        # loop over atom species
        atom_names_unique, atom_names_unique_indices = unique(char.strip(char.decode(frag.names)), return_index=True)
        for i, atom_name in enumerate(atom_names_unique):

            # default basis set name
            name = default

            znum = frag.znums[atom_names_unique_indices[i]]

            if type(self.name) is str:
                name = self.name
            elif type(self.name) is dict:
                for key, val in self.name.items():
                    # str type (atom name)
                    try:
                        # strictly match (default)
                        if strict:
                            if atom_name.lower() == key.strip().lower():
                                name = val
                        # loosely match (not recommended)
                        else:
                            if atom_name.lower().startswith(key.strip().lower()):
                                name = val
                    # int type (znum)
                    except AttributeError:
                        if znum == key:
                            name = val

            # run the main BSE routine to get a returned string of basis set
            returned_str = bse.get_basis(name, elements=[znum], fmt=fmt, header=True)

            # element name from znum
            element_name = atomicdata.symbols_lower[znum]

            # post-process the str to replace with the correct atom names
            lines = returned_str.splitlines()
            for iline, line in enumerate(lines):
                # skip comment lines
                if line.strip().startswith('#'):
                    strbuff += line.strip() + '\n'
                    continue
                # skip blank lines
                if not line.strip():
                    continue
                # YL FIXME 25/07/2022: there is a risk when sulpher is the element with name, e.g., S1, the "S" (orbital label) will be replaced!
                index_alpha = line.lower().find(element_name)
                # not found in the current line or no whole word matching
                if index_alpha == -1 or not element_name in line.lower().split():
                    strbuff += line + '\n'
                    continue
                # restruct the line with the actual atom name
                index_omega = index_alpha + len(element_name)
                strbuff += line[:index_alpha] + atom_name + line[index_omega:] + '\n'

#            strbuff += returned_str
        strbuff += '\n'

        return strbuff
    
    
