#!/usr/env python

import os

import subprocess

import ROOT

import xboa.Bunch

from xboa.Bunch import Bunch

import xboa.Common as Common

IDEAL_FIELD = None

def run_sim():

    erit_root_dir = os.getenv('ERIT_ROOT')

    sim_path = os.path.join(erit_root_dir, 'maus', 'simulate_erit.py')

    conf_path = os.path.join(erit_root_dir, 'maus', 'tests', 'test_field', 'field_configuration.py')

    print sim_path

    print conf_path

    subproc = subprocess.Popen(['python', sim_path, '-configuration_file', \

                                conf_path])

    subproc.wait()

def plot_output():

    global IDEAL_FIELD

    primary_bunch = Bunch.new_from_read_builtin('maus_primary', 'simulation.out')

    primary_bunch.root_scatter_graph('x', 'z', 'm', 'm')

    virtual_hits = Bunch.new_from_read_builtin('maus_virtual_hit', 'simulation.out')

    virtual_hits.root_scatter_graph('x', 'z', 'm', 'm')

    virtual_hits_by_station = Bunch.new_list_from_read_builtin('maus_virtual_hit', 'simulation.out')

    for bunch in virtual_hits_by_station:

        if abs(bunch.mean(['z'])['z']) < 1e-9:

            (canvas, hist, graph) = bunch.root_scatter_graph('x', 'by', 'm', 'T', ymin=0.5, ymax=1.0)

            graph.Draw('l')

            IDEAL_FIELD = ROOT.TF1("ideal_field","0.727*(2.350/x)**1.92",2.000,2.700)

            IDEAL_FIELD.Draw('lsame')

            canvas.Update()

            for hit in bunch:

                print hit['x'], hit['z'], hit['by']

    #three_d_plot(virtual_hits, 'x', 'z', 'by', 'm', 'm', 'T')

def three_d_plot(bunch, x_axis, y_axis, z_axis, x_units='', y_units='', z_units=''):

    canvas_name = x_axis+'-'+y_axis+'-'+z_axis

    canvas = Common.make_root_canvas(canvas_name)

    x_list, y_list, z_list = [], [], []

    for hit in bunch:

        x_list.append(hit[x_axis]/Common.units[x_units])

        y_list.append(hit[y_axis]/Common.units[y_units])

        z_list.append(hit[z_axis]/Common.units[z_units])

        print x_list[-1], y_list[-1], z_list[-1]

    (hist, graph_3D, graph_2D) = Common.make_root_graph_2d(canvas_name, x_list, x_axis+' ['+x_units+']', y_list, y_axis+' ['+y_units+']', z_list, z_axis+' ['+z_units+']')

    canvas.cd()

    hist.Draw()

    graph_3D.Draw('CONTZ')

    graph_2D.Draw('')

    canvas.Update()

    print 'done'

def main():

    run_sim()

    plot_output()

    raw_input()

if __name__ == "__main__":

    main()