#  Copyright (C) 2021 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 sys   import modules
from numpy import array, char, divide, rint


class __Wrapper(object):
    '''A wrapper class that will be disguised as a module'''


    @property
    def scontrol_show_partition(self):
        '''Output of command `scontrol show partition` as an array of lines of split str'''

        from subprocess import check_output
   
        try:
            return self._scontrol_show_partition
        except:
            self._scontrol_show_partition = array(check_output([ 'scontrol', 'show', 'partition' ]).decode().split())
            return self._scontrol_show_partition


    @property
    def scontrol_show_reservation(self):
        '''Output of command `scontrol show reservation` as an array of lines of split str'''

        from subprocess import check_output
   
        try:
            return self._scontrol_show_reservation
        except:
            self._scontrol_show_reservation = array(check_output([ 'scontrol', 'show', 'reservation' ]).decode().split())
            return self._scontrol_show_reservation


    @property
    def sacctmgr_show_qos(self):
        '''Output of command `sacctmgr show qos` as an list of lines of split str'''

        from subprocess import check_output
   
        try:
            return self._sacctmgr_show_qos
        except:
            self._sacctmgr_show_qos = check_output([ 'sacctmgr', 'show', 'qos' ]).decode().splitlines()
            return self._sacctmgr_show_qos


    @property
    def qos(self):
        '''A array of available quality of service (QoS)'''

        try:
            return self._qos
        except:
            self._qos = char.partition(char.strip((self.sacctmgr_show_qos[2:])), ' ')[:,0]
            return self._qos


    def getPartitionFromKey(self, key):
        '''Return array of values for a given key'''

        if not key.endswith('='):
            key = key.strip() + '='

        masks = char.startswith(self.scontrol_show_partition, key)
        lines = self.scontrol_show_partition[masks]

        return char.lstrip(lines, key)


    def getReservationFromKey(self, key):
        '''Return array of values for a given key'''

        if not key.endswith('='):
            key = key.strip() + '='

        masks = char.startswith(self.scontrol_show_reservation, key)
        lines = self.scontrol_show_reservation[masks]

        return char.lstrip(lines, key)


    @property
    def partitions(self):
        '''A array of all partitions'''

        try:
            return self._partitions
        except:
            self._partitions = self.getPartitionFromKey('PartitionName=')
            return self._partitions


    @property
    def reservations(self):
        '''A array of all reservations'''

        try:
            return self._reservations
        except:
            self._reservations = self.getReservationFromKey('ReservationName=')
            return self._reservations


    @property
    def max_partition_str_len(self):
        '''Maximum string len of partition name'''

        return char.str_len(self.partitions).max()


    @property
    def walltimes(self):
        '''A array of all maximum walltimes'''

        try:
            return self._walltimes
        except:
            self._walltimes = self.getPartitionFromKey('MaxTime=')
            return self._walltimes


    @property
    def sinfo_show_features(self):
        '''Output of command `scontrol show reservation` as an array of lines of split str'''

        from subprocess import check_output
   
        try:
            return self._sinfo_show_features
        except:
            # this command can only be run once because every time it outputs result in a radom order
            self._sinfo_show_features = check_output([ 'sinfo', '-o', '"%N  %c  %f"' ]).decode().replace('"', '')
            return self._sinfo_show_features


    # this is just a mean value
    @property
    def ncpus_per_node(self):
        '''A array of all number of CPU cores per node'''

        try:
            return self._ncpus_per_node
        except:
            ncpus  = self.getPartitionFromKey('TotalCPUs=').astype(int)
            nnodes = self.getPartitionFromKey('TotalNodes=').astype(int)
            self._ncpus_per_node = rint(ncpus/nnodes).astype(int)
            return self._ncpus_per_node


    @property
    def ncpus_per_subsection(self):
        '''Number of CPUs from output of command `sinfo -o "%N  %c  %f"` as an array'''

        from numpy      import genfromtxt
        from io         import StringIO
   
        try:
            return self._ncpus_per_subsection
        except:
            with StringIO(self.sinfo_show_features) as sbuff:
                abuff = genfromtxt(sbuff, skip_header=1, dtype=int)
                if len(abuff.shape) > 1:
                    self._ncpus_per_subsection = abuff[:,1]
                # YL 13/06/2023: 0-dim array
                elif len(abuff.shape) == 1:
                    self._ncpus_per_subsection = abuff[None,:][:,1]
                else:
                    return []
            return self._ncpus_per_subsection


    @property
    def subsections(self):
        '''Subsections from output of command `sinfo -o "%N  %c  %f"` as an array'''

        from numpy      import genfromtxt
        from io         import StringIO
   
        try:
            return self._subsections
        except:
            with StringIO(self.sinfo_show_features) as sbuff:
                abuff = genfromtxt(sbuff, skip_header=1, dtype=str)
                if len(abuff.shape) > 1:
                    self._subsections = abuff[:,2]
                # YL 13/06/2023: 0-dim array
                elif len(abuff.shape) == 1:
                    self._subsections = abuff[None,:][:,2]
                else:
                    return []
            return self._subsections


    @property
    def nodes_per_subsection(self):
        '''Nodes list of each subsection from output of command `sinfo -o "%N  %c  %f"` as an array'''

        from numpy      import array, genfromtxt
        from io         import StringIO
   
        try:
            return self._nodes_per_subsection
        except:
            with StringIO(self.sinfo_show_features) as sbuff:
                abuff = genfromtxt(sbuff, skip_header=1, dtype=str)
                if len(abuff.shape) > 1:
                    self._nodes_per_subsection = char.strip(abuff[:,0], '"')
                # YL 13/06/2023: 0-dim array
                elif len(abuff.shape) == 1:
                    self._nodes_per_subsection = char.strip(abuff[None,:][:,0], '"')
                else:
                    return array([], dtype=str)
            return self._nodes_per_subsection


    @property
    def partitions2walltimes(self):
        '''A dictionary of partition to maximum walltime mapping'''

        try:
            return self._partitions2walltimes
        except:
            self._partitions2walltimes = dict(zip(self.partitions, self.walltimes))
            return self._partitions2walltimes


    @property
    def partitions2ncpus_per_node(self):
        '''A dictionary of partition mapped to number of CPUs per node'''

        try:
            return self._partitions2ncpus_per_node
        except:
            self._partitions2ncpus_per_node = dict(zip(self.partitions, self.ncpus_per_node))
            return self._partitions2ncpus_per_node


    @property
    def subsections2ncpus_per_node(self):
        '''A dictionary of subsection (constraint/feature) mapped to number of CPUs per node'''

        try:
            return self._subsections2ncpus_per_node
        except:
            self._subsections2ncpus_per_node = dict(zip(self.subsections, self.ncpus_per_node))
            return self._subsections2ncpus_per_node


modules[__name__] = __Wrapper()
