#!/usr/bin/env python
from optparse import OptionParser

"""
Need to keep this script compatible with Python 2.4
and only dependent on the default PyROOT, not rootpy.
This script must be able to run alone in a very limited environment.
"""

parser = OptionParser()
parser.add_option('-r','--branches-regex',action="store",type="str",dest="branch_regex",default=None,
                  help="file containing one regex per line for branch selection")
parser.add_option('-g','--branches-glob',action="store",type="str",dest="branch_glob",default=None,
                  help="file containing one glob per line for branch selection")
parser.add_option('-b','--branches',action="store",type="str",dest="branch_literal",default=None,
                  help="file containing one branch name per line for branch selection")
parser.add_option('-t','--tree',action="store",type="str",dest="tree",default=None,
                  help="name of tree (including path) to optimize common to all files")
parser.add_option('--prefix',action="store",type="str",dest="prefix",default="skimmed",
                  help="prefix of output file/directory name")
parser.add_option('-s','--selection',action="store",type="str",dest="selection",default='',
                  help="only entries satisfying this cut will be kept")
parser.add_option('-o','--or-if-exists',action="store",type="str",dest="or_if_exists",default=None,
                  help="OR conditions on all branches existing in current file. " \
                       "List one branch name and condition on branch separated by whitespace per line")
parser.add_option('--grid',action="store_true",dest="grid",default=False,
                  help="Running on the grid mode. Input files are separated by commas.")
parser.add_option('--overwrite',action="store_true",dest="overwrite",default=False,
                  help="Overwrite existing output.")

options, args = parser.parse_args()

import sys

if options.tree is None:
    sys.exit("You must specify a tree name")

if not options.prefix:
    sys.exit("Invalid prefix")

if not args:
    sys.exit("No input ROOT files specified")

if options.grid:
    if len(args) != 1:
        sys.exit("Expected exactly one argument containing list of files separated by commas")
    args = args[0].split(',')

import ROOT
ROOT.PyConfig.IgnoreCommandLineOptions = True
import re
import os
import fnmatch
import uuid

branches_keep_regex = []
if options.branch_regex is not None:
    try:
        branches_file = open(options.branch_regex,'r')
        branches_keep_regex += [re.compile(regex.strip()) for regex in branches_file.readlines()]
        branches_file.close()
    except Exception, ex:
        print ex
        sys.exit("Could not parse branch regex file %s"% options.branch_regex)

if options.branch_glob is not None:
    try:
        branches_file = open(options.branch_glob,'r')
        branches_keep_regex += [re.compile(fnmatch.translate(regex.strip())) for regex in branches_file.readlines()]
        branches_file.close()
    except Exception, ex:
        print ex
        sys.exit("Could not parse branch glob file %s"% options.branch_regex)

branches_keep_literal = []
if options.branch_literal is not None:
    try:
        branches_file = open(options.branch_literal,'r')
        branches_keep_literal = [line.strip() for line in branches_file.readlines()]
        branches_file.close()
    except Exception, ex:
        print ex
        sys.exit("Could not parse branch file %s"% options.branch_literal)

if options.or_if_exists:
    or_file = open(options.or_if_exists, 'r')
    options.or_if_exists = [line.strip() for line in or_file.readlines()]
    or_file.close()
        
infiles = args

for infilename in infiles:
    sub_infiles = [(infilename, "%s."% options.prefix)]
    isdir = False
    if os.path.isdir(infilename):
        isdir = True
        infilename = os.path.normpath(infilename)
        sub_infiles = []
        for root, dirnames, filenames in os.walk(infilename):
            base = ("%s."% options.prefix) + os.path.basename(root)
            for filename in fnmatch.filter(filenames, '*.root*'):
                sub_infiles.append((os.path.join(root, filename), os.path.join(os.path.dirname(root), base) + os.sep))
    for filename, out_path in sub_infiles:
        if isdir:
            if not os.path.exists(out_path):
                os.makedirs(out_path)
        outfilename = out_path+os.path.basename(filename)
        if os.path.exists(outfilename):
            if options.overwrite:
                print "overwriting %s" % outfilename
                outFile = ROOT.TFile(outfilename,'RECREATE')

            else:
                print "skipping %s which already exists"% outfilename
                continue
        else:
            outFile = ROOT.TFile(outfilename,'NEW')
        inFile = ROOT.TFile.Open(filename)
        if not inFile:
            print "skipping %s which could not be opened"% filename
            continue
        print "skimming %s ..." % filename
        try:
            tree = inFile.Get(options.tree)
            if not tree:
                print "skipping %s where tree %s could not be found" % (
                        filename, options.tree)
                inFile.Close()
                outFile.Close()
                os.remove(outfilename)
                continue
            if tree.GetNbranches() == 0:
                print "skipping %s where tree %s does not contain any branches"% (
                        filename, options.tree)
                inFile.Close()
                outFile.Close()
                os.remove(outfilename)
                continue
            if branches_keep_regex:
                tree.SetBranchStatus('*',0)
                branches = [b.GetName() for b in tree.GetListOfBranches()]
                for branch in branches:
                    for branch_regex in branches_keep_regex:
                        if re.search(branch_regex, branch):
                            print "keeping branch %s" % branch
                            tree.SetBranchStatus(branch, 1)
                            break
            if branches_keep_literal:
                tree.SetBranchStatus('*',0)
                for branch in branches_keep_literal:
                    if tree.GetBranch(branch):
                        print "keeping branch %s" % branch
                        tree.SetBranchStatus(branch, 1)
            selection = options.selection
            if options.or_if_exists:
                for branch in options.or_if_exists:
                    if tree.GetBranch(branch):
                        tree.SetBranchStatus(branch, 1)
                        if selection:
                            selection = '(%s)||(%s)' % (selection, branch)
                        else:
                            selection = branch
            if selection:
                tempfile = ROOT.TFile(uuid.uuid4().hex+".root",'NEW')
                tempfile.cd()
                print "applying selection %s on tree %s in %s ..." % \
                      (selection, tree.GetName(), filename)
                tree = tree.CopyTree(selection)
                tempfile.Write()
                tempfile.Flush()
            outFile.cd()
            print "optimizing tree %s in %s ..." % (tree.GetName(), filename)
            newTree = tree.CloneTree(-1, "fast SortBasketsByEntry")
            newTree.OptimizeBaskets()
            newTree.Write()
            outFile.Close()
            inFile.Close()
            if options.selection:
                tempfile.Close()
                os.remove(tempfile.GetName())
        except KeyboardInterrupt, SystemExit:
            print "cleaning up ..."
            try:
                inFile.Close()
                outFile.Close()
                os.remove(outfilename)
                tempfile.Close()
                os.remove(tempfile.GetName())
            except:
                pass
            sys.exit()
