from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import time
import copy

from DIRAC import S_OK, S_ERROR, gLogger
from DIRAC.AccountingSystem.private.DBUtils import DBUtils
from DIRAC.Core.Utilities.Plotting import gDataCache
from DIRAC.Core.Utilities.Plotting.Plots import generateNoDataPlot,\
    generateTimedStackedBarPlot,\
    generateQualityPlot,\
    generateCumulativePlot,\
    generatePiePlot,\
    generateStackedLinePlot,\
    generateHistogram


class BaseReporter(DBUtils):

  _PARAM_CHECK_FOR_NONE = 'checkNone'
  _PARAM_CALCULATE_PROPORTIONAL_GAUGES = 'calculateProportionalGauges'
  _PARAM_CONVERT_TO_GRANULARITY = 'convertToGranularity'
  _VALID_PARAM_CONVERT_TO_GRANULARITY = ('sum', 'average')
  _PARAM_CONSOLIDATION_FUNCTION = "consolidationFunction"

  _EA_THUMBNAIL = 'thumbnail'
  _EA_WIDTH = 'width'
  _EA_HEIGHT = 'height'
  _EA_THB_WIDTH = 'thbWidth'
  _EA_THB_HEIGHT = 'thbHeight'
  _EA_PADDING = 'figurePadding'
  _EA_TITLE = 'plotTitle'

  _RATE_UNITS = {
      'time': (
          ('seconds / s',
           1,
           24),
          ('hours / s',
           3600,
           24),
          ('days / s',
           86400,
           15),
          ('weeks / s',
           86400 * 7,
           10),
          ('months / s',
           86400 * 30,
           12),
          ('years / s',
           86400 * 365,
           1)),
      'cpupower': (
          ('HS06',
           1,
           750),
          ('kHS06',
           1000,
           750),
          ('MHS06',
           10 ** 6,
           1)),
      'bytes': (
          ('MB / s',
           10 ** 6,
           1024),
          ('GB / s',
           10 ** 9,
           1024),
          ('TB / s',
           10 ** 12,
           1024),
          ('PB / s',
           10 ** 15,
           1)),
      'jobs': (
          ('jobs / hour',
           1 / 3600.0,
           1000),
          ('kjobs / hour',
           (10 ** 3) / 3600.0,
           1000),
          ('Mjobs / hour',
           (10 ** 6) / 3600.0,
           1)),
      'files': (
          ('files / hour',
           1 / 3600.0,
           1000),
          ('kfiles / hour',
           (10 ** 3) / 3600.0,
           1000),
          ('Mfiles / hour',
           (10 ** 6) / 3600.0,
           1))}

  _UNITS = {'time': (('seconds', 1, 24),
                     ('hours', 3600, 24),
                     ('days', 86400, 15),
                     ('weeks', 86400 * 7, 10),
                     ('months', 86400 * 30, 12),
                     ('years', 86400 * 365, 1)),
            'cpupower': (('HS06 hours', 3600, 24),
                         ('HS06 days', 86400, 750),
                         ('kHS06 days', 86400 * 1000, 750),
                         ('MHS06 days', 86400 * 10**6, 1)),
            'bytes': (('MB', 10 ** 6, 1024),
                      ('GB', 10 ** 9, 1024),
                      ('TB', 10 ** 12, 1024),
                      ('PB', 10 ** 15, 1)),
            'jobs': (('jobs', 1, 1000),
                     ('kjobs', 10 ** 3, 1000),
                     ('Mjobs', 10 ** 6, 1)),
            'files': (('files', 1, 1000),
                      ('kfiles', 10 ** 3, 1000),
                      ('Mfiles', 10 ** 6, 1))
            }

  # To be defined in the derived classes
  _typeKeyFields = []
  _typeName = ''

  def __init__(self, db, setup, extraArgs=None):
    DBUtils.__init__(self, db, setup)
    if isinstance(extraArgs, dict):
      self._extraArgs = extraArgs
    else:
      self._extraArgs = {}
    reportsRevMap = {}
    for attr in dir(self):
      if attr.find("_report") == 0:
        if attr.find('Name', len(attr) - 4) == len(attr) - 4:
          reportId = attr[7:-4]
          reportName = getattr(self, attr)
          reportsRevMap[reportId] = reportName
        else:
          reportId = attr[7:]
          if reportId not in reportsRevMap:
            reportsRevMap[reportId] = reportId
    self.__reportNameMapping = {}
    for rId in reportsRevMap:
      self.__reportNameMapping[reportsRevMap[rId]] = rId

  def _translateGrouping(self, grouping):
    return ("%s", [grouping])

  def _averageConsolidation(self, total, count):
    if count == 0:
      return 0
    else:
      return float(total) / float(count)

  def _efficiencyConsolidation(self, total, count):
    if count == 0:
      return 0
    else:
      return (float(total) / float(count)) * 100.0

  def generate(self, reportRequest):
    reportRequest['groupingFields'] = self._translateGrouping(reportRequest['grouping'])
    reportHash = reportRequest['hash']
    reportName = reportRequest['reportName']
    if reportName in self.__reportNameMapping:
      reportRequest['reportName'] = self.__reportNameMapping[reportName]
    gLogger.info("Retrieving data for %s:%s" % (reportRequest['typeName'], reportRequest['reportName']))
    sT = time.time()
    retVal = self.__retrieveReportData(reportRequest, reportHash)
    reportGenerationTime = time.time() - sT
    if not retVal['OK']:
      return retVal
    if not reportRequest['generatePlot']:
      return retVal
    reportData = retVal['Value']
    gLogger.info("Plotting data for %s:%s" % (reportRequest['typeName'], reportRequest['reportName']))
    sT = time.time()
    retVal = self.__generatePlotForReport(reportRequest, reportHash, reportData)
    plotGenerationTime = time.time() - sT
    gLogger.info(
        "Time for %s:%s - Report %.2f Plot %.2f (%.2f%% r/p)" %
        (reportRequest['typeName'],
         reportRequest['reportName'],
         reportGenerationTime,
         plotGenerationTime,
         ((reportGenerationTime * 100 / plotGenerationTime) if plotGenerationTime else 0.)))
    if not retVal['OK']:
      return retVal
    plotDict = retVal['Value']
    if 'retrieveReportData' in reportRequest['extraArgs'] and reportRequest['extraArgs']['retrieveReportData']:
      plotDict['reportData'] = reportData
    return S_OK(plotDict)

  def plotsList(self):
    return sorted([k for k in self.__reportNameMapping])

  def __retrieveReportData(self, reportRequest, reportHash):
    funcName = "_report%s" % reportRequest['reportName']
    try:
      funcObj = getattr(self, funcName)
    except Exception:
      return S_ERROR("Report %s is not defined" % reportRequest['reportName'])
    return gDataCache.getReportData(reportRequest, reportHash, funcObj)

  def __generatePlotForReport(self, reportRequest, reportHash, reportData):
    funcName = "_plot%s" % reportRequest['reportName']
    try:
      funcObj = getattr(self, funcName)
    except Exception:
      return S_ERROR("Plot function for report %s is not defined" % reportRequest['reportName'])
    return gDataCache.getReportPlot(reportRequest, reportHash, reportData, funcObj)

###
# Helper functions for reporters
###

  def _getTimedData(self, startTime, endTime, selectFields, preCondDict, groupingFields, metadataDict):
    condDict = {}
    # Check params
    if self._PARAM_CHECK_FOR_NONE not in metadataDict:
      metadataDict[self._PARAM_CHECK_FOR_NONE] = False
    if self._PARAM_CONVERT_TO_GRANULARITY not in metadataDict:
      metadataDict[self._PARAM_CONVERT_TO_GRANULARITY] = "sum"
    elif metadataDict[self._PARAM_CONVERT_TO_GRANULARITY] not in self._VALID_PARAM_CONVERT_TO_GRANULARITY:
      return S_ERROR("%s field metadata is invalid" % self._PARAM_CONVERT_TO_GRANULARITY)
    if self._PARAM_CALCULATE_PROPORTIONAL_GAUGES not in metadataDict:
      metadataDict[self._PARAM_CALCULATE_PROPORTIONAL_GAUGES] = False
    # Make safe selections
    for keyword in self._typeKeyFields:
      if keyword in preCondDict:
        condDict[keyword] = preCondDict[keyword]
    # Query!
    timeGrouping = ("%%s, %s" % groupingFields[0], ['startTime'] + groupingFields[1])
    retVal = self._retrieveBucketedData(self._typeName,
                                        startTime,
                                        endTime,
                                        selectFields,
                                        condDict,
                                        timeGrouping,
                                        ('%s', ['startTime'])
                                        )
    if not retVal['OK']:
      return retVal
    dataDict = self._groupByField(0, retVal['Value'])
    coarsestGranularity = self._getBucketLengthForTime(self._typeName, startTime)
    # Transform!
    for keyField in dataDict:
      if metadataDict[self._PARAM_CHECK_FOR_NONE]:
        dataDict[keyField] = self._convertNoneToZero(dataDict[keyField])
      if metadataDict[self._PARAM_CONVERT_TO_GRANULARITY] == "average":
        dataDict[keyField] = self._averageToGranularity(coarsestGranularity, dataDict[keyField])
      if metadataDict[self._PARAM_CONVERT_TO_GRANULARITY] == "sum":
        dataDict[keyField] = self._sumToGranularity(coarsestGranularity, dataDict[keyField])
      if self._PARAM_CONSOLIDATION_FUNCTION in metadataDict:
        dataDict[keyField] = self._executeConsolidation(
            metadataDict[self._PARAM_CONSOLIDATION_FUNCTION], dataDict[keyField])
    if metadataDict[self._PARAM_CALCULATE_PROPORTIONAL_GAUGES]:
      dataDict = self._calculateProportionalGauges(dataDict)
    return S_OK((dataDict, coarsestGranularity))

  def _executeConsolidation(self, functor, dataDict):
    for timeKey in dataDict:
      dataDict[timeKey] = [functor(*dataDict[timeKey])]
    return dataDict

  def _getSummaryData(
          self,
          startTime,
          endTime,
          selectFields,
          preCondDict,
          groupingFields,
          metadataDict=None,
          reduceFunc=False):
    condDict = {}
    # Make safe selections
    for keyword in self._typeKeyFields:
      if keyword in preCondDict:
        condDict[keyword] = preCondDict[keyword]
    # Query!
    retVal = self._retrieveBucketedData(self._typeName,
                                        startTime,
                                        endTime,
                                        selectFields,
                                        condDict,
                                        groupingFields,
                                        []
                                        )
    if not retVal['OK']:
      return retVal
    dataDict = self._groupByField(0, retVal['Value'])
    for key in dataDict:
      if not reduceFunc:
        dataDict[key] = dataDict[key][0][0]
      else:
        dataDict[key] = reduceFunc(*dataDict[key][0])
    # HACK to allow serialization of the type because MySQL decides to return data in a non python standard format
    for key in dataDict:
      dataDict[key] = float(dataDict[key])
    return S_OK(dataDict)

  def _getBucketData(self, startTime, endTime, selectFields, preCondDict):
    """
    It retrieves data for histogram.

    :param int startTime: epoch time
    :param int endTime: epoch time
    :param list selectFields: the value being plotted
    :param dict preCondDict: plot attributes (conditions)

    """

    condDict = {}
    # Make safe selections
    for keyword in self._typeKeyFields:
      if keyword in preCondDict:
        condDict[keyword] = preCondDict[keyword]
    retVal = self._retrieveBucketedData(self._typeName,
                                        startTime,
                                        endTime,
                                        selectFields,
                                        condDict)
    if not retVal['OK']:
      return retVal
    dataDict = self._groupByField(0, retVal['Value'])
    data = {}
    # convert values to the correct format
    for key, values in dataDict.items():
      data[key] = [float(i[0]) for i in values]
    return S_OK(data)

  def _getSelectStringForGrouping(self, groupingFields):
    if len(groupingFields) == 3:
      return groupingFields[2]
    if len(groupingFields[1]) == 1:
      # If there's only one field, then we send the sql representation in pos 0
      return groupingFields[0]
    else:
      return "CONCAT( %s )" % ", ".join(["%s, '-'" % sqlRep for sqlRep in groupingFields[0]])

  def _findSuitableRateUnit(self, dataDict, maxValue, unit):
    return self._findUnitMagic(dataDict, maxValue, unit, self._RATE_UNITS)

  def _findSuitableUnit(self, dataDict, maxValue, unit):
    return self._findUnitMagic(dataDict, maxValue, unit, self._UNITS)

  def _findUnitMagic(self, reportDataDict, maxValue, unit, selectedUnits):
    if unit not in selectedUnits:
      raise AttributeError("%s is not a known rate unit" % unit)
    baseUnitData = selectedUnits[unit][0]
    if 'staticUnits' in self._extraArgs and self._extraArgs['staticUnits']:
      unitData = selectedUnits[unit][0]
    else:
      unitList = selectedUnits[unit]
      unitIndex = -1
      for unitName, unitDivFactor, unitThreshold in unitList:
        unitIndex += 1
        if maxValue / unitDivFactor < unitThreshold:
          break
      unitData = selectedUnits[unit][unitIndex]
    # Apply divFactor to all units
    graphDataDict, maxValue = self._divideByFactor(copy.deepcopy(reportDataDict), unitData[1])
    if unitData == baseUnitData:
      reportDataDict = graphDataDict
    else:
      reportDataDict, dummyMaxValue = self._divideByFactor(reportDataDict, baseUnitData[1])
    return reportDataDict, graphDataDict, maxValue, unitData[0]
##
# Plotting
##

  def __checkPlotMetadata(self, metadata):
    if self._EA_WIDTH in self._extraArgs and self._extraArgs[self._EA_WIDTH]:
      try:
        metadata[self._EA_WIDTH] = min(1600, max(200, int(self._extraArgs[self._EA_WIDTH])))
      except Exception:
        pass
    if self._EA_HEIGHT in self._extraArgs and self._extraArgs[self._EA_HEIGHT]:
      try:
        metadata[self._EA_HEIGHT] = min(1600, max(200, int(self._extraArgs[self._EA_HEIGHT])))
      except Exception:
        pass
    if self._EA_TITLE in self._extraArgs and self._extraArgs[self._EA_TITLE]:
      metadata['title'] = self._extraArgs[self._EA_TITLE]

  def __checkThumbnailMetadata(self, metadata):
    if self._EA_THUMBNAIL in self._extraArgs and self._extraArgs[self._EA_THUMBNAIL]:
      thbMD = dict(metadata)
      thbMD['legend'] = False
      if self._EA_THB_HEIGHT in self._extraArgs:
        thbMD[self._EA_HEIGHT] = self._extraArgs[self._EA_THB_HEIGHT]
      else:
        thbMD[self._EA_HEIGHT] = 125
      if self._EA_THB_WIDTH in self._extraArgs:
        thbMD[self._EA_WIDTH] = self._extraArgs[self._EA_THB_WIDTH]
      else:
        thbMD[self._EA_WIDTH] = 200
      thbMD[self._EA_PADDING] = 20
      for key in ('title', 'ylabel', 'xlabel'):
        if key in thbMD:
          del(thbMD[key])
      return thbMD
    return False

  def __plotData(self, filename, dataDict, metadata, funcToPlot):
    self.__checkPlotMetadata(metadata)
    if not dataDict:
      funcToPlot = generateNoDataPlot
    plotFileName = "%s.png" % filename
    finalResult = funcToPlot(plotFileName, dataDict, metadata)
    if not finalResult['OK']:
      return finalResult
    thbMD = self.__checkThumbnailMetadata(metadata)
    if not thbMD:
      return S_OK({'plot': True, 'thumbnail': False})
    thbFilename = "%s.thb.png" % filename
    retVal = funcToPlot(thbFilename, dataDict, thbMD)
    if not retVal['OK']:
      return retVal
    return S_OK({'plot': True, 'thumbnail': True})

  def _generateTimedStackedBarPlot(self, filename, dataDict, metadata):
    return self.__plotData(filename, dataDict, metadata, generateTimedStackedBarPlot)

  def _generateQualityPlot(self, filename, dataDict, metadata):
    return self.__plotData(filename, dataDict, metadata, generateQualityPlot)

  def _generateCumulativePlot(self, filename, dataDict, metadata):
    return self.__plotData(filename, dataDict, metadata, generateCumulativePlot)

  def _generatePiePlot(self, filename, dataDict, metadata):
    return self.__plotData(filename, dataDict, metadata, generatePiePlot)

  def _generateStackedLinePlot(self, filename, dataDict, metadata):
    return self.__plotData(filename, dataDict, metadata, generateStackedLinePlot)

  def _generateHistogram(self, filename, dataDict, metadata):
    return self.__plotData(filename, dataDict, metadata, generateHistogram)
