#!/usr/bin/env python

"""  Check tracker efficiency """

import sys

import os

import ROOT

import libMausCpp #pylint: disable = W0611

#pylint: disable = R0902

#pylint: disable = R0912

#pylint: disable = R0914

#pylint: disable = R0915

#pylint: disable = C0103

#pylint: disable = W0612

class PatRecEfficiency():

    """ Class to check Pattern Recognition efficiency with real data """

    def __init__(self):

        """ Initialise member variables """

        self.check_helical = False

        self.check_straight = True

        self.cut_on_tof = True

        self.cut_on_trackers = True

        self.root_files = []

        self.num_total_events = 0

        self.num_12spoint_events = 0

        self.num_10spoint_events = 0

        self.num_5spoint_tkus_events = 0

        self.num_5spoint_tkds_events = 0

        self.num_10spoint_tracks = 0

        self.num_5spoint_tkus_tracks = 0

        self.num_5spoint_tkds_tracks = 0

        self.num_3to5spoint_tkus_tracks = 0

        self.num_3to5spoint_tkds_tracks = 0

    def run(self, files):

        """ Calculate the Pattern Recognition efficiency for the given files """

        print '\nPattern Recognition Efficiency Calculator'

        print '*****************************************\n'

        print 'Parameters:'

        print 'Check helical\t' + str(self.check_helical)

        print 'Check straight\t' + str(self.check_straight)

        print 'Cut on TOF\t' + str(self.cut_on_tof)

        print 'Cut on trackers\t' + str(self.cut_on_trackers) + '\n'

        self.root_files = self.load_data(files)

        if len(self.root_files) < 1:

            return False

        print '\nFile\t\t\t#_Recon_Events\tTkUS_5pt\tTkUS_3-5_pt\t',

        print 'TkDS_5pt\tTkDS_3-5pt'

        for root_file_name in self.root_files:

            self.calculate_efficiency(root_file_name)

            self.print_file_info(root_file_name)

        return True

    def load_data(self, files):

        """ Load data from files. If a dir is given, search recursively """

        root_files = []

        if type(files) is not list:

            files = [files]

        for file_name in files:

            # Check if file_name is a ROOT file

            if os.path.isfile(file_name):

                file_suffix, file_extension = os.path.splitext(file_name)

                if file_extension == '.root':

                    root_files.append(file_name)

                else:

                    print 'Bad file name, aborting'

                    return root_files

            # If file_name is a directory, walk it and save any ROOT files found

            if os.path.isdir(file_name):

                self.root_files_dir_search(file_name, root_files)

            if len(root_files) < 1:

                print 'No data files found'

                return root_files

        print '\nFound ' + str(len(root_files)) + ' ROOT files:'

        for f in root_files:

            print f

        return root_files

    @staticmethod

    def root_files_dir_search(top_dir, root_files):

        """ Appends any ROOT files found in the directory to root_files """

        for dir_name, subdir_list, file_list in os.walk(top_dir):

            print('Searching directory: %s' % dir_name)

            for fname in file_list:

                file_suffix, file_extension = os.path.splitext(fname)

                if file_extension == '.root':

                    print('\t%s' % fname)

                    root_files.append(dir_name + '/' + fname)

    def calculate_efficiency(self, root_file_name):

        """ Calculate the Pattern Recognition efficiency for the input file """

        self.clear_counters() # Start clean each time

        # Load the ROOT file

        # print "Loading ROOT file", root_file_name

        root_file = ROOT.TFile(root_file_name, "READ") # pylint: disable = E1101

        # Set up the data tree to be filled by ROOT IO

        # print "Setting up data tree"

        tree = root_file.Get("Spill")

        data = ROOT.MAUS.Data() # pylint: disable = E1101

        tree.SetBranchAddress("data", data)

        # Loop over spills

        for i in range(tree.GetEntries()):

            tree.GetEntry(i)

            spill = data.GetSpill()

            # Print some basic information about the spill

            # print "Found spill number", spill.GetSpillNumber(),

            # print "in run number", spill.GetRunNumber()

            if spill.GetDaqEventType() != "physics_event":

                continue

            for i in range(spill.GetReconEvents().size()):

                self.num_total_events += 1

                # Pull out tof data

                #------------------

                tof_evt = spill.GetReconEvents()[i].GetTOFEvent()

                tof1 = tof_evt.GetTOFEventSpacePoint().GetTOF1SpacePointArray()

                tof2 = tof_evt.GetTOFEventSpacePoint().GetTOF1SpacePointArray()

                bool_2tof_spoint_event = True

                if ((len(tof1) != 1) or (len(tof2) != 1)):

                    bool_2tof_spoint_event = False

                    if self.cut_on_tof:

                        continue

                # Pull out tracker data

                #----------------------

                # All spacepoint data

                tk_evt = spill.GetReconEvents()[i].GetSciFiEvent()

                tk_spoints = tk_evt.spacepoints()

                bool_10spoint_event = True

                bool_tkus_5spoint_event = True

                bool_tkds_5spoint_event = True

                for i in range(2):

                    tracker = [sp for sp in tk_spoints if sp.get_tracker() == i]

                    for j in range(1, 6):

                        station = \

                          [sp for sp in tracker if sp.get_station() == j]

                        if len(station) != 1:

                            bool_10spoint_event = False

                            if i == 0:

                                bool_tkus_5spoint_event = False

                            elif i == 1:

                                bool_tkds_5spoint_event = False

                if self.cut_on_trackers and (not bool_10spoint_event):

                    continue

                if bool_2tof_spoint_event and bool_10spoint_event:

                    self.num_12spoint_events += 1

                if bool_10spoint_event:

                    self.num_10spoint_events += 1

                if bool_tkus_5spoint_event:

                    self.num_5spoint_tkus_events += 1

                if bool_tkds_5spoint_event:

                    self.num_5spoint_tkds_events += 1

                # Pat Rec seed data

                if (self.check_helical) and (not self.check_straight):

                    prtracks = tk_evt.helicalprtracks()

                    # print 'Looking at helical tracks'

                elif (not self.check_helical) and (self.check_straight):

                    prtracks = tk_evt.straightprtracks()

                    # print 'Looking at straight tracks'

                elif (self.check_helical) and (self.check_straight):

                    prtracks = []

                    prtracks.append(tk_evt.helicalprtracks())

                    prtracks.append(tk_evt.straightprtracks())

                    # print 'Looking at helical and straight tracks'

                else:

                    print 'Both track options set, aborting'

                    return

                # Is there at least 1 track present in either tracker

                bool_tkus_1track = False

                bool_tkds_1track = False

                if len(prtracks) < 1:

                    continue

                num_tkus_tracks = 0

                num_tkds_tracks = 0

                # Require 1 and only 1 track in a tracker

                for trk in prtracks:

                    if trk.get_tracker() == 0:

                        num_tkus_tracks += 1

                    elif trk.get_tracker() == 1:

                        num_tkds_tracks += 1

                if num_tkus_tracks == 1:

                    bool_tkus_1track = True

                if num_tkds_tracks == 1:

                    bool_tkds_1track = True

                # If there were 5 spacepoints in a tracker and 1 track only was

                # found in it, increment the 3 to 5 spoint track counter

                if bool_tkus_5spoint_event and bool_tkus_1track:

                    self.num_3to5spoint_tkus_tracks += 1

                if bool_tkds_5spoint_event and bool_tkds_1track:

                    self.num_3to5spoint_tkds_tracks += 1

                # Now check the tracks found had 5 spoints, 1 from each station

                bool_tkus_5spoint_track = False

                bool_tkds_5spoint_track = False

                for trk in prtracks:

                    if trk.get_tracker() == 0:

                        if (len(trk.get_spacepoints()) == 5) \

                          and bool_tkus_1track:

                            self.num_5spoint_tkus_tracks += 1

                            bool_tkus_5spoint_track = True

                    elif trk.get_tracker() == 1:

                        if (len(trk.get_spacepoints()) == 5) \

                          and bool_tkds_1track:

                            self.num_5spoint_tkds_tracks += 1

                            bool_tkds_5spoint_track = True

                    # Lastly see if we found one 5 spoint track in BOTH trackers

                    if (len(prtracks) != 2):

                        continue

                    if (bool_tkus_5spoint_track and bool_tkds_5spoint_track):

                        self.num_10spoint_tracks += 1

        #print 'Total recon events: ' + str(num_total_events)

        #print '12 spacepoint events: ' + str(num_12spoint_events)

        #print '10 spacepoint events: ' + str(num_10spoint_events)

        #print 'TkUS 5 spoint events: ' + str(num_5spoint_tkus_events)

        #print 'TkDS 5 spoint events: ' + str(num_5spoint_tkds_events)

        #print '10 spoint PR tracks: ' + str(num_10spoint_tracks)

        #print 'TkUS 5 spoint PR tracks: ' + str(num_5spoint_tkus_tracks)

        #print 'TkDS 5 spoint PR tracks: ' + str(num_5spoint_tkds_tracks)

        #print 'TkUS 3to5 spoint PR tracks: ' + str(num_3to5spoint_tkus_tracks)

        #print 'TkDS 3to5 spoint PR tracks: ' + str(num_3to5spoint_tkds_tracks)

        # Close the ROOT file

        # print "Closing root file"

        root_file.Close()

    def clear_counters(self):

        """ Set the internal counters to zero """

        self.num_total_events = 0

        self.num_12spoint_events = 0

        self.num_10spoint_events = 0

        self.num_5spoint_tkus_events = 0

        self.num_5spoint_tkds_events = 0

        self.num_10spoint_tracks = 0

        self.num_5spoint_tkus_tracks = 0

        self.num_5spoint_tkds_tracks = 0

        self.num_3to5spoint_tkus_tracks = 0

        self.num_3to5spoint_tkds_tracks = 0

    def print_file_info(self, root_file_name):

        """ Calculate the efficiencies and print """

        try:

            us_5pt = float(self.num_5spoint_tkus_tracks) \

              / float(self.num_5spoint_tkus_events)

        except ZeroDivisionError:

            us_5pt = 0.0

        try:

            us_3to5pt = float(self.num_3to5spoint_tkus_tracks) \

              / float(self.num_5spoint_tkus_events)

        except ZeroDivisionError:

            us_3to5pt = 0.0

        try:

            ds_5pt = float(self.num_5spoint_tkds_tracks) \

              / float(self.num_5spoint_tkds_events)

        except ZeroDivisionError:

            ds_5pt = 0.0

        try:

            ds_3to5pt = float(self.num_3to5spoint_tkds_tracks) \

              / float(self.num_5spoint_tkds_events)

        except ZeroDivisionError:

            ds_3to5pt = 0.0

        print os.path.basename(root_file_name) + '\t',

        print str(self.num_total_events) + '\t\t',

        print '%.5f \t%.5f \t%.5f \t%.5f' % \

          (us_5pt, us_3to5pt, ds_5pt, ds_3to5pt)

if __name__ == "__main__":

    eff = PatRecEfficiency()

    args = sys.argv

    args.pop(0)

    eff.run(args)