source: TOOLS/ConsoGENCI/trunk/bin/libconso_cpt.py @ 2775

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# ==================================================================== #
# Author: Sonia Labetoulle                                             #
# Contact: sonia.labetoulle _at_ ipsl.jussieu.fr                       #
# Created: 2016                                                        #
# History:                                                             #
# Modification:                                                        #
# ==================================================================== #

# =================================================================== #
# ssh readonly@prodiguer-test-db.ipsl.upmc.fr                         #
# psql -U prodiguer_db_user prodiguer                                 #
#                                                                     #
# ssh readonly@prodiguer-test-db.ipsl.upmc.fr -L 5432:localhost:5432  #
# =================================================================== #


# This must come first
from __future__ import print_function, unicode_literals, division

# Standard library imports
import os
import glob
# import psycopg2
# import psycopg2.extras
import pprint

# Application library imports

pp = pprint.PrettyPrinter(indent=2)


#######################################################################
class ProjectBloc(object):

  #--------------------------------------------------------------------
  def __init__(
    self,
    name,
    center,
    machine=None,
    node=None,
    date_beg=None,
    date_end=None,
    alloc=None,
    line_beg=None,
    line_end=None,
  ):
    self.name     = name
    self.center   = center
    self.machine  = machine
    self.node     = node
    self.date_beg = date_beg
    self.date_end = date_end
    self.alloc    = alloc
    self.line_beg = line_beg
    self.line_end = line_end
    self.alloc_id = None

  #--------------------------------------------------------------------
  def __repr__(self):
    return "{}/{} = {}".format(
      self.name,
      self.center,
      self.alloc,
    )


#######################################################################
class ConsoItem(object):

  #--------------------------------------------------------------------
  def __init__(self, login, date, conso):
    self.login = login
    self.date  = date
    self.conso = conso

  #--------------------------------------------------------------------
  def __repr__(self):
    return "{}/{:%Y%m%d} = {}".format(
      self.login,
      self.date,
      self.conso,
    )


#######################################################################
class AllocRow(object):

  #--------------------------------------------------------------------
  def __init__(self, row):
    self.name     = row["project"]
    self.alloc_id = row["id"]
    self.machine  = row["machine"]
    self.centre   = row["centre"]
    self.node     = row["node_type"]
    self.date_beg = row["start_date"]
    self.date_end = row["end_date"]
    self.alloc    = row["total_hrs"]

  #--------------------------------------------------------------------
  def __repr__(self):
    return "{}/{}/{:%Y%m%d}/{:%Y%m%d} ({})".format(
      self.machine,
      self.node,
      self.date_beg,
      self.date_end,
      self.alloc_id,
    )


#######################################################################
def get_project_id(bloc, allocs):

  res = None

  for alloc in allocs:
    # print(
    #   bloc.cpt_date,
    #   alloc.date_end,
    #   bloc.cpt_date <= alloc.date_end
    # )
    if bloc.machine == alloc.machine   and \
       bloc.node == alloc.node         and \
       bloc.cpt_date >= alloc.date_beg and \
       bloc.cpt_date <= alloc.date_end:
      res = alloc.alloc_id
      break

  return res


#######################################################################
def cpt_pattern(project, center):

  return "cpt_{}_{}_*.dat".format(center, project)


#######################################################################
def parse_input_cpt(filename, project, center, mode_conso=False):

  if center == "idris":
    return parse_idris_cpt(filename, project)
  elif center == "tgcc":
    return parse_tgcc_cpt(filename, project, mode_conso)
  else:
    print("Unknown center {}".format(center))
    exit()


#######################################################################
def parse_tgcc_cpt(filename, project, mode_conso=False):

  import datetime as dt

  # Determine first and last lines of project bloc(s),
  # and basic project info
  blocs = []
  with open(filename, "r") as filein:
    for num, ligne in enumerate(filein):
      if "Accounting" in ligne:
        _, _, _, name, _, machine, node, _, day = ligne.split()
        bloc = ProjectBloc(
          name.lower(),
          "tgcc",
          machine.lower(),
          node.lower(),
        )
        bloc.cpt_date = dt.datetime.strptime(
          "{} {}".format(day, "23:59"), "%Y-%m-%d %H:%M"
        )
        bloc.line_beg = num + 1
      if "Allocated" in ligne:
        bloc.alloc = float(ligne.split()[-1])
      if "deadline" in ligne:
        bloc.date_end = dt.datetime.strptime(
          "{}".format(ligne.split()[-1]), "%Y-%m-%d"
        )
        bloc.date_end = dt.date(
          bloc.date_end.year,
          bloc.date_end.month,
          bloc.date_end.day,
        )
        bloc.line_end = num + 1
        if bloc.name == project:
          blocs.append(bloc)

  if mode_conso:
    for bloc in blocs:
      bloc.consos = set()
      with open(filename, "r") as filein:
        for _ in xrange(bloc.line_beg - 1):
          next(filein)
        for num, ligne in enumerate(filein):
          if num > bloc.line_end - bloc.line_beg:
            break
          if len(ligne.split()) == 2 and \
             ligne.split()[0] != "Allocated":
            login, conso = ligne.lower().split()
            conso = float(conso)
            # bloc["consos"].add((login.lower(), conso))
            bloc.consos.add(
              ConsoItem(
                login.lower(),
                bloc.cpt_date,
                conso,
              )
            )

  return blocs[0].cpt_date, blocs


#######################################################################
def parse_idris_cpt(filename, project):

  import datetime as dt

  print(filename)

  blocs = []

  with open(filename, "r") as filein:
    # bloc = {}
    bloc = ProjectBloc(
      project,
      "idris",
    )
    bloc.consos = set()
    for num, ligne in enumerate(filein):
      if "mise a jour" in ligne:
        jour, heure = ligne.strip().split()[-2:]
        bloc.cpt_date = dt.datetime.strptime(
          "{} {}".format(jour, heure), "%d/%m/%Y %H:%M"
        )
        bloc.date_end = dt.date(bloc.cpt_date.year, 12, 31)
        # bloc["date_end"] = dt.datetime.strptime(
        #   "{} {}".format(jour, heure), "%d/%m/%Y %H:%M"
        # )
      elif "Comptabilite" in ligne:
        bloc.machine = ligne.strip().split()[-1].strip(".")
        bloc.node = "standard"
      elif "Heures" in ligne:
        bloc.alloc = int(ligne.strip().split()[-1])
      elif "Totaux" in ligne:
        conso = float(ligne.strip().split()[1])
        # bloc.consos.add(("total", conso))
        bloc.consos.add(
          ConsoItem(
            "total",
            bloc.cpt_date,
            conso,
          )
        )
      elif project in ligne:
        login = ligne.strip().split()[0]
        try:
          conso = float(ligne.strip().split()[2])
        except ValueError:
          conso = float("nan")
        # bloc.consos.add((login, conso))
        bloc.consos.add(
          ConsoItem(
            login.lower(),
            bloc.cpt_date,
            conso,
          )
        )

    blocs.append(bloc)

  return blocs[0].cpt_date, blocs


#######################################################################
def date_from_filename(filename):
  """
  """
  return os.path.basename(filename).split("_")[-2]


#######################################################################
def find_input_files(data_dir, pattern, date_range):
  """
  """
  file_list = glob.glob(os.path.join(data_dir, pattern))

  if date_range:
    res = []
    for filename in file_list:
      file_date = date_from_filename(filename)
      if file_date >= min(date_range) and \
         file_date <= max(date_range):
        res.append(filename)
  else:
    res = file_list

  # return set(res)
  return sorted(res)


#######################################################################
if __name__ == "__main__":
  pass