#! /usr/bin/env python
import lighthisto
import logging
import sys

if sys.version_info[:3] < (2,4,0):
    print "rivet scripts require Python version >= 2.4.0... exiting"
    sys.exit(1)

import os, copy, re
from math import sqrt

## Try to load faster but non-standard cElementTree module
try:
    import xml.etree.cElementTree as ET
except ImportError:
    try:
        import cElementTree as ET
    except ImportError:
        try:
            import xml.etree.ElementTree as ET
        except:
            sys.stderr.write("Can't load the ElementTree XML parser: please install it!\n")
            sys.exit(1)



# #############################################

def fillAbove(desthisto, sourcehistosbyptmin):
    for i,b in enumerate(desthisto.getBins()):
        ## Fill bins with pT-ordered histos (so that 'highest always wins')
        for ptmin, h in sorted(sourcehistosbyptmin.iteritems()):
            newb = h.getBin(i)
            if b.xlow >= float(ptmin) :
                b.val =  newb.val
                b.errplus =  newb.errplus
                b.errminus =  newb.errminus
                b._focus= newb._focus

def merge(hpath):
    global inhistos
    global outhistos
    try:
        fillAbove(outhistos[hpath], inhistos[hpath])
    except:
        pass

def useOne(hpath, sqrts):
    global inhistos
    global outhistos
    try:
        outhistos[hpath] =  inhistos[hpath][float(sqrts)]
    except:
        pass
if __name__ == "__main__":
    import logging
    from optparse import OptionParser, OptionGroup
    parser = OptionParser(usage="%prog name")
    verbgroup = OptionGroup(parser, "Verbosity control")
    verbgroup.add_option("-v", "--verbose", action="store_const", const=logging.DEBUG, dest="LOGLEVEL",
                         default=logging.INFO, help="print debug (very verbose) messages")
    verbgroup.add_option("-q", "--quiet", action="store_const", const=logging.WARNING, dest="LOGLEVEL",
                         default=logging.INFO, help="be very quiet")
    parser.add_option_group(verbgroup)
    (opts, args) = parser.parse_args()
    logging.basicConfig(level=opts.LOGLEVEL, format="%(message)s")


    ## Prefix used in dat file headers
    headerprefix = "# "


    ## Check args
    if len(args) < 1:
        logging.error("Must specify at least the name of the files")
        sys.exit(1)

# #######################################

aidafiles=["-e--LowQ2.aida","-e+-LowQ2.aida","-e+-HighQ2.aida"]

## Get histos
inhistos = {}
outhistos={}
weights = {}

for f in aidafiles:
    file = args[0]+f
    if(file.find("Low")> 0) :
        q2=0
    elif(file.find("High")>0) :
        q2=50

    if not os.access(file, os.R_OK):
        logging.error("%s can not be read" % file)
        break
    try:
        tree = ET.parse(file)
    except:
        logging.error("%s can not be parsed as XML" % file)
        break
    tree = ET.parse(file)
    ## Get histos from this AIDA file
    for dps in tree.findall("dataPointSet"):
        h = lighthisto.Histo.fromDPS(dps)
        # di-jet decorrelations
        # jet shapes
        if(h.fullPath().find("4129130")>0 ) :
           if not inhistos.has_key(h.fullPath()):
               inhistos[h.fullPath()] = {}
           tmpE = inhistos[h.fullPath()]
           if not tmpE.has_key(q2):
               tmpE[q2] = h
           else:
                raise Exception("A set with q2    = %s already exists" % ( q2   ))
        elif(h.fullPath().find("2919893")>0 ) :
            outhistos[h.fullPath()] = h
        else :
            print '!!!!! DIDNT FIND',h.fullPath()

## Make empty output histos if needed
for hpath,hsets in inhistos.iteritems():
    if( hpath.find("4129130")>0) :
        outhistos[hpath] = copy.deepcopy(hsets.values()[0])
        ## Empty the bin set for histos which we're going to merge
        for i,b in enumerate(outhistos[hpath].getBins()):
            b.val = 0
            b.errplus = 0
            b.errminus = 0
            b._focus= None

# H1 transverse energy
useOne("/H1_2000_S4129130/d01-x01-y01","0")
useOne("/H1_2000_S4129130/d02-x01-y01","0")
useOne("/H1_2000_S4129130/d03-x01-y01","0")
useOne("/H1_2000_S4129130/d04-x01-y01","0")
useOne("/H1_2000_S4129130/d05-x01-y01","0")
useOne("/H1_2000_S4129130/d06-x01-y01","0")
useOne("/H1_2000_S4129130/d07-x01-y01","0")
useOne("/H1_2000_S4129130/d08-x01-y01","0")
useOne("/H1_2000_S4129130/d09-x01-y01","0")
useOne("/H1_2000_S4129130/d10-x01-y01","0")
useOne("/H1_2000_S4129130/d11-x01-y01","0")
useOne("/H1_2000_S4129130/d12-x01-y01","0")
useOne("/H1_2000_S4129130/d13-x01-y01","0")
useOne("/H1_2000_S4129130/d14-x01-y01","0")
useOne("/H1_2000_S4129130/d15-x01-y01","0")
useOne("/H1_2000_S4129130/d16-x01-y01","50")
useOne("/H1_2000_S4129130/d17-x01-y01","50")
useOne("/H1_2000_S4129130/d18-x01-y01","50")
useOne("/H1_2000_S4129130/d19-x01-y01","50")
useOne("/H1_2000_S4129130/d20-x01-y01","50")
useOne("/H1_2000_S4129130/d21-x01-y01","50")
useOne("/H1_2000_S4129130/d22-x01-y01","50")
useOne("/H1_2000_S4129130/d23-x01-y01","50")
useOne("/H1_2000_S4129130/d24-x01-y01","50")
useOne("/H1_2000_S4129130/d25-x01-y01","0")
useOne("/H1_2000_S4129130/d26-x01-y01","0")
useOne("/H1_2000_S4129130/d27-x01-y01","0")
useOne("/H1_2000_S4129130/d28-x01-y01","0")
useOne("/H1_2000_S4129130/d29-x01-y01","50")
useOne("/H1_2000_S4129130/d30-x01-y01","50")
useOne("/H1_2000_S4129130/d31-x01-y01","50")
useOne("/H1_2000_S4129130/d32-x01-y01","50")

merge("/H1_2000_S4129130/d33-x01-y01")
merge("/H1_2000_S4129130/d34-x01-y01")

# Choose output file
name = args[0]+".aida"
out = open(name, "w")
## Write out merged histos
out.write('<?xml version="1.0" encoding="ISO-8859-1" ?>\n')
out.write('<!DOCTYPE aida SYSTEM "http://aida.freehep.org/schemas/3.3/aida.dtd">\n')
out.write('<aida version="3.3">\n')
out.write('  <implementation version="1.1" package="FreeHEP"/>\n')
for hpath, h in sorted(outhistos.iteritems()):
    logging.debug("hpath = %s" % hpath)
    out.write(h.asAIDA() + "\n\n")
out.write('</aida>\n')

sys.exit(0)
