#!/usr/bin/env python

import os
import types

import ROOT
from ROOT import gROOT, TH1F
try:
    RATROOT = os.environ['RATROOT']
    from ROOT import RAT
except ImportError:
    # Detect if we already loaded the dictionary, for example inside RAT already
    if ROOT.TClassTable.GetID("RAT::DS::Entry") == -1:
        ROOT.gROOT.ProcessLine('.x ' + os.path.join(RATROOT, 'rootinit.C'))
    from ROOT import RAT

import rat.parser


def dsreader(*filenames):
    """Read the RAT::DS data structures from a ROOT file(s)

        dsreader(filename)
        dsreader(filename0, filename1, filename2, [...], filenameN)

    Returns an iterator over the RAT::DS objects stored in the files
    given comma separated. Assumes the files were generated by rat and
    has the TTree named 'T' and the event branch named 'ds', and all
    filenames given are related.
    """
    for filename in filenames:
        reader = RAT.DU.DSReader(filename)
        for ievent in range(0, reader.GetEntryCount()):
            yield reader.GetEntry(ievent), reader.GetRun()


def socreader(*filenames):
    """Read the RAT::DS::Run data structure from a SOC ROOT file(s)

    Returns the DS::Run instance. Assumes all filenames given are
    related.
    """
    for filename in filenames:
        reader = RAT.DU.SOCReader(filename)
        for isoc in range(0, reader.GetSOCCount()):
            yield reader.GetSOC(isoc), reader.GetRun()


def utility():
    """ Return the rat data utility class, a convienience method."""
    return RAT.DU.Utility.Get()


def ratiter(element, selector=''):
    """Creates iterator over arbitrary member of data structure.

    The selector describes the element to loop over.  It is a
    dotted hierarchy of data structure names.  For example:
      mc.particle.px
      ev.pmt
      ev.efit.ke

    The selector is applied to element.  If element is a string, then
    it is assumed to be a filename, and that file is opened, and the
    selector applied to each RAT.DS event in the file.  You may also
    pass an object in the data structure (like a RAT.MC), in which
    case the selector should be relative to that.  The previous example
    would become:
      particle.px
    in such a scenario.

    The empty selector "" will generate an iterator that just contains
    one item, the element, if it is not a C++ vector.  If it is a vector
    then iterator will return each item in the vector sequentially.
    """
    evaluation_tree = rat.parser.create_evaluation_tree(*selector.split(':'))

    if isinstance(element, types.StringTypes):
        # If element is a string, assume it is a filename and
        # start iterating through events
        for ds, run in dsreader(element):
            for row in evaluation_tree.eval(ds):
                if len(row) == 1:
                    yield row[0]
                else:
                    yield row
    else:
        for row in evaluation_tree.eval(element):
            if len(row) == 1:
                yield row[0]
            else:
                yield row

# Only create these functions if numpy exists
try:
    import numpy as np

    def ratarray(element, selector="", dtype=float):
        """Returns a 1D or 2D numpy array of the values that would be
        returned by ratiter(element, selector)."""
        if isinstance(element, (list, tuple)):
            return parallel_ratarray(element, selector, dtype)
        else:
            return np.array(list(ratiter(element, selector)), dtype=dtype)

    # This function requires the multiprocessing module (Python 2.6 and later)
    import multiprocessing

    def parallel_ratarray(filenames, selector, dtype=float):
        """Helper function to ratiter() that is called when a list of
        filenames is passed to ratiter().

        Spawns multiple processes to loop over the files in parallel.
        Warm up your SSDs!
        """
        RATITER_POOL = multiprocessing.Pool()

        results = RATITER_POOL.map(get_array, [(f,selector,dtype) 
                                               for f in filenames])
        return np.concatenate(results)

    def get_array(args):
        filename, selector, dtype = args
        return ratarray(filename, selector, dtype=dtype)

except ImportError:
    pass


def lookup(name):
    """Looks up an object in gDirectory"""
    return ROOT.gDirectory.Get(name)


def browse():
    """Create a new TBrowser and return it"""
    return ROOT.TBrowser()