from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import sys
from array import array

from ratproc.base import Processor
from rat.parser import create_evaluation_tree
from rat import ROOT, RAT

class NtupleWriter(object):
    '''
    Writes out a ROOT ntuple.
    '''
    def __init__(self, ntuple, field_selectors, double_prec):
        self.ntuple = ntuple
        self.field_selectors = field_selectors
        if double_prec:
            self.field_contents = array('d', [0.0] * len(self.field_selectors))
        else:
            self.field_contents = array('f', [0.0] * len(self.field_selectors))
        self.count = 0
        self.eval_tree = create_evaluation_tree(*field_selectors)

    def fill(self, ds):
        '''
        Fill the ntuple given a RAT::DS::Entry object.
        '''
        # Add a try block to avoid silent errors.
        try:
            for row in self.eval_tree.eval(ds):
                for i, value in enumerate(row):
                    if value is None:
                        value = -9999999.0
                    self.field_contents[i] = value
                self.ntuple.Fill(self.field_contents)
                self.count += 1

        except IOError as e:
            print("ratproc.NtupleWriter : I/O error while processing DS entry : [{0}] = {1}".format(e.errno, e.strerror))
            print("ratproc.NtupleWriter : ERROR: The ntuple will be empty.\n\n")
            raise e

        except Exception as e:
            print("ratproc.NtupleWriter : Unexpected error:", sys.exc_info()[0])
            raise e

    def write(self):
        '''
        Write the ntuple to a file.
        Assumes that one is in the directory of the TFile to be written to.
        '''
        self.ntuple.Write()

class Ntuple(Processor):
    '''
    Creates a list of ROOT ntuples.
    '''
    # Using kwargs because Python 2...
    def __init__(self, *ntuples, **kwargs):
        Processor.__init__(self)

        # Get the file name, use default if not specified.
        filename = kwargs.get("filename")
        if filename is None:
            db = RAT.DB.Get()
            dblink = db.GetLink("IO")
            filename = dblink.GetS("default_output_filename")

        self.double_prec = kwargs.get("d_prec", True)

        print('ratproc.Ntuple : Writing to {}'.format(filename))
        self.f = ROOT.TFile.Open(filename, "RECREATE")

        # Create the list of Ntuple objects for each ntuple name.
        self.ntuples = [self.create_ntuple(name) for name in ntuples]

    def __del__(self):
        if self.f.IsOpen():
            self.f.Close()

    def create_ntuple(self, name):
        '''
        Create a ROOT ntuple and an NtupleWriter object.
        '''
        db = RAT.DB.Get()
        lntuple = db.GetLink("NTUPLE", name)

        # Python lists are better than STL vectors
        fields = list(lntuple.GetSArray("fields"))

        # Deinterleave field names and selectors
        field_names = fields[::2]
        field_selectors = fields[1::2]

        assert len(field_names) == len(field_selectors)

        if not self.double_prec:
            N = ROOT.TNtuple(name, "RAT reduced ntuple", ":".join(field_names))
        else:
            N = ROOT.TNtupleD(name, "RAT reduced ntupleD", ":".join(field_names))

        return NtupleWriter(ntuple=N, field_selectors=field_selectors, double_prec=self.double_prec)

    def dsevent(self, run, ds):
        for writer in self.ntuples:
            writer.fill(ds)

        return Processor.OK

    def finish(self):
        self.f.cd()

        for writer in self.ntuples:
            print('ratproc.Ntuple : Wrote {} entries to {}'.format(writer.count, writer.ntuple.GetName()))
            writer.write()

        self.f.Close()