#!/bin/env python
#
# Copyright (c) 2009-2022 CERN. All rights nots expressly granted are
# reserved.
#
# This file is part of iLCDirac
# (see ilcdirac.cern.ch, contact: ilcdirac-support@cern.ch).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# In applying this licence, CERN does not waive the privileges and
# immunities granted to it by virtue of its status as an
# Intergovernmental Organization or submit itself to any jurisdiction.
#
# This program 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
"""Prepare the production summary tables.

Options:
   -P, --prods prodID            Productions: greater than with gt1234, range with 32-56, list with 34,56
   -p, --precise_detail          Precise detail, slow
   -t, --types prodTypeList      Production Types, comma separated, default all
   -S, --Statuses statusList     Statuses, comma separated, default all
   -N, --sample_size             The number of generator files to sample for cross-section and "luminosity, default 100
"""
from __future__ import absolute_import
from six.moves import range
__RCSID__ = "$Id$"

from collections import defaultdict
import os
import random
import textwrap

from DIRAC.Core.Base import Script
from DIRAC.Core.Utilities.DIRACScript import DIRACScript
from DIRAC.Core.Utilities.List import breakListIntoChunks
from DIRAC import S_OK, exit as dexit, S_ERROR


def _getFileInfo(lfn, fc):
  """Retrieve the file info."""
  from DIRAC.Core.Utilities import DEncode
  from DIRAC import gLogger
  lumi = 0.0
  nbevts = 0
  res = fc.getFileUserMetadata(lfn)
  if not res['OK']:
    gLogger.error("Failed to get metadata of %s" % lfn)
    return (0, 0, {})

  lumi += float(res['Value'].get('Luminosity', 0.0))
  nbevts += int(res['Value'].get('NumberOfEvents', 0))

  addinfo = {}
  if 'AdditionalInfo' in res['Value']:
    addinfo = res['Value']['AdditionalInfo']
    if addinfo.count("{"):
      addinfo = eval(addinfo)
    else:
      addinfo = DEncode.decode(addinfo)[0]
  xsec = addinfo.get('xsection', {}).get('sum', {}).get('xsection', 0.0)
  return (float(lumi), int(nbevts), float(xsec))


def _translate(detail):
  """Replace whizard naming convention by human conventions."""
  detail = detail.replace('v', 'n1:n2:n3:N1:N2:N3')
  detail = detail.replace('qli', 'u:d:s:U:D:S')
  detail = detail.replace('ql', 'u:d:s:c:b:U:D:S:C:B')
  detail = detail.replace('q', 'u:d:s:c:b:t')
  detail = detail.replace('Q', 'U:D:S:C:B:T')
  detail = detail.replace('e1', 'e-')
  detail = detail.replace('E1', 'e+')
  detail = detail.replace('e2', 'mu-')
  detail = detail.replace('E2', 'mu+')
  detail = detail.replace('e3', 'tau-')
  detail = detail.replace('E3', 'tau+')
  detail = detail.replace('n1', 'nue')
  detail = detail.replace('N1', 'nueb')
  detail = detail.replace('n2', 'numu')
  detail = detail.replace('N2', 'numub')
  detail = detail.replace('n3', 'nutau')
  detail = detail.replace('N3', 'nutaub')
  detail = detail.replace('U', 'ubar')
  detail = detail.replace('C', 'cbar')
  detail = detail.replace('T', 'tbar')
  detail = detail.replace('tbareV', 'TeV')
  detail = detail.replace('D', 'dbar')
  detail = detail.replace('S', 'sbar')
  detail = detail.replace('B', 'bbar')
  detail = detail.replace('Z0', 'Z')
  detail = detail.replace('Z', 'Z0')
  detail = detail.replace('gghad', 'gamma gamma -> hadrons')
  detail = detail.replace(',', '')
  detail = detail.replace('n N', 'nu nub')
  detail = detail.replace('se--', 'seL-')
  detail = detail.replace('se-+', 'seL+')
  detail = detail.replace(' -> ', '->')
  detail = detail.replace('->', ' -> ')
  detail = detail.replace(' H ->', ', H ->')
  detail = detail.replace(' Z0 ->', ', Z0 ->')
  detail = detail.replace(' W ->', ', W ->')

  return detail

# def getAncestor(lfn):
#  from DIRAC.Resources.Catalog.FileCatalogClient import FileCatalogClient
#  fc = FileCatalogClient()
#  res = fc.getFileAncestors([lfn],1)
#  if not res['OK']:
#    return S_ERROR("Failed getting ancestor")
#  for ancestor in res['Value']['Successful'][lfn].keys():
#    if not ancestor.count("stdhep"):
#      res = getAncestor(ancestor)
#      if not res['OK']:
#        return S_ERROR("Failed geting ancestor")
#    else:
#      return S_OK(ancestor)


class _Params(object):
  """CLI Parameters class."""

  def __init__(self):
    """Initialize."""
    self.prod = []
    self.minprod = 0
    self.full_det = False
    self.sampleSize = 100
    self.ptypes = ['MCGeneration', 'MCSimulation', 'MCReconstruction', "MCReconstruction_Overlay"]
    self.statuses = ['Active', 'Stopped', 'Completed', 'Archived']

  def setProdID(self, opt):
    """Set the prodID to use.

    Can be a range, a list, a unique value and a 'greater than' value.
    """
    if opt.count(","):
      parts = opt.split(",")
    else:
      parts = [opt]
    prods = []
    try:
      for part in parts:
        if part.count("gt"):
          self.minprod = int(part.replace("gt", ""))
          continue
        if part.count("-"):
          prods.extend(list(range(int(part.split("-")[0]), int(part.split("-")[1]) + 1)))
        else:
          prods.append(int(part))
    except ValueError:
      return S_ERROR('ProductionID: bad syntax')
    self.prod = prods

    return S_OK()

  def setFullDetail(self, _opt):
    """Get every individual file's properties, makes this very very slow."""
    self.full_det = True
    return S_OK()

  def setSampleSize(self, opt):
    """Set the number of generator files to sample for cross-section and luminosity."""
    try:
      self.sampleSize = int(opt)
    except ValueError:
      return S_ERROR('SampleSize: bad value')
    return S_OK()

  def setProdTypes(self, opt):
    """Set the prod types to consider."""
    self.ptypes = opt.split(",")
    return S_OK()

  def setStatuses(self, opt):
    """Set the prod statuses."""
    self.statuses = opt.split(",")
    return S_OK()

  def registerSwitches(self):
    """Register all CLI switches."""
    Script.registerSwitch(
        "P:",
        "prods=",
        "Productions: greater than with gt1234, range with 32-56, list with 34,56",
        self.setProdID)
    Script.registerSwitch("p", "precise_detail", "Precise detail, slow", self.setFullDetail)
    Script.registerSwitch("N:", "sample_size=", "The number of generator files to sample for cross-section and "
                          "luminosity, default 100", self.setSampleSize)
    Script.registerSwitch("t:", "types=", "Production Types, comma separated, default all", self.setProdTypes)
    Script.registerSwitch("S:", "Statuses=", "Statuses, comma separated, default all", self.setStatuses)
    Script.setUsageMessage('\n'.join([__doc__.split('\n')[1],
                                         '\nUsage:',
                                         '  %s [option|cfgfile] ...\n' % Script.scriptName]))


@DIRACScript()
def main():
  clip = _Params()
  clip.registerSwitches()
  Script.parseCommandLine()
  from ILCDIRAC.Core.Utilities.HTML import Table
  from ILCDIRAC.Core.Utilities.ProcessList import ProcessList
  from DIRAC.TransformationSystem.Client.TransformationClient import TransformationClient
  from DIRAC.Resources.Catalog.FileCatalogClient import FileCatalogClient
  from DIRAC import gConfig, gLogger
  prod = clip.prod
  full_detail = clip.full_det
  fc = FileCatalogClient()

  processlist = gConfig.getValue('/LocalSite/ProcessListPath', '')
  prl = ProcessList(processlist)
  processesdict = prl.getProcessesDict()

  trc = TransformationClient()
  prodids = []
  if not prod:
    conddict = {}
    conddict['Status'] = clip.statuses
    if clip.ptypes:
      conddict['Type'] = clip.ptypes
    res = trc.getTransformations(conddict)
    if res['OK']:
      for transfs in res['Value']:
        prodids.append(transfs['TransformationID'])
  else:
    prodids.extend(prod)

  metadata = []

  gLogger.info("Will run on prods %s" % str(prodids))

  for prodID in prodids:
    if prodID < clip.minprod:
      continue
    meta = {}
    meta['ProdID'] = prodID
    res = trc.getTransformation(str(prodID))
    if not res['OK']:
      gLogger.error("Error getting transformation %s" % prodID)
      continue
    prodtype = res['Value']['Type']
    proddetail = res['Value']['Description']
    meta['Datatype'] = 'DST' if prodtype in ('MCReconstruction', 'MCReconstruction_Overlay') \
                       else 'gen' if prodtype == 'MCGeneration' \
                       else 'SIM' if prodtype == 'MCSimulation' else None
    if prodtype in ['Split', 'Merge']:
      gLogger.warn("Invalid query for %s productions" % prodtype)
      continue
    if not meta['Datatype']:
      gLogger.error("Unknown production type %s" % prodtype)
      continue
    gLogger.verbose('Looking for files: %s' % meta)
    res = fc.findFilesByMetadata(meta)
    if not res['OK']:
      gLogger.error(res['Message'])
      continue
    lfns = res['Value']
    nb_files = len(lfns)
    if not lfns:
      gLogger.warn("No files found for prod %s" % prodID)
      continue
    path = os.path.dirname(lfns[0])
    gLogger.verbose('Looking for user metadata')
    res = fc.getDirectoryUserMetadata(path)
    if not res['OK']:
      gLogger.warn('No meta data found for %s' % path)
      continue
    dirmeta = {'proddetail': proddetail,
               'prodtype': prodtype,
               'nb_files': nb_files,
               }
    dirmeta.update(res['Value'])
    lumi = 0.
    nbevts = 0
    files = 0
    xsec = 0.0
    if not full_detail:
      info = _getFileInfo(lfns[0], fc)
      lumi = info[0] * len(lfns)
      nbevts = info[1] * len(lfns)
      if info[2]:
        xsec += info[2]
        files += 1
    else:
      for lfn in lfns:
        info = _getFileInfo(lfn, fc)
        lumi += info[0]
        nbevts += info[1]
        if info[2]:
          xsec += info[2]
          files += 1
        gLogger.verbose('Found lumi: %s, xsec: %s' % (lumi, xsec))
        if not lumi:  # we are guessing none of the files have that information at this point
          gLogger.notice('Files do not have xsection info, looking at ancestors now')
          nbevts = info[1] * len(lfns)
          break

    if not lumi:
      xsec = 0
      files = 0
      depthDict = defaultdict(set)
      gLogger.verbose('Looking for Ancestors of %s files' % len(lfns))
      for index, lfnChunk in enumerate(breakListIntoChunks(lfns[:1000], 1000)):
        gLogger.verbose('Looking at chunk %s/%s' % (index, len(lfns) // 1000))
        res = fc.getFileAncestors(lfnChunk, [1, 2, 3, 4])
        if not res['OK']:
          gLogger.error('failed to find ancenstors')
          return res
        for lfn, ancestorsDict in res['Value']['Successful'].items():
          for ancestor, dep in ancestorsDict.items():
            depthDict[dep].add(ancestor)
      gLogger.verbose('Found all the ancestors')
      oldestAncestorID = sorted(depthDict)[-1]
      oldestAncestors = depthDict[oldestAncestorID]
      numberOfAncestors = len(oldestAncestors)
      sampleSize = clip.sampleSize if not full_detail else numberOfAncestors
      gLogger.verbose('Sampling %s from %s' % (sampleSize, numberOfAncestors))
      for index, ancestor in enumerate(random.sample(oldestAncestors, sampleSize)):
        gLogger.verbose('Getting info for ancestor %s / %s' % (index, sampleSize))
        info = _getFileInfo(ancestor, fc)
        # lumi is not used later on
        lumi += info[0] * float(numberOfAncestors) / float(sampleSize)
        if info[2]:
          xsec += info[2]
          files += 1
        gLogger.verbose('Found %s, %s, %s' % (info[0], info[2], files))

    dirmeta['CrossSection'] = xsec / float(files) if xsec and files else 0.0
    dirmeta['NumberOfEvents'] = nbevts if nbevts else 0

    dirmeta['detail'] = ''
    if dirmeta['EvtType'] in processesdict:
      if 'Detail' in processesdict[dirmeta['EvtType']]:
        detail = processesdict[dirmeta['EvtType']]['Detail']
    else:
      detail = dirmeta['EvtType']

    dirmeta['MomProdID'] = 0
    if prodtype != 'MCGeneration':
      res = trc.getTransformationMetaDataQuery(str(prodID), 'Input')
      if res['OK']:
        if 'ProdID' in res['Value']:
          dirmeta['MomProdID'] = res['Value']['ProdID']

    dirmeta['detail'] = _translate(detail)
    metadata.append(dirmeta)

  detectors = {'ILD': {'SIM': [], 'REC': [], 'title': 'ILD CDR'},
               'SID': {'SIM': [], 'REC': [], 'title': 'SID CDR'},
               'sid': {'SIM': [], 'REC': [], 'title': 'sid dbd prod'},
               'FCCee_o1_v04': {'SIM': [], 'REC': [], 'title': 'FCC studies'},
               'gen': [],
               }

  corres = {'MCGeneration': 'gen',
            'MCSimulation': 'SIM',
            'MCReconstruction': 'REC',
            'MCReconstruction_Overlay': 'REC'}

  for channel in metadata:
    if 'DetectorType' not in channel:
      detectors['gen'].append((channel['detail'],
                               channel['Energy'],
                               channel['ProdID'],
                               channel['nb_files'],
                               channel['NumberOfEvents'] / channel['nb_files'],
                               channel['NumberOfEvents'],
                               channel['CrossSection'], str(channel['proddetail'])))
    else:
      if channel['DetectorType'] not in detectors:
        gLogger.error("This is unknown detector", channel['DetectorType'])
        continue
      detectors[channel['DetectorType']][corres[channel['prodtype']]].append((channel['detail'],
                                                                              channel['Energy'],
                                                                              channel['DetectorType'],
                                                                              channel['ProdID'],
                                                                              channel['nb_files'],
                                                                              channel['NumberOfEvents'] / channel['nb_files'],
                                                                              channel['NumberOfEvents'],
                                                                              channel['CrossSection'],
                                                                              channel['MomProdID'],
                                                                              str(channel['proddetail'])))
  header_row_tuple_gen = ('Channel', 'Energy', 'ProdID', 'Tasks', 'Average Evts/task', 'Statistics',
                          'Cross Section (fb)', 'Comment')

  header_row_tuple_det = ('Channel', 'Energy', 'Detector', 'ProdID', 'Number of Files', 'Events/File',
                          'Statistics', 'Cross Section (fb)', 'Origin ProdID', 'Comment')

  with open('tables.html', 'w') as of:
    of.write(textwrap.dedent("""
                             <!DOCTYPE html>
                             <html>
                               <head>
                                 <title> Production summary </title>
                               </head>
                               <body>
                             """).strip())
    if detectors['gen']:
      of.write("<h1>gen prods</h1>\n")
      table = Table(header_row=header_row_tuple_gen)
      table.rows.extend(detectors['gen'])
      of.write(str(table))
      gLogger.info("Gen prods")
      gLogger.info(str(table))

    detectors.pop('gen')

    for detName, infos in detectors.items():
      if any(infos.get(pt) for pt in ('SIM', 'REC')):
        of.write('<h1>%s prods</h1>\n' % infos['title'])
        for ptype in ('SIM', 'REC'):
          if infos[ptype]:
            of.write('<h2>%s</h2>\n' % ptype)
            table = Table(header_row=header_row_tuple_det)
            table.rows.extend(infos[ptype])
            of.write(str(table))
            gLogger.info('%s prods %s' % (infos['title'], ptype))
            gLogger.info(str(table))

    of.write(textwrap.dedent("""
                               </body>
                             </html>"""))
  gLogger.notice("Check ./tables.html in any browser for the results")
  dexit(0)


if __name__ == "__main__":
  main()
