###############################################################################
# (c) Copyright 2000-2022 CERN for the benefit of the LHCb Collaboration      #
#                                                                             #
# This software is distributed under the terms of the GNU General Public      #
# Licence version 3 (GPL Version 3), copied verbatim in the file "COPYING".   #
#                                                                             #
# 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.                                       #
###############################################################################

import GaudiPython as GP
from GaudiConf.reading import load_manifest as load_tck, do_unpacking
from Configurables import (ApplicationMgr, LHCbApp, IODataManager,
                           HltDecReportsDecoder)
from GaudiConf import IOHelper
from PyConf.application import configured_ann_svc
import os
import operator
from collections import Counter
import json
import re
import argparse
import csv
'''

    Run snippet with 'python line-rates.py and [1] <MDF file name> [2] <TCK config file name> [3] <JSON file name specifying configuration>'
    + '--c' flag with 'five or 'sixteen' for 5 or 16 stream configuration
    + '--s' flag with 'line' or 'stream' to get calculate per stream or per line
    Note: '--c sixteen --s line' is not a valid combination of arguments

    When running 5-stream config, returns:

        Per line (in form of single HTML table):
            1. Inclusive retention
            2. Inclusive rate
            3. Exclusive retention
            4. Exclusive rate
            5. Average DstData bank size
            6. DstData bandwidth
            7. Average event size (all banks in particular stream)
            8. Bandwidth

        Per stream in Turbo/Full/Turcal/Monitoring/IFT
            1. Inclusive retention
            2. Inclusive rate
            3. Average DstData bank size
            4. DstData bandwidth
            5. Average event size (all banks in particular stream)
            6. Bandwidth

    When running 16-stream config, returns same figures as above (only per stream)

'''

LHCb = GP.gbl.LHCb
RAW_BANK_TYPES = [(i, LHCb.RawBank.typeName(i))
                  for i in range(LHCb.RawBank.LastType)]
banks_all = [(16, 'ODIN'), (17, 'HltDecReports'), (53, 'HltRoutingBits'),
             (56, 'HltLumiSummary'), (60, 'DstData')]
banks_turcal_ift = [(9, 'Rich'), (13, 'Muon'), (16, 'ODIN'),
                    (17, 'HltDecReports'), (21, 'EcalPacked'),
                    (35, 'HcalPackedError'), (53, 'HltRoutingBits'),
                    (56, 'HltLumiSummary'), (60, 'DstData'), (63, 'VP'),
                    (64, 'FTCluster'), (66, 'UT'), (73, 'VPRetinaCluster'),
                    (77, 'Calo'), (84, 'Plume')]


def rawbank_sizes(rawevent, lst):
    """Return (name, size) for each raw bank type."""
    if rawevent:

        def size(i):
            return sum(bank.totalSize() for bank in rawevent.banks(i))
    else:

        def size(i):
            return 0

    return [(name, size(i)) for i, name in lst]


def processing_events_per_line(evt_max, all_banks, lines):
    '''
        Return, per line:
        i) How many events triggered on
        ii) Average DstData size of all events
        iii) Average size of all events

    '''

    # Stores how many events each line fired on
    event_stats = {
        line: []
        for line in [line + 'Decision' for line in list(lines)]
    }

    # Stores whole event size size
    raw = {line: 0 for line in [line + 'Decision' for line in list(lines)]}

    # Stores DstData bank size
    dst = {line: 0 for line in [line + 'Decision' for line in list(lines)]}

    exclusive = {}

    # Loop over all events
    analysed = 0
    while analysed < evt_max:
        analysed += 1

        exclusive.update({analysed: 0})

        # Run an event
        #print('Event: ', analysed, '\n')
        appMgr.run(1)
        report = evt['/Event/Hlt2/DecReports']
        rawevent = evt['/Event/DAQ/RawEvent']

        if all_banks:
            bank_sizes = rawbank_sizes(rawevent, banks_turcal_ift)
        else:
            bank_sizes = rawbank_sizes(rawevent, banks_all)

        dst_sizes = rawbank_sizes(rawevent, [(60, 'DstData')])

        evtsize = sum(bank[1] for bank in bank_sizes)
        dstsize = sum(bank[1] for bank in dst_sizes)

        # Will quit running if there are no more events in the input file
        if report:
            for line in event_stats.keys():
                if report.decReport(line):
                    if report.decReport(line).decision() == 1:
                        event_stats[line].append(analysed)
                        exclusive[analysed] += 1
                        raw[line] += evtsize
                        dst[line] += dstsize
        else:
            break

    return event_stats, exclusive, raw, dst


def processing_events_per_stream(evt_max, all_banks):
    '''
        Returns number of events, cumulative event size
        and cumulative dst size for specified stream

    '''

    events = 0
    raw_size = 0
    dst_size = 0

    # Loop over all events
    analysed = 0
    while analysed < evt_max:
        analysed += 1

        # Run an event
        appMgr.run(1)
        report = evt['/Event/Hlt2/DecReports']
        rawevent = evt['/Event/DAQ/RawEvent']

        if all_banks:
            bank_sizes = rawbank_sizes(rawevent, banks_turcal_ift)
        else:
            bank_sizes = rawbank_sizes(rawevent, banks_all)

        dst_sizes = rawbank_sizes(rawevent, [(60, 'DstData')])

        evtsize = sum(bank[1] for bank in bank_sizes)
        dstsize = sum(bank[1] for bank in dst_sizes)

        if report:
            events += 1
            raw_size += evtsize
            dst_size += dstsize
        else:
            break

    return events, raw_size, dst_size


def rates_per_line(event_stats, exclusive, raw, dst, configname, streamname):

    data = []

    # Compute exclusive rate
    sort = dict(
        sorted(
            {k: v
             for (k, v) in exclusive.items() if v > 0}.items(),
            key=operator.itemgetter(1),
            reverse=True))

    unique_events = [key for key, value in sort.items() if value == 1]

    for line, val in event_stats.items():
        events_all = val + unique_events
        num_events = len(event_stats[line])
        row_values = (
            line,
            num_events / LHCbApp().EvtMax * 100
            if num_events else 0,  # Inclusive Retention (expressed as %)
            num_events / LHCbApp().EvtMax * 1e3
            if num_events else 0,  # Inclusive Rate (in kHz)
            len([
                key for key, value in Counter(events_all).items() if value > 1
            ]) / LHCbApp().EvtMax * 100
            if num_events else 0,  # Exclusive retention (expressed as %)
            len([
                key for key, value in Counter(events_all).items() if value > 1
            ]) / LHCbApp().EvtMax * 1e3
            if num_events else 0,  # Exclusive rate (in kHz)
            raw[line] / num_events * 1e-3
            if num_events else 0,  # Average event size (in kB)
            (num_events / LHCbApp().EvtMax * raw[line] / num_events) / 1e3
            if num_events else 0,  # Event bandwidth (in GB/s)
            dst[line] / len(event_stats[line]) * 1e-3
            if num_events else 0,  # Average DstData size (in kB)
            (num_events / LHCbApp().EvtMax * dst[line] / num_events) / 1e3
            if num_events else 0)  # DstData Bandwidth (in GB/s)

        data.append(row_values)

    with open(
            f'tmp/Output/Inter/rates-all-lines-{configname}-{streamname}.csv',
            'w') as f:
        csv_out = csv.writer(f)
        for tup in data:
            csv_out.writerow(tup)

    return


def rates_per_stream(events, raw_size, dst_size, configname, streamname):

    data = []

    row_values = (
        streamname,
        events / LHCbApp().EvtMax * 100
        if events else 0,  # Inclusive Retention (expressed as %)
        events / LHCbApp().EvtMax * 1e3
        if events else 0,  # Inclusive Rate (in kHz)
        raw_size / events * 1e-3
        if events else 0,  # Average event size (in kB)
        (events / LHCbApp().EvtMax * raw_size / events) / 1e3
        if events else 0,  # Event bandwidth (in GB/s)
        dst_size / events * 1e-3
        if events else 0,  # Average DstData size (in kB)
        (events / LHCbApp().EvtMax * dst_size / events) / 1e3 if events else 0
    )  # DstData Bandwidth (in GB/s)

    data.append(row_values)

    with open(
            f'tmp/Output/Inter/rates-per-stream-{configname}-{streamname}.csv',
            'w') as f:
        csv_out = csv.writer(f)
        for tup in data:
            csv_out.writerow(tup)

    return


parser = argparse.ArgumentParser(description='Inspect Moore output')
parser.add_argument(
    '-i', '--input', type=str, help='MDF input file', required=True)
parser.add_argument(
    '-t',
    '--tck',
    type=str,
    help='Manifest file for chosen MDF',
    required=True)
parser.add_argument(
    '-j',
    '--json',
    type=str,
    help='Stream configuration specified as JSON',
    required=True)
parser.add_argument(
    '-c',
    '--config',
    type=str,
    help='Choose 5 (five) or 16 (sixteen) stream configuration',
    choices=['five', 'sixteen'],
    required=True)
parser.add_argument(
    '-s',
    '--setting',
    type=str,
    help='Choose line or stream',
    choices=['line', 'stream'],
    required=True)

args = parser.parse_args()

# EvtMax can be set arbitrarily high, will quit if there are no more events
LHCbApp(
    DataType="Upgrade",
    Simulation=True,
    DDDBtag="dddb-20171126",
    CondDBtag="sim-20171127-vc-md100",
    EvtMax=int(os.environ.get('MOORE_EVTMAX')))

IODataManager(DisablePFNWarning=True)
manifest = load_tck(args.tck)
algs = do_unpacking(manifest, input_process='Hlt2')
algsnew = algs[0:4] + [algs[-1]]

decode_hlt2 = HltDecReportsDecoder(
    name="Hlt2DecReportsDecoder",
    SourceID="Hlt2",
    OutputHltDecReportsLocation="/Event/Hlt2/DecReports",
)

appMgr = ApplicationMgr(TopAlg=algsnew)
appMgr.ExtSvc += [configured_ann_svc()]

file = args.input
IOHelper("MDF").inputFiles([file])

with open(args.json) as f:
    config = json.load(f)

take_all_banks = ['turcal', 'ift', 'pid', 'trackeff']

if args.config == 'five' and args.setting == 'line':

    configname = '5streams'
    stream = str(re.search("-(?!.*-)(.*).mdf", file).group(
        1))  # Finds string between last - and .mdf suffix = stream identifier

    lines = config[stream]

    if stream in take_all_banks:
        all_banks = 1
    else:
        all_banks = 0

    appMgr = GP.AppMgr()
    evt = appMgr.evtsvc()

    # Calculate rates per-line for all lines in specified stream/file
    event_stats, exclusive, raw, dst = processing_events_per_line(
        LHCbApp().EvtMax, all_banks, lines)
    rates_per_line(event_stats, exclusive, raw, dst, configname, stream)

elif args.config == 'five' and args.setting == 'stream':

    configname = '5streams'
    stream = str(re.search("-(?!.*-)(.*).mdf", file).group(
        1))  # Finds string between last - and .mdf suffix = stream identifier

    lines = config[stream]

    if stream in take_all_banks:
        all_banks = 1
    else:
        all_banks = 0

    appMgr = GP.AppMgr()
    evt = appMgr.evtsvc()

    # Calculate rates per stream
    events, raw_size, dst_size = processing_events_per_stream(
        LHCbApp().EvtMax, all_banks)
    rates_per_stream(events, raw_size, dst_size, configname, stream)

else:

    configname = '16streams'
    stream = str(re.search("-(?!.*-)(.*).mdf", file).group(
        1))  # Finds string between last - and .mdf suffix = stream identifier

    lines = config[stream]

    if stream in take_all_banks:
        all_banks = 1
    else:
        all_banks = 0

    appMgr = GP.AppMgr()
    evt = appMgr.evtsvc()

    # Calculate rates per stream
    events, raw_size, dst_size = processing_events_per_stream(
        LHCbApp().EvtMax, all_banks)
    rates_per_stream(events, raw_size, dst_size, configname, stream)