import numpy as np

import read_ttree as read

from ROOT import TFile, TCanvas, gDirectory, TTree, TGraph, TGraphErrors, TMath, TLegend, TMultiGraph, TH1F, TH2F, TColor, gStyle, gROOT, TF2, TEllipse

import matplotlib.pyplot as pyplot

from mpl_toolkits.mplot3d import Axes3D

def main():

    filename = '~/Documents/Work/MICE Data/Data run 7469/run7469_extracted_data__MAUS2pt5_b.root'

    fig, ax = pyplot.subplots(5, 5) # makes a 5x5 set of plots

    fig2 = pyplot.figure(2)

    ax2 = fig2.gca(projection='3d')

    ax2.set_xlabel('x')

    ax2.set_ylabel('y')

    ax2.set_zlabel('z, beam')

    f = TFile(filename, 'read')

    tree = f.Get("T")

    row, column = 0, 0

    for entry in tree:

        if entry.cut_allPassed == 1 and row*column < 16:

            x1 = entry.TKU_s1_x

            x2 = entry.TKU_s2_x

            x3 = entry.TKU_s3_x

            x4 = entry.TKU_s4_x

            x5 = entry.TKU_s5_x

            y1 = entry.TKU_s1_y

            y2 = entry.TKU_s2_y

            y3 = entry.TKU_s3_y

            y4 = entry.TKU_s4_y

            y5 = entry.TKU_s5_y

            z1 = entry.TKU_s1_z

            z2 = entry.TKU_s2_z

            z3 = entry.TKU_s3_z

            z4 = entry.TKU_s4_z

            z5 = entry.TKU_s5_z

            x = [x5, x4, x3, x2, x1]

            y = [y5, y4, y3, y2, y1]

            z = [z5, z4, z3, z2, z1]

            ax[row, column].plot(x1, y1, 'ko')

            ax[row, column].plot(x2, y2, 'bo')

            ax[row, column].plot(x3, y3, 'co')

            ax[row, column].plot(x4, y4, 'go')

            ax[row, column].plot(x5, y5, 'ro')

            print "(row, column) = ", row, column

            if row == 0 and column == 0:

                ax2.plot(x, y, z)

                ax2.scatter(x5, y5, z5, c='r')

                ax2.scatter(x4, y4, z4, c='g')

                ax2.scatter(x3, y3, z3, c='c')

                ax2.scatter(x2, y2, z2, c='b')

                ax2.scatter(x1, y1, z1, c='k')

            if column == 4:

                row += 1

                column = 0

            else:

                column += 1

        if row*column > 16:

            break

    pyplot.show()

main()