#!/usr/bin/env python
# =============================================================================
# $Id: summary.py,v 1.11 2010-04-30 13:50:50 rlambert Exp $
# =============================================================================
"""
*******************************************************************************
*        Class to contain the python xmlsummary object                        *
*        this simplifies the interface to the schema parser                   *
*        it can also be taken as an example of how to interact                *
*        with the xml schema/VTree object                                     *
*                                                                             *
*        By necessity, the summary object hard-codes the name of several      *
*        of the tags and attributes of the LHCb schema.                       *
*        This is required for merging and intelligent handling of the objects *
*                                                                             *
*                                                                             *
* Print: returns an xml string                                                *
* write: sends output to a file                                               *
* set_step: is used to set 'initialize'/'finalize' etc.                       *
* fill_file:    used to add information about a file                          *
* fill_counter: used to create a counter in the xml                           *
* fill_lumi:    used to create a counter to be stored under lumiCounters      *
* fill_memory:  used to add memory usage information                          *
* file_dict:    returns dictionary of input/output files by name and GUID     *
* count_dict:   returns dictionary of counters/lumiCounters by name           *
*******************************************************************************
"""
# =============================================================================
__author__ = 'Rob Lambert Rob.Lambert@cern.ch'
# =============================================================================

import os as __os__
# print __file__
__schema__ = None

try:
    import imp as __imp__
    __schema__ = __imp__.load_source('schema', __os__.sep.join(__file__.split(__os__.sep)[:-1] + ['schema.py']))
    # print 'loaded from file'
except (ImportError, IOError):
    try:
        from XMLSummaryBase import schema as __schema__
    except ImportError:
        import schema as __schema__


# import the schema-handling
# def cleverImport(path=None):
#    global __schema__
#    if __schema__ is not None:
#        return
#    if path is None:
# try:
#    from XMLSummaryBase import schema as __schema__
# except ImportError:
#    import schema as __schema__
    # finally fail here if module cannot be found!

#    import imp
#    __schema__=imp.load_module('schema',__os__.path.expanduser(__os__.path.expandvars(path)))
#    return

# from __schema_module__ import *

__default_schema__ = '$XMLSUMMARYBASEROOT/xml/XMLSummary.xsd'
__file_tag__ = 'file'
__count_tag__ = 'counter'

# =============================================================================
# Simple algorithm which manipulates with counters
# =============================================================================


class Summary(__schema__.VTree):
    """ Simple helper functions for writing the xml summary in python
    This class inherits from the VTree object, and has some hard coded element names
    merging might also be done here"""
    # a dictionary of name:object
    # __count_dict__={}
    # a dictionary of name:object
    # __file_dict__={}

    def __init__(self, schemafile=__default_schema__, construct_default=True):
        """ Constructor. Variables in schemafile are expanded.
        construct_default should always be left true unless you wish to parse a file straight away.
        module path optionally imports the schema module from a different path if it isn't already here"""
        # print 'crap'
        if 'Schema' not in str(type(schemafile)):
            # print str(type(schemafile))
            # print "I think it's a string"
            self.__schema__ = __schema__.Schema(schemafile)
        else:
            self.__schema__ = schemafile
        if construct_default:
            self.__element__ = self.__schema__.create_default(self.__schema__.root()).__element__
        else:
            self.__element__ = None
        self.__count_dict__ = {}
        for mother in self.__schema__.Tag_mothers(__count_tag__):
            self.__count_dict__[mother] = {}
        self.__file_dict__ = {}
        for mother in self.__schema__.Tag_mothers(__file_tag__):
            self.__file_dict__[mother] = {}
        super(Summary, self).__init__(element=self.__element__, schema=self.__schema__, mother=None, check=False)

    def __file_exists__(self, GUID, filename, isOutput=False):
        '''internal method, return the file if it already exists'''
        # mothers=self.__schema__.Tag_mothers(__file_tag__)
        self.file_dict()
        mother = 'input'
        if isOutput:
            mother = 'output'

        # if it exists, return it
        if (GUID is not None) and (GUID in self.__file_dict__[mother]):
            return self.__file_dict__[mother][GUID]
        elif (filename is not None) and (filename in self.__file_dict__[mother]):
            return self.__file_dict__[mother][filename]

        return None

    def __file_merger__(self, destination, filename=None, GUID=None, status=None, addevents=0, isOutput=False):
        '''merge two file objects'''
        mothers = self.__schema__.Tag_mothers(__file_tag__)
        mother = 'input'
        if isOutput:
            mother = 'output'

        # update filename or GUID if only one is present
        if (GUID is not None) and (filename is not None):
            if (destination.attrib('name') == ""):
                destination.attrib('name', filename)
                self.__file_dict__[mother][filename] = destination
            if (destination.attrib('GUID') == ""):
                destination.attrib('GUID', GUID)
                self.__file_dict__[mother][GUID] = destination

        # because it's a reference I can do this:
        if status is not None:
            if destination.__element__.attrib['status'] == 'fail':
                pass
                # a fail is a fail
            elif status == 'fail':
                # a fail is a fail
                destination.__element__.attrib['status'] = status
            elif destination.__element__.attrib['status'] == 'mult':
                # mult is mult!
                pass
            elif destination.__element__.attrib['status'] == 'full':  # and status!='full':
                # more than one lot of this file!
                destination.__element__.attrib['status'] = 'mult'
            elif status == 'none':
                # none is default, don't overwrite anything
                pass
                # destination.attrib('status',status)
            else:
                destination.__element__.attrib['status'] = status
        # print 'at end of merge for file', filename, 'adding', addevents, 'to', destination.value()
        # destination.value(destination.value()+addevents)
        # print destination.value()
        destination.__element__.text = str(int(float(destination.__element__.text)) + int(addevents))
        return True

    def __fill_file__(self, filename=None, GUID=None, status=None, addevents=0, isOutput=False):
        '''Internal common method to fill the file processed information
        filename and GUID are the file to find/create
        status is a new status to set
        the number of events is taken from addevents'''
        if filename == "":
            filename = None
        if GUID == "":
            GUID = None
        if status == "":
            status = None

        if filename is None and GUID is None:
            return False

        mothers = self.__schema__.Tag_mothers(__file_tag__)
        mother = 'input'
        if isOutput:
            mother = 'output'

        if filename is not None:
            if "LFN:" not in filename.upper():
                if "PFN:" not in filename.upper():
                    filename = "PFN:" + filename
        open_file = self.__file_exists__(GUID, filename, isOutput)

        # hardcoded tag for file entry :S

        if open_file is None:
            # file needs to be created
            open_file = self.__schema__.create_default(__file_tag__)
            # print 'default created'
            if (status is None):
                open_file.attrib('status', 'none')
            else:
                open_file.attrib('status', status)
            open_file.value(addevents)
            files = self.children('input')[0]
            if isOutput:
                files = self.children('output')[0]

            if filename is not None:
                open_file.attrib('name', filename)
                self.__file_dict__[mother][filename] = open_file
            if GUID is not None:
                open_file.attrib('GUID', GUID)
                self.__file_dict__[mother][GUID] = open_file
            files.__append_element__(open_file)
            return True

        return self.__file_merger__(open_file, filename,
                                    GUID, status,
                                    addevents, isOutput)

    def __recache__(self):
        self.counter_dict(True)
        self.file_dict(True)

    def __buildcounter__(self, name, flag, nEntries=None, flag2=None, min=None, max=None):
        '''method to build a VTree counter'''
        # pick the daughter to fill
        type = __count_tag__
        attrib = "name"

        # if isLumi:
        #    if nEntries is not None: type="lumiStatEntity"
        #    else: type="lumiCounter"
        if nEntries is not None:
            type = "statEntity"

        # create the xml object to fill
        cnt = self.__schema__.create_default(type)
        cnt.attrib(attrib, name)
        # if it's a simple counter...
        if nEntries is None:
            cnt.value(flag)
            return cnt
        # else it's a statEntity
        val = []
        for f in cnt.attrib('format'):
            if f == 'Flag':
                val.append(flag)
            elif f == 'Flag2':
                val.append(flag2)
            elif f == 'Entries':
                val.append(nEntries)
        cnt.value(val)
        cnt.attrib('max', max)
        cnt.attrib('min', min)

        return cnt

    def fill_VTree_file(self, file, isOutput=False):
        '''append or merge a vtree file into the tree'''

        mothers = self.__schema__.Tag_mothers(__file_tag__)
        mother = 'input'
        if isOutput:
            mother = 'output'

        if file.__element__.attrib['name'] != "":
            if "LFN:" not in file.__element__.attrib['name'].upper():
                if "PFN:" not in file.__element__.attrib['name'].upper():
                    file.__element__.attrib['name'] = "PFN:" + file.__element__.attrib['name']

        filename = file.__element__.attrib['name']
        GUID = file.__element__.attrib['GUID']
        open_file = self.__file_exists__(GUID, filename, isOutput)

        # if it doesn't exist, just add it
        if open_file is None:
            files = self.children(mother)[0]
            a = file.clone()
            files.__append_element__(a)
            self.__file_dict__[mother][filename] = a
            self.__file_dict__[mother][GUID] = a
            return True

        # else merge:
        return self.__file_merger__(open_file, filename,
                                    GUID, file.__element__.attrib['status'],
                                    int(file.__element__.text), isOutput)

    def isFailure(self):
        """the opposite of isSuccess, negates what is stored by the success tag"""
        return (self.children('success')[0].value() is False)

    def set_step(self, step="initialize", success=False):
        """ Set the step entry, and the success entry """
        test = True
        istep = self.children('step')[0]
        test = (istep.value(step) and test)

        # example setting success
        isuccess = self.children('success')[0]
        test = (isuccess.value(success) and test)

        return test

    def parse(self, xmlfile):
        """ parse an existing xml file, and load it into THIS object """
        self.__element__ = self.__schema__.parse(xmlfile).__element__
        self.__recache__()

    def fill_output(self, filename=None, GUID=None, status=None, addevents=0):
        '''Method to fill the file processed information
        filename and GUID are the file to find/create
        status is a new status to set
        the number of events is taken from addevents'''
        return self.__fill_file__(filename, GUID, status, addevents, True)

    def fill_input(self, filename=None, GUID=None, status=None, addevents=0):
        '''Method to fill the file processed information
        filename and GUID are the file to find/create
        status is a new status to set
        the number of events is taken from addevents'''
        return self.__fill_file__(filename, GUID, status, addevents, False)

    def fill_VTree_counter(self, counter, isLumi=False):
        ''' fill a basic or complex stat counter, or lumi counter, from a VTree
        will merge the counter if it already exists.
        If you want to add, but not merge use the add method'''
        if 'VTree' not in str(type(counter)):
            raise TypeError('expected VTree, got ' + str(type(counter)) + ' instead')
        if counter.tag() not in ['statEntity', __count_tag__]:
            raise TypeError('expected counter, got ' + counter.tag() + ' instead')
        mother = 'counters'  # self.__schema__.Tag_mothers(counter.tag())[0]
        if isLumi:
            mother = 'lumiCounters'  # self.__schema__.Tag_mothers(counter.tag())[1]

        # this takes a lot of time, why is it needed??
        # self.counter_dict(True)

        # hard coded, no way around it :S
        # attrib='name'
        name = counter.__element__.attrib['name']

        # check if it exists. If not, just add it
        # print self.__count_dict__
        # print self.__count_dict__[mother]
        if name not in self.__count_dict__[mother]:
            counters = __schema__.VTree([bt for bt in self.__element__.getchildren()
                                         if mother in bt.tag][0], self.__schema__, self, False)
            # counter=counter.clone()
            counters.__append_element__(counter)
            self.__count_dict__[mother][name] = counter
            return True

        # counters=self.children(mother)[0]
        # check if it exists. If not, just add it
        # cnt=counters.children(counter.tag(),attrib={attrib:name})
        # if cnt is None or len(cnt)==0:
        #    counters.add(counter)
        #    return True
        cnt = self.__count_dict__[mother][name]
        # else, merge the two
        if 'statentity' not in counter.tag().lower():
            cnt.__element__.text = str(int(cnt.__element__.text) + int(counter.__element__.text))
            return True
        # else be clever about the format
        val = cnt.value()
        for i in range(len(cnt.attrib('format'))):
            for j in range(len(counter.attrib('format'))):
                if cnt.attrib('format')[i] == counter.attrib('format')[j]:
                    # all entries are additive
                    val[i] = val[i] + counter.value()[j]
                    break
        cnt.value(val)
        if cnt.attrib('max') < counter.attrib('max'):
            cnt.attrib('max', counter.attrib('max'))
        if cnt.attrib('min') > counter.attrib('min'):
            cnt.attrib('min', counter.attrib('min'))

    def fill_counter(self, name, flag, nEntries=None, flag2=None, min=None, max=None):
        ''' fill a basic or complex stat counter '''
        cnt = self.__buildcounter__(name, flag, nEntries, flag2, min, max)
        return self.fill_VTree_counter(cnt, False)

    def fill_lumi(self, name, flag, nEntries=None, flag2=None, min=None, max=None):
        ''' fill a basic or complex Lumi stat counter '''
        cnt = self.__buildcounter__(name, flag, nEntries, flag2, min, max)
        return self.fill_VTree_counter(cnt, True)

    def fill_memory(self, memory, unit='b'):
        '''Method to fill the memory usage information'''
        usage = self.children('usage')[0]
        astat = usage.children('stat', attrib={'useOf': 'MemoryMaximum'})
        if not astat or len(astat) == 0:
            astat = self.__schema__.create_default('stat')
            astat.attrib('useOf', 'MemoryMaximum')
            astat.attrib('unit', unit)
            astat.value(memory)
            usage.add(astat)
        else:
            astat = astat[0]
            if astat.attrib('unit') != unit:
                raise AttributeError(
                    'I cannot compare two MemoryMaxima when they have different units! '
                    + astat.attrib('unit')
                    + " "
                    + unit)
            if astat.value() < memory:
                astat.value(memory)

    def counter_dict(self, recache=False):
        '''return a name:VTree dictionary of counters'''
        mothers = self.__schema__.Tag_mothers(__count_tag__)
        empty = True
        if self.__count_dict__ != {}:
            for mother in mothers:
                if self.__count_dict__[mother] != {}:
                    empty = False
        if empty or recache:
            for mother in mothers:
                self.__count_dict__[mother].clear()
                for amother in self.children(mother):
                    for count in amother.children():
                        self.__count_dict__[mother][count.attrib('name')] = count
        return self.__count_dict__

    def file_dict(self, recache=False):
        '''return a name:VTree dictionary of all input and output files'''
        mothers = self.__schema__.Tag_mothers(__file_tag__)
        empty = True
        if self.__file_dict__ != {}:
            for mother in mothers:
                if self.__file_dict__[mother] != {}:
                    empty = False
        if empty or recache:
            for mother in mothers:
                self.__file_dict__[mother].clear()
                for amother in self.children(mother):
                    for file in amother.children():
                        if (file.attrib('name') != ""):
                            self.__file_dict__[mother][file.attrib('name')] = file
                        if (file.attrib('GUID') != ""):
                            self.__file_dict__[mother][file.attrib('GUID')] = file
        return self.__file_dict__


def Merge(summaries, schemafile=__default_schema__):
    '''Merge a list of summaries, return a new summary
    summaries can be a list of xml files to be parsed, or a list of summary objects'''
    if isinstance("", type(schemafile)):
        schemafile = __schema__.Schema(schemafile)
    if not isinstance([], type(summaries)):
        raise TypeError('you should send a list into the merger, I got a ' + str(type(summaries)) + ' object instead')
    sum_objects = []

    if isinstance("", type(summaries[0])) and str(type(Summary(schemafile=schemafile))) == str(type(summaries[0])):
        raise TypeError('you should send strings or Summaries into the merger, I got a '
                        + str(type(summaries[0])) + ' object instead')

    if str(type(Summary(schemafile=schemafile))) == str(type(summaries[0])):
        sum_objects = summaries
        # check they all have the same schema
        schema = sum_objects[0].__schema__.__schemafile_short__
        for asummary in sum_objects:
            if asummary.__schema__.__schemafile_short__ != schema:
                raise AttributeError('these xml files have different schema so cannot be merged')

    # try-except, in case you haven't given real summaries!

    try:

        if isinstance("", type(summaries[0])):
            for asummary in summaries:
                # print asummary
                sum_object = Summary(schemafile, construct_default=False)
                # don't check every summary, unless the merging fails
                # sum_object.parse(asummary)
                asummary = __os__.path.expanduser(__os__.path.expandvars(asummary))
                if not __os__.path.exists(asummary):
                    raise IOError('file does not exist ' + str(asummary))
                sum_object.__element__ = schemafile.__fast_parse__(asummary).__element__
                sum_object.__recache__()
                # print sum_object
                sum_objects.append(sum_object)

        # make default object
        merged = Summary(sum_objects[0].__schema__)
        # print 'made default object'
        # merge success
        # print 'merge success'
        flag = True
        for asummary in sum_objects:
            if asummary.isFailure():
                flag = False
                break
        merged.children('success')[0].value(flag)
        # merge step
        # print 'merge step'
        # enum will be in order
        steps = merged.__schema__.Tag_enumeration(merged.children('step')[0].tag())
        flag = ''
        for asummary in sum_objects:
            if flag == asummary.children('step')[0].__element__.text:
                continue
            for step in steps:
                # retain the lowest possible step
                if flag == step:
                    break
                if step == asummary.children('step')[0].__element__.text:
                    flag = step
                    break
        merged.children('step')[0].__element__.text = flag
        # merge input/output, simple counters, usage
        # print 'merge ip/op simple counters, usage'
        counters = {}
        # lumiCounters={}
        for asummary in sum_objects:
            # print sum_objects.index(asummary)
            # print 'usage'
            for stat in asummary.children('usage')[0].children('stat'):
                merged.fill_memory(stat.value(), stat.__element__.attrib['unit'])
            # print 'input'
            for file in asummary.children('input')[0].children():
                # merged.fill_VTree_file(file,False)
                merged.fill_input(file.__element__.attrib['name'],
                                  file.__element__.attrib['GUID'],
                                  file.__element__.attrib['status'],
                                  int(file.__element__.text))
            # print 'output'
            for file in asummary.children('output')[0].children():
                # merged.fill_VTree_file(file,True)
                merged.fill_output(file.__element__.attrib['name'],
                                   file.__element__.attrib['GUID'],
                                   file.__element__.attrib['status'],
                                   int(file.__element__.text))
            # merge counters
            # print 'counters'
            for cnt in asummary.children('counters')[0].children(__count_tag__):
                # print 'counter'
                # print cnt
                name = cnt.__element__.attrib['name']
                if name not in counters:
                    counters[name] = int(cnt.__element__.text)
                else:
                    counters[name] = (counters[name] + int(cnt.__element__.text))
                # if name not in merged.__count_dict__['counters']:
                #    cnt=cnt.clone()
                # merged.fill_VTree_counter(cnt)
            # merge counters
            # print 'lumiCounters'
            for cnt in asummary.children('lumiCounters')[0].children(__count_tag__):
                name = cnt.__element__.attrib['name']
                if name not in merged.__count_dict__['lumiCounters']:
                    cnt = cnt.clone()
                merged.fill_VTree_counter(cnt, isLumi=True)
        # merge collected counters
        for c in counters:
            merged.fill_counter(c, counters[c])
        # print 'merge statCounters'
        for asummary in sum_objects:
            # print sum_objects.index(asummary)
            # merge statCounters
            for cnt in asummary.children('counters')[0].children('statEntity'):
                merged.fill_VTree_counter(cnt.clone())
            # merge counters
            for cnt in asummary.children('lumiCounters')[0].children('statEntity'):
                merged.fill_VTree_counter(cnt.clone(), isLumi=True)
        if merged.__schema__.__check__(merged.__element__):
            return merged
        else:
            raise AttributeError('merged file could not be verified')
    except BaseException:
        print('Something went horribly wrong, I will now check all your summaries')
        for asummary in sum_objects:
            asummary.__schema__.__check__(asummary.__element__)
        raise
    return None
