import os

import subprocess

import json

import operator

import glob

import sys

import numpy

import ROOT

import Configuration

import maus_cpp.globals

import xboa.Common as common

from xboa.Bunch import Bunch

from xboa.Hit import Hit

from xboa.tracking import MAUSTracking

class Tracking:

    def amp(self, ellipse_inv, hit):

        arr = numpy.matrix([hit['x'], hit['px'], hit['y'], hit['py']])

        amplitude = arr*ellipse_inv*arr.T

        return amplitude[0, 0]

    def amp_list(self, bunch, beta, alpha, bz):

        ellipse = Bunch.build_penn_ellipse(1., self.m_mu, beta, alpha, self.p, 0., bz, 1.)

        ellipse_inv = numpy.linalg.inv(ellipse)

        amplitude_list = [self.amp(ellipse_inv, hit) for hit in bunch]

        return amplitude_list

    def amp_list_6d(self, bunch):

        bunch.set_covariance_matrix(True)

        amp_list = [bunch.get_amplitude(bunch, hit, ['x', 'y', 'ct']) for hit in bunch]

        bunch.set_covariance_matrix(False)

        return amp_list

    def scatter_plot_mean_rms(self, x_values, y_values, xmin, xmax, xstep):

        nbins = 9

        bin_low = 0.

        bin_step = 10.

        hist_data = []

        for x, y in zip(x_values, y_values):

            bin = int((x-bin_low)/bin_step)

            while len(hist_data) - 1 < bin:

                hist_data.append([])

            hist_data[bin].append(y)

        hist = ROOT.TH1D("Mean", "", nbins, bin_low, bin_low+(nbins+1)*bin_step)

        for i, data in enumerate(hist_data):

            mean, sigmaSqu = 0., 0.

            if len(data) > 0:

                mean = sum(data)/len(data)

            if len(data) > 1:

                sigmaSqu = sum([x**2 for x in data])/len(data)-mean**2

            hist.SetBinContent(i+1, mean)

            hist.SetBinError(i+1, sigmaSqu**0.5)

        common._hist_persistent.append(hist)

        hist.SetLineWidth(2)

        hist.SetLineColor(8)

        return hist

    def amplitude_scatter_plot(self, bunch_in, bunch_out, predicate):

        """predicate takes arguments (spill_index, spills_in_bunch_out)"""

        spills_out = set([(hit['spill'], hit['event_number']) for hit in bunch_out])

        hits_in = [hit for hit in bunch_in if (hit['spill'], hit['event_number']) in spills_out]

        hits_out = [hit for hit in bunch_out if (hit['spill'], hit['event_number']) in spills_out]

        amp_in = self.amp_list(hits_in, self.beta, 0., self.bz_us)

        amp_out = self.amp_list(hits_out, self.beta, 0., self.bz_ds)

        delta_amp = [(amp_out[i]-amp_in[i])/amp_in[i] for i, a_in in enumerate(amp_in)]

        if len(amp_in) == 0 or len(amp_in) != len(amp_out):

            print "Failed to make graph:", len(amp_in), len(amp_out)

            hist, graph = common.make_root_graph('amplitude x y', [0.], 'x-y amplitude in [mm]', [0.], 'x-y amplitude out [mm]')

        else:

            hist, graph = common.make_root_graph('amplitude x y', amp_in, 'x-y amplitude in [mm]', delta_amp, '(A_{out} - A_{in})/A_{in} [mm]', xmin=0., xmax=200.)

        hist_means = self.scatter_plot_mean_rms(amp_in, delta_amp, 0., 200., 10.)

        graph.SetMarkerStyle(7)

        return hist, hist_means, graph

    def amplitude_p_scatter_plot(self, bunch_in):

        """predicate takes arguments (spill_index, spills_in_bunch_out)"""

        hits_in = [hit for hit in bunch_in]

        amp_in = self.amp_list(hits_in, self.beta, 0., self.bz_us)

        p_in = [hit['p'] for hit in hits_in]

        if len(amp_in) == 0:

            print "Failed to make graph:", len(amp_in)

            hist, graph = common.make_root_graph('amplitude x y', [0.], 'x-y amplitude in [mm]', [0.], 'P [MeV/c]', xmin=0., xmax=200.)

        else:

            hist, graph = common.make_root_graph('amplitude x y', amp_in, 'x-y amplitude in [mm]', p_in, 'P [MeV/c]', xmin=0., xmax=200.)

        hist_means = self.scatter_plot_mean_rms(amp_in, p_in, 0., 200., 10.)

        graph.SetMarkerStyle(7)

        return hist, hist_means, graph

    def amplitude_t_scatter_plot(self, bunch_in, bunch_out, weight_cut):

        """predicate takes arguments (spill_index, spills_in_bunch_out)"""

        spills_out = set([(hit['spill'], hit['event_number']) for hit in bunch_out])

        hits_in_cut = [hit for hit in bunch_in if (hit['spill'], hit['event_number']) in spills_out and hit['weight'] > weight_cut]

        hits_out_cut = [hit for hit in bunch_out if (hit['spill'], hit['event_number']) in spills_out and hit['weight'] > weight_cut]

        amp_in_cut = self.amp_list(hits_in_cut, self.beta, 0., self.bz_us)

        amp_out_cut = self.amp_list(hits_out_cut, self.beta, 0., self.bz_ds)

        delta_t_cut = [h_o['t']-h_i['t'] for h_i, h_o in zip(hits_in_cut, hits_out_cut)]

        if len(amp_in_cut) == 0:

            print "Failed to make graph:", len(amp_in_cut)

            hist, graph = common.make_root_graph('amplitude x y', [0.], 'x-y amplitude in [mm]', [0.], 'dt [ns]', xmin=0., xmax=200.)

        else:

            hist, graph = common.make_root_graph('amplitude x y', amp_in_cut, 'x-y amplitude in [mm]', delta_t_cut, 'dt [ns]', xmin=0., xmax=200.)

        hist_means = self.scatter_plot_mean_rms(amp_in_cut, delta_t_cut, 0., 200., 10.)

        graph.SetMarkerStyle(7)

        return hist, hist_means, graph

    def amplitude_6d_scatter_plot(self, bunch_in, bunch_out):

        """predicate takes arguments (spill_index, spills_in_bunch_out)"""

        amp_in = self.amp_list_6d(bunch_in)

        amp_out = self.amp_list_6d(bunch_out)

        hist_in = common.make_root_histogram('A_{6D} in', amp_in, 'A_{6D} [mm]', 20, xmin=0., xmax=200.)

        hist_out = common.make_root_histogram('A_{6D} out', amp_out, 'A_{6D} [mm]', 20, xmin=0., xmax=200.)

        hist_in.SetName('A_{6D} in')

        hist_out.SetName('A_{6D} out')

        return hist_in, hist_out

    def amplitude_6d_capture_plot(self, bunch_in, bunch_out):

        amp_in = sorted(self.amp_list_6d(bunch_in))

        amp_out = sorted(self.amp_list_6d(bunch_out))

        bins = [i*10. for i in range(11)]+[i*50. for i in range(2, 5)]

        contents_in = []

        contents_out = []

    def amplitude_scatter(self, bunch_in, bunch_out):

        canvas = common.make_root_canvas('amplitude 6d hist')

        hist_in, hist_out = self.amplitude_6d_scatter_plot(bunch_in, bunch_out)

        hist_out.SetLineColor(4)

        hist_out.Draw()

        hist_in.SetLineColor(2)

        hist_in.Draw("SAME")

        common.make_root_legend(canvas, [hist_in, hist_out])

        canvas.Update()

        self.print_canvas(canvas, 'acceptance_6d_hist')

        spills_out = set([hit['spill'] for hit in bunch_out])

        canvas = common.make_root_canvas('amplitude x y scatter')

        hist, hist_means, graph = self.amplitude_scatter_plot(bunch_in, bunch_out, lambda x: True)

        hist.Draw()

        graph.Draw('p')

        hist_means.Draw("same")

        canvas.Update()

        self.print_canvas(canvas, 'acceptance')

        canvas = common.make_root_canvas('amplitude x y p scatter')

        hist, hist_means, graph = self.amplitude_p_scatter_plot(bunch_in)

        hist.Draw()

        graph.Draw('p')

        hist_means.Draw("same")

        canvas.Update()

        self.print_canvas(canvas, 'acceptance_p')

        canvas = common.make_root_canvas('amplitude x y t scatter')

        hist, hist_means, graph_cut = self.amplitude_t_scatter_plot(bunch_in, bunch_out, 0.1)

        hist.Draw()

        dummy_1, dummy_2, graph_all = self.amplitude_t_scatter_plot(bunch_in, bunch_out, -0.1)

        graph_all.SetMarkerColor(2)

        graph_all.Draw('p')

        graph_cut.Draw('p')

        hist_means.Draw("same")

        canvas.Update()

        self.print_canvas(canvas, 'acceptance_t')

    def two_d_plots(self, bunch_list, colour, canvas_list = None):

        n_graphs = 4

        if canvas_list == None:

            canvas_list = [None]*n_graphs

        graph_list = [None]*n_graphs

        canvas_list[0], hist, graph_list[0] = bunch_list[0].root_scatter_graph("t", "energy", "ns", "MeV", include_weightless=False, canvas=canvas_list[0])

        canvas_list[1], hist, graph_list[1] = bunch_list[-1].root_scatter_graph("t", "energy", "ns", "MeV", include_weightless=False, canvas=canvas_list[1])

        canvas_list[2], hist, graph_list[2] = bunch_list[0].root_scatter_graph("r", "pt", "mm", "MeV/c", include_weightless=False, canvas=canvas_list[2])

        canvas_list[3], hist, graph_list[3] = bunch_list[-1].root_scatter_graph("r", "pt", "mm", "MeV/c", include_weightless=False, canvas=canvas_list[3])

        for i in range(n_graphs):

            canvas_list[i].cd()

            graph_list[i].SetMarkerStyle(6)

            graph_list[i].SetMarkerColor(colour)

            graph_list[i].Draw("p")

            canvas_list[i].Update()

        self.print_canvas(canvas_list[0], "time_energy_at_start")

        self.print_canvas(canvas_list[1], "time_energy_at_end")

        self.print_canvas(canvas_list[2], "r_pt_at_start")

        self.print_canvas(canvas_list[3], "r_pt_at_end")

        return canvas_list

    def cut(self, bunch_list):

        print "    weights in    "       

        print "      no cut:", bunch_list[0].bunch_weight(), bunch_list[-1].bunch_weight()

        for bunch in bunch_list:

            bunch.transmission_cut(bunch_list[-1], global_cut=True)

            bunch.transmission_cut(bunch_list[-1], global_cut=True, test_variable = ['spill', 'event_number', 'particle_number', 'pid'])

            bunch.cut({'pid':-13}, operator.ne, global_cut=True)

        print '      transmission cut:', bunch_list[0].bunch_weight(), bunch_list[-1].bunch_weight(), "**"

        bunches = [bunch_list[0], bunch_list[-1]]

        cut_stream = ['upstream', 'downstream']

        for i in [0, 1]:

            print '   ', cut_stream[i]

            print '                           ', 'cut  ', '       u/s', '       d/s'

            bunches[i].cut({'r':self.r_cut[i]}, operator.gt, global_cut=True)

            print '      cut r:               ', self.r_cut[i], str(bunches[0].bunch_weight()).rjust(10), str(bunches[1].bunch_weight()).rjust(10)

            bunches[i].cut({'amplitude x y':self.amplitude_trans_cut[i]}, operator.gt, global_cut=True)

            print '      cut amp_trans:       ', str(self.amplitude_trans_cut[i]).rjust(5), str(bunches[0].bunch_weight()).rjust(10), str(bunches[1].bunch_weight()).rjust(10)

            bunches[i].cut({'amplitude x y':self.amplitude_trans_cut[i]}, operator.gt, global_cut=True)

            print '      cut amp_trans:       ', str(self.amplitude_trans_cut[i]).rjust(5), str(bunches[0].bunch_weight()).rjust(10), str(bunches[1].bunch_weight()).rjust(10)

            if self.do_long_cut:

                bunches[i].cut({'amplitude ct':self.amplitude_ct_cut[i]}, operator.gt, global_cut=True)

                print '      cut amp_long:        ', str(self.amplitude_ct_cut[i]).rjust(5), str(bunches[0].bunch_weight()).rjust(10), str(bunches[1].bunch_weight()).rjust(10)

                bunches[i].cut({'amplitude ct':self.amplitude_ct_cut[i]}, operator.gt, global_cut=True)

                print '      cut amp_long:        ', str(self.amplitude_ct_cut[i]).rjust(5), str(bunches[0].bunch_weight()).rjust(10), str(bunches[1].bunch_weight()).rjust(10)

                mean_p = bunches[i].mean(['p'])['p']

                bunches[i].cut({'p':mean_p+self.p_cut[i][1]}, operator.gt, global_cut=True)

                bunches[i].cut({'p':mean_p-self.p_cut[i][0]}, operator.lt, global_cut=True)

                print '      cut',  round(mean_p-self.p_cut[i][0], 1), '< p <', round(mean_p+self.p_cut[i][1], 1), '     ', str(bunches[0].bunch_weight()).rjust(10), str(bunches[1].bunch_weight()).rjust(10)

    def z_plot(self, bunch_list):

        print "Start MOMENTS"

        print bunch_list[0].moment(['t', 't'])

        print bunch_list[0].moment(['pz', 't'])

        print bunch_list[0].moment(['pz', 'pz'])

        """

        for item in ['t', 'pz', 'energy']:

            canvas = common.make_root_canvas('sigma_'+item)

            z_list = [bunch.mean(['z'])['z'] for bunch in bunch_list]

            item_list = [bunch.moment([item, item])**0.5 for bunch in bunch_list]

            hist, graph = common.make_root_graph('sigma_'+item, z_list, 'mean z [mm]', item_list, '#sigma('+item+')')

            hist.Draw()

            graph.Draw('l')

            self.print_canvas(canvas, 's_'+item+'_vs_z')

        """

        canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'mean', ['p'], 'mm', 'MeV/c')

        self.print_canvas(canvas, 'momentum_vs_z')

        canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'beta', ['x', 'y'], 'mm', 'mm', ymin=0., ymax=2000.)

        hist.GetXaxis().SetTitle("z [mm]")

        hist.GetYaxis().SetTitle("Transverse #beta [mm]")

        self.set_style(graph)

        self.get_verticals(canvas)

        self.print_canvas(canvas, 'beta_trans_vs_z')

        canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'emittance', ['x', 'y'], 'mm', 'mm')

        hist.GetXaxis().SetTitle("z [mm]")

        hist.GetYaxis().SetTitle("Transverse emittance [mm]")

        self.set_style(graph)

        self.get_verticals(canvas)

        self.print_canvas(canvas, 'emittance_trans_vs_z')

        canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'beta', ['t'], 'mm', 'mm')

        self.print_canvas(canvas, 'beta_long_vs_z')

        canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'emittance', ['ct'], 'mm', 'mm')

        self.print_canvas(canvas, 'emittance_long_vs_z')

        canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'emittance', ['ct', 'x', 'y'], 'mm', 'mm')

        self.print_canvas(canvas, 'emittance_6d_vs_z')

        ref_bunch_list = [Bunch.new_from_hits([bunch[0]]) for bunch in bunch_list]

        canvas, hist, graph = Bunch.root_graph(ref_bunch_list, 'mean', ['z'], 'mean', ['bz'], 'mm', 'T')

        hist.GetXaxis().SetTitle("z [mm]")

        hist.GetYaxis().SetTitle("B_{z} [T]")

        self.set_style(graph)

        self.get_verticals(canvas)

        self.print_canvas(canvas, 'bz_vs_z')

    def grok_filename(self, filename):

        base = os.path.basename(filename)[12:-5]

        emittance = base.split('_')[0]

        emittance = float(emittance.split('=')[1])

        a_dir = os.path.dirname(filename).split('/')

        index = a_dir[-1].split('_')[-1]

        base += '_'+index

        a_dir = "/".join(a_dir[:-1])

        print "    directory", a_dir, "name", base

        self.current = {

            "dir":a_dir,

            "name":base,

            "full":filename,

            "emittance":emittance,

            "index":index

        }

        return self.current

    def print_canvas(self, canvas, name):

        for format in ["png", "root"]:

            canvas.Print(self.current["dir"]+"/"+name+"_"+self.current["name"]+"."+format)

    def single_analysis(self, filename):

        print "Loading data", filename

        metadata = self.grok_filename(filename)

        bunch_list = Bunch.new_list_from_read_builtin('maus_root_virtual_hit', filename)

        len_cut = max([len(bunch) for bunch in bunch_list])/20

        print "    Discarding bunches with fewer than", len_cut, "particles"

        bunch_list = [bunch for bunch in bunch_list if len(bunch) > len_cut]

        if self.z_max != None:

            bunch_list = [bunch for bunch in bunch_list if bunch[0]['z'] < self.z_max]

        print "    Printing transmission (before cut)"

        canvas, hist, graph = Bunch.root_graph(bunch_list, 'mean', ['z'], 'bunch_weight', [], 'mm', '')

        self.print_canvas(canvas, 'weight_vs_z')

        print "    Making cuts"

        #canvas_list = self.two_d_plots(bunch_list, 2)

        self.cut(bunch_list)

        #self.amplitude_scatter(bunch_list[0], bunch_list[-1])

        print "    Making 2d plots"

        #self.two_d_plots(bunch_list, 1, canvas_list)

        print "    Making z plots"

        self.z_plot(bunch_list)

        print "    Making acceptance plots"

        self.data.append((metadata, bunch_list))

        Bunch.clear_global_weights()

    def multi_analysis(self):

        multi_data = open(self.data[0][0]["dir"]+"/multi_analysis_"+self.current["index"]+".json", "w")

        delta_transmission = []

        delta_emittance_trans = []

        nominal_emittance_trans = []

        actual_emittance_trans = []

        actual_emittance_long = []

        delta_emittance_long = []

        actual_emittance_6d = []

        delta_emittance_6d = []

        sigma_energy = []

        for item in self.data:

            self.cut(item[1])

            try:

                delta_trans = float(len(item[1][-1]))/float(len(item[1][0]))

                emit_in = item[1][0].get_emittance(['x', 'y'])

                emit_out = item[1][-1].get_emittance(['x', 'y'])

                delta_emit_t = emit_out/emit_in-1.

                s_energy = (item[1][0].moment(['energy', 'energy']))**0.5

                emit_long_in = item[1][0].get_emittance(['ct'])

                emit_long_out = item[1][-1].get_emittance(['ct'])

                delta_emit_l = emit_long_out/emit_long_in-1.

                emit_6d_in = item[1][0].get_emittance(['x', 'y', 'ct'])

                emit_6d_out = item[1][-1].get_emittance(['x', 'y', 'ct'])

                delta_emit_6d = emit_6d_out/emit_6d_in-1.

                sigma_energy.append(s_energy)

                actual_emittance_long.append(emit_long_in)

                delta_emittance_long.append(delta_emit_l)

                actual_emittance_trans.append(emit_in)

                delta_emittance_trans.append(delta_emit_t)

                actual_emittance_6d.append(emit_6d_in)

                delta_emittance_6d.append(delta_emit_6d)

                delta_transmission.append(delta_trans)

                nominal_emittance_trans.append(item[0]['emittance'])

            except Exception:

                sys.excepthook(*sys.exc_info())

            Bunch.clear_global_weights()

        json_doc = {

            "delta_transmission":delta_transmission,

            "nominal_emittance_trans":nominal_emittance_trans,

            "delta_emittance_trans":delta_emittance_trans,

            "actual_emittance_trans":actual_emittance_trans,

            "delta_emittance_long":delta_emittance_long,

            "actual_emittance_long":actual_emittance_long,

            "delta_emittance_6d":delta_emittance_6d,

            "actual_emittance_6d":actual_emittance_6d,

            "sigma_energy":sigma_energy,

        }

        print >> multi_data, json.dumps(json_doc)

    def set_style(self, graph, is_alternative = None):

        if is_alternative == None:

            is_alternative = self.current["full"].find("alternative") > -1

        if is_alternative:

            graph.SetLineColor(ROOT.kBlue-4)

            graph.SetLineStyle(2)

            graph.SetLineWidth(1)

            graph.SetMarkerStyle(ROOT.kFullSquare)

            graph.SetMarkerSize(1)

        else:

            graph.SetLineColor(1)

            graph.SetLineStyle(1)

            graph.SetLineWidth(1)

            graph.SetMarkerStyle(ROOT.kFullCircle)

            graph.SetMarkerSize(1)

    def get_verticals(self, canvas, is_alternative = None):            

        print "get_verticals", self.current["full"]

        if is_alternative == None:

            is_alternative = self.current["full"].find("alternative") > -1

        canvas.cd()

        if is_alternative:

            afc_centre = [-1090., +1090.]

            absorber_centre = [0.]

        else:

            afc_centre = [-862.75, +862.75]

            absorber_centre = [0.]

        canvas.cd()

        for z in afc_centre:

            dummy, graph = common.make_root_graph("afc", [z]*2, "", [-1.e6, 1.e6], "")

            graph.SetLineColor(ROOT.kRed)

            graph.SetLineStyle(3)

            graph.SetLineWidth(1)

            graph.Draw('l')

        for z in absorber_centre:

            dummy, graph = common.make_root_graph("abs", [z]*2, "", [-1.e6, 1.e6], "")

            graph.SetLineColor(ROOT.kRed+2)

            graph.SetLineStyle(3)

            graph.SetLineWidth(1)

            graph.Draw('l')

        canvas.Update()

    def __init__(self, momentum, amplitude):

        self.data = []

        self.current = {}

        self.p = momentum

        self.m_mu = common.pdg_pid_to_mass[13]

        self.bz_us = 4.e-3

        self.bz_ds = -4.e-3

        self.beta = self.p/self.bz_us*2./common.constants['c_light']

        self.z_max = 4050.1

        self.scraped = {'__any__':[]}

        self.r_cut = [150., 150.]

        self.amplitude_trans_cut = [72., 72.]

        self.do_long_cut = True

        self.amplitude_ct_cut = [150., 150.]

        self.p_cut = [[50., 50.], [50., 50.]]

        file_list = glob.glob("output/output_reference/output*_13/*.root")

        batch = len(file_list) > 2

        ROOT.gROOT.SetBatch(batch)

        print file_list

        for filename in sorted(file_list):

            try:

                self.single_analysis(filename)

            except Exception:

                sys.excepthook(*sys.exc_info())

        if batch:

            self.multi_analysis()

        print "Done\n\n"

        if not batch:

            raw_input()

def main():

    Tracking(200., 1.)

if __name__ == "__main__":

    main()