
from . import rich

import re

# ==================================================================================================

def colors(theme=None):
  r'''
Return dictionary of colors.

.. code-block:: python

  {
    'selection' : '...',
    'free'      : '...',
    'error'     : '...',
    'warning'   : '...',
    'low'       : '...',
  }

:options:

  **theme** ([``'dark'``] | ``<str>``)
    Select color-theme.
  '''

  if theme == 'dark':
    return \
      {
        'selection' : '1;32;40',
        'free'      : '1;32',
        'error'     : '9;31',
        'warning'   : '1;31',
        'low'       : '1;36',
      }

  return \
  {
    'selection' : '',
    'free'      : '',
    'error'     : '',
    'warning'   : '',
    'low'       : '',
  }

# ==================================================================================================

def read(data=None):
  r'''
Read ``sinfo -o "%all"``.

:options:

  **data** (``<str>``)
    For debugging: specify the output of ``sinfo -o "%all"`` as string.

:returns:

  **lines** ``<list<dict>>``
    A list of dictionaries, that contain the different fields. All data are strings.
  '''

  import subprocess

  # get live info
  if data is None:
    data = subprocess.check_output('sinfo -o "%all"',shell=True).decode('utf-8')

  # extract the header and the info
  head,data = data.split('\n',1)
  data      = list(filter(None,data.split('\n')))

  # get the field-names
  head = head.split('|')

  # convert to list of dictionaries
  # - initialize
  lines = []
  # - loop over lines
  for line in data:
    # -- initialize empty dictionary
    info = {}
    # -- fill dictionary
    for key,val in zip(head,line.split('|')):
      if len(key.strip()) > 0:
        info[key.strip()] = val.strip()
    # -- store to list of lines
    lines += [info]

  # return output
  return lines

# ==================================================================================================

def cpu_score(CPU_LOAD,CPUS_A,**kwargs):

  if not CPU_LOAD.isnumeric(): return rich.Float('')
  if int(CPUS_A) <= 0        : return rich.Float('')

  return rich.Float(float(CPU_LOAD)/float(CPUS_A), precision=2)

# ==================================================================================================

def mem_score(MEMORY,FREE_MEM,CPUS_A,CPUS_T,**kwargs):

  if not FREE_MEM.isnumeric(): return rich.Float('')
  if not CPUS_T  .isnumeric(): return rich.Float('')
  if int(MEMORY) <= 0        : return rich.Float('')
  if int(CPUS_A) <= 0        : return rich.Float('')

  MEM_USED = float(MEMORY) - float(FREE_MEM)

  return rich.Float(float(MEM_USED)/float(MEMORY) * float(CPUS_T)/float(CPUS_A), precision=2)

# ==================================================================================================

def interpret(lines, theme=colors()):
  r'''
Interpret the output of ``GooseSLURM.sinfo.read``. All fields are converted to the
``GooseSLURM.rich`` classes adding useful colors in the process.

:arguments:

  **lines** ``<list<dict>>``
    The output of ``GooseSLURM.sinfo.read``

:options:

  **theme** (``<dict>``)
    The color-theme, as selected by ``GooseSLURM.sinfo.colors``.

:returns:

  **lines** (``<list<dict>>``)
    A list of dictionaries, that contain the different fields. All data are
    ``GooseSLURM.rich.String`` or derived types.
  '''

  # interpret input as list
  if type(lines) != list: lines = [lines]

  # loop over all lines
  for line in lines:

    # convert fields to string
    for key in line:
      if not isinstance(line[key], rich.String):
        line[key] = rich.String(line[key])

	# covert to float
    for key in ['CPU_LOAD']:
      line[key] = rich.Float(str(line[key]))

    # "days-hours:mins:secs" (e.g. "1-4:18:13") -> seconds
    for key in ['TIMELIMIT']:
      line[key] = rich.Duration(str(line[key]))

    # convert memory (e.g. "4G") -> bytes
    for key in ['MEMORY','FREE_MEM']:
      line[key] = rich.Memory(str(line[key]), default_unit=1e6)

    # CPUs
    # - split: allocated/idle/other/total
    a,i,o,t = str(line['CPUS(A/I/O/T)']).split('/')
    # - store extra fields
    line['CPUS_A'] = rich.Integer(a)
    line['CPUS_I'] = rich.Integer(i)
    line['CPUS_O'] = rich.Integer(o)
    line['CPUS_T'] = rich.Integer(t)

    # compute scores
    line['CPU_RELJOB'] = cpu_score(**line)
    line['MEM_RELJOB'] = mem_score(**line)

    # default: CPUs down/online
    line['CPUS_D'] = rich.Integer(0)
    line['CPUS_O'] = rich.Integer(str(line['CPUS_T']))

    # node down: mark all fields
    if re.match(r'down,*' ,str(line['STATE'])) or \
       re.match(r'maint.*',str(line['STATE'])) or \
       re.match(r'drain.*',str(line['STATE'])):

      line['CPUS_I'] = rich.Integer(0)
      line['CPUS_D'] = rich.Integer(str(line['CPUS_T']))
      line['CPUS_O'] = rich.Integer(0)

      for key in line:
        line[key].color = theme['error']

    # highlight 'scores'
    if   int  (line['CPUS_I']    ) > 0   : line['CPUS_I'    ].color = theme['free'   ]
    if   float(line['CPU_RELJOB']) > 1.05: line['CPU_RELJOB'].color = theme['warning']
    elif float(line['CPU_RELJOB']) < 0.95: line['CPU_RELJOB'].color = theme['low'    ]

    # highlight 'scores'
    if float(line['MEM_RELJOB']) > 0.9: line['MEM_RELJOB'].color = theme['warning']

  return lines

# ==================================================================================================

def read_interpret(data=None, theme=colors()):
  r'''
Read and interpret ``sinfo -o "%all"``.

:returns:

  **lines** (``<list<dict>>``)
    A list of dictionaries, that contain the different fields. All data are
    ``GooseSLURM.rich.String`` or derived types.
  '''

  return interpret(read(data), theme)

# ==================================================================================================
