#!/usr/bin/env python3
#NOTE: For python 2.7 remove single quites around numbers compared to user input ex. choice=='1' to choice==1
import json 
import sys
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d

#creates dynamic attributes with the objects keys recursively
#conversion has 3 cases 1)lists 2)dicts (new object) 3)bool, int, float and str
class AppConfiguration(object):
    def __init__(self, data=None):
        if data is None:
            with open("data_mvtx10.json") as fh:
            #with open("D0toPiKInAcceptanceSignal_Iteration2_3001.cfg_HFMLTriggerInterface.json") as fh:
                data = json.load(fh)
                #print (data) #DEBUG
        else:
            data = dict(data)
            #print (data) #DEBUG

        for key, val in data.items():
            setattr(self, key, self.compute_attr_value(val))

    def compute_attr_value(self, value):
        if type(value) is list:
            return [self.compute_attr_value(x) for x in value]
        elif type(value) is dict:
            return AppConfiguration(value)
        else:
            return value

nEvents = 0
if (len(sys.argv)>1): nEvents = int(sys.argv[1])
instance = AppConfiguration()

layer2color = ['yellow', 'blue', 'green']
chip2color = ['blue', 'orange', 'green', 'red', 'purple', 'brown', 'pink', 'olive', 'cyan']

if (nEvents==0):
    hitLists = [event.RawHit.MVTXHits for event in instance.Events]
else:
    hitLists = [event.RawHit.MVTXHits for event in instance.Events[:nEvents]]

#hits = instance.Events[0].RawHit.MVTXHits
hits = sum(hitLists,[]) #append all hit lits to the empty list

hits_event0 = instance.Events[0].RawHit.MVTXHits

def Layers2D():
    x = [hit.Coordinate[0] for hit in hits]
    y = [hit.Coordinate[1] for hit in hits]
    colors = [layer2color[hit.ID.Layer] for hit in hits]
    #markers = [layer2marker[hit.ID.Layer] for hit in hits]
    plt.scatter(x, y, c=colors)
    plt.title('MVTX Layers in 2D')        
    plt.xlabel('X axis')
    plt.ylabel('Y axis')
    plt.show()

def Layers2D_old():
    for i in range (len(instance.Events[0].RawHit.MVTXHits)):
        if(instance.Events[0].TruthTriggerFlag.Flags.D0toPiKInAcceptance == True):
            if(instance.Events[0].RawHit.MVTXHits[i].ID.Layer == 0):
            #if(instance.Events[0].RawHit.MVTXHits[i].ID.Layer == 0 and instance.Events[0].RawHit.MVTXHits[i].ID.Stave == 0):
                x = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
                y = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
                plt.scatter(x, y, marker='d', color='red')
            if(instance.Events[0].RawHit.MVTXHits[i].ID.Layer == 1):
            #if(instance.Events[0].RawHit.MVTXHits[i].ID.Layer == 1 and instance.Events[0].RawHit.MVTXHits[i].ID.Stave == 0):
                x1 = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
                y1 = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
                plt.scatter(x1, y1, marker='o', color='blue')
            if(instance.Events[0].RawHit.MVTXHits[i].ID.Layer == 2):
            #if(instance.Events[0].RawHit.MVTXHits[i].ID.Layer == 2 and instance.Events[0].RawHit.MVTXHits[i].ID.Stave == 0):
                x2 = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
                y2 = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
                plt.scatter(x2, y2, marker='^', color='green')
    
    plt.title('MVTX Layers in 2D')        
    plt.xlabel('X axis')
    plt.ylabel('Y axis')
    plt.show()

def Layers():
    fig = plt.figure()
    ax = plt.axes(projection="3d")
    x = [hit.Coordinate[0] for hit in hits]
    y = [hit.Coordinate[1] for hit in hits]
    z = [hit.Coordinate[2] for hit in hits]
    colors = [layer2color[hit.ID.Layer] for hit in hits]
    ax.scatter3D(x, y, z, c=colors)
    ax.set_title('MVTX Layers 3D')
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')
    ax.set_zlabel('Z axis')
    ax.view_init(90, 0) 
    plt.show()
    
def Layers_old():
    fig = plt.figure()
    ax = plt.axes(projection="3d")
    
    for i in range(len(instance.Events[0].RawHit.MVTXHits)):
        if(instance.Events[0].RawHit.MVTXHits[i].ID.Layer == 0):
            x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
            y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
            z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
            ax.scatter3D(x_points, y_points, z_points, marker='d', edgecolor='orange', facecolor="pink");
        if(instance.Events[0].RawHit.MVTXHits[i].ID.Layer == 1):
            x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
            y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
            z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
            ax.scatter3D(x_points, y_points, z_points, marker='o', edgecolor='blue', facecolor="red");
        if(instance.Events[0].RawHit.MVTXHits[i].ID.Layer == 2):
            x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
            y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
            z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
            ax.scatter3D(x_points, y_points, z_points, marker='^', edgecolor='green', facecolor="yellow");
    
    ax.set_title('MVTX Layers 3D')
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')
    ax.set_zlabel('Z axis')
    ax.view_init(90, 0) 
    plt.show()

def Chips():
    fig = plt.figure()
    ax = plt.axes(projection="3d")
    x = [hit.Coordinate[0] for hit in hits]
    y = [hit.Coordinate[1] for hit in hits]
    z = [hit.Coordinate[2] for hit in hits]
    colors = [chip2color[hit.ID.Chip] for hit in hits]
    ax.scatter3D(x, y, z, c=colors)
    ax.view_init(0, -132)
    ax.set_title('MVTX Stave ALPIDE chips')
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')
    ax.set_zlabel('Z axis')
    plt.show()
    
def Chips_old():
    fig = plt.figure()
    ax = plt.axes(projection="3d")

    for i in range(len(instance.Events[0].RawHit.MVTXHits)):
        if(instance.Events[0].RawHit.MVTXHits[i].ID.Chip == 0):
            x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
            y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
            z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
            ax.scatter3D(x_points, y_points, z_points, marker='d', edgecolor='pink', facecolor='green');
        if(instance.Events[0].RawHit.MVTXHits[i].ID.Chip == 1):
            x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
            y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
            z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
            ax.scatter3D(x_points, y_points, z_points, marker='o', edgecolor='yellow', facecolor="blue");
        if(instance.Events[0].RawHit.MVTXHits[i].ID.Chip == 2):
            x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
            y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
            z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
            ax.scatter3D(x_points, y_points, z_points, marker='^', edgecolor='green', facecolor="red");
        if(instance.Events[0].RawHit.MVTXHits[i].ID.Chip == 3):
            x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
            y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
            z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
            ax.scatter3D(x_points, y_points, z_points, marker='D', edgecolor='blue', facecolor="red");
        if(instance.Events[0].RawHit.MVTXHits[i].ID.Chip == 4):
            x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
            y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
            z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
            ax.scatter3D(x_points, y_points, z_points, marker='8', edgecolor='pink', facecolor="red");
        if(instance.Events[0].RawHit.MVTXHits[i].ID.Chip == 5):
            x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
            y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
            z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
            ax.scatter3D(x_points, y_points, z_points, marker='s', edgecolor='orange', facecolor="red");
        if(instance.Events[0].RawHit.MVTXHits[i].ID.Chip == 6):
            x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
            y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
            z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
            ax.scatter3D(x_points, y_points, z_points, marker='p', edgecolor='violet', facecolor="red");
        if(instance.Events[0].RawHit.MVTXHits[i].ID.Chip == 7):
            x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
            y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
            z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
            ax.scatter3D(x_points, y_points, z_points, marker='H', edgecolor='purple', facecolor="red");
        if(instance.Events[0].RawHit.MVTXHits[i].ID.Chip == 8):
            x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
            y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
            z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
            ax.scatter3D(x_points, y_points, z_points, marker='+', edgecolor='red', facecolor="red");
        if(instance.Events[0].RawHit.MVTXHits[i].ID.Chip == 9):
            x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
            y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
            z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
            ax.scatter3D(x_points, y_points, z_points, marker='^', edgecolor='red', facecolor="red");
    
    ax.view_init(0, -132)
    ax.set_title('MVTX Stave ALPIDE chips')
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')
    ax.set_zlabel('Z axis')

    plt.show()

def PixelX():
    fig = plt.figure()
    ax = plt.axes(projection="3d")
    hits_chip0 = [hit for hit in hits if hit.ID.Chip==0]
    x = [hit.Coordinate[0] for hit in hits_chip0]
    y = [hit.Coordinate[1] for hit in hits_chip0]
    z = [hit.Coordinate[2] for hit in hits_chip0]
    colors = [hit.ID.Pixel/458 for hit in hits_chip0] #color gradient
    ax.scatter3D(x, y, z, c=colors, cmap='rainbow')

    ax.set_title('MVTX ALPIDE Pixel X')
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')
    ax.set_zlabel('Z axis')
    
    ax.view_init(0, -132)
    plt.show()

def PixelY():
    fig = plt.figure()
    ax = plt.axes(projection="3d")
    hits_layer0 = [hit for hit in hits if hit.ID.Layer==0]
    x = [hit.Coordinate[0] for hit in hits_layer0]
    y = [hit.Coordinate[1] for hit in hits_layer0]
    z = [hit.Coordinate[2] for hit in hits_layer0]
    colors = [hit.ID.Pixel%458 for hit in hits_layer0] #color gradient
    ax.scatter3D(x, y, z, c=colors, cmap='rainbow')

    ax.set_title('MVTX ALPIDE Pixel Y')
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')
    ax.set_zlabel('Z axis')
    
    ax.view_init(90, 0)
    plt.show()

def HalfChips():
    fig = plt.figure()
    ax = plt.axes(projection="3d")
    hits_chip0 = [hit for hit in hits if hit.ID.Chip==0]
    x = [hit.Coordinate[0] for hit in hits_chip0]
    y = [hit.Coordinate[1] for hit in hits_chip0]
    z = [hit.Coordinate[2] for hit in hits_chip0]
    colors = ['blue' if (hit.ID.Pixel/458<500) else 'green' for hit in hits_chip0]
    ax.scatter3D(x, y, z, c=colors)

    ax.set_title('MVTX Stave AlPIDE Chips Half Colors')
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')
    ax.set_zlabel('Z axis')
    
    ax.view_init(0, -132)
    plt.show()

def HalfChips_old():
    fig = plt.figure()
    ax = plt.axes(projection="3d")
    
    for i in range(len(instance.Events[9].RawHit.MVTXHits)):
        x_points = instance.Events[9].RawHit.MVTXHits[i].Coordinate[0]
        y_points = instance.Events[9].RawHit.MVTXHits[i].Coordinate[1]
        z_points = instance.Events[9].RawHit.MVTXHits[i].Coordinate[2]
        if (instance.Events[9].RawHit.MVTXHits[i].ID.Chip == 0):
            if(instance.Events[9].RawHit.MVTXHits[i].ID.Pixel/458 < 500):
                ax.scatter3D(x_points, y_points, z_points, marker='x', cmap='hsv', color='blue');
            else:
                ax.scatter3D(x_points, y_points, z_points, marker='x', cmap='hsv', color='green');
    
    ax.set_title('MVTX Stave AlPIDE Chips Half Colors')
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')
    ax.set_zlabel('Z axis')
    
    plt.show()

def SingleTrack():
    fig = plt.figure()
    ax = plt.axes(projection="3d")
    hitIdsThisTrack = instance.Events[0].TruthHit.TruthTracks[0].HitSequenceInEvent
    hitsThisTrack = [hit for hit in hits_event0 if hit.ID.HitSequenceInEvent in hitIdsThisTrack]
    x = [hit.Coordinate[0] for hit in hitsThisTrack]
    y = [hit.Coordinate[1] for hit in hitsThisTrack]
    z = [hit.Coordinate[2] for hit in hitsThisTrack]
    ax.scatter3D(x, y, z, marker='x', color='red');
    ax.set_title('MVTX Single Track')
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')
    ax.set_zlabel('Z axis')
    plt.show()

def SingleTrack_old():
    #DEBUG
    #print("HitSequenceInEvent: ",instance.Events[0].TruthHit.TruthTracks[0].HitSequenceInEvent) 
    #    for i in range(len(HitSequenceInEvent_1)):
    #        print(HitSequenceInEvent_1[i])
    #    print("ALEX this works",instance.Events[0].RawHit.MVTXHits[0].ID.HitSequenceInEvent)
    
    fig = plt.figure()
    ax = plt.axes(projection="3d")
    
    if(instance.Events[0].TruthHit.TruthTracks[1].TrackSequenceInEvent == 8):
        HitSequenceInEvent_1=instance.Events[0].TruthHit.TruthTracks[0].HitSequenceInEvent
        #for i in range(len(HitSequenceInEvent_1)):  #DEBUG
            #print(HitSequenceInEvent_1[i])		#DEBUG
        for i in range(len(instance.Events[0].RawHit.MVTXHits)):
            for j in range(len(HitSequenceInEvent_1)):
                if(instance.Events[0].RawHit.MVTXHits[i].ID.HitSequenceInEvent==HitSequenceInEvent_1[j]):
                    #print(HitSequenceInEvent_1[j])		                #DEBUG
                    #print(instance.Events[0].RawHit.MVTXHits[i].Coordinate) #DEBUG
                    x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
                    y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
                    z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
                    #ax.scatter3D(x_points, y_points, z_points, marker='o', edgecolor='yellow', facecolor="blue");
                    ax.scatter3D(x_points, y_points, z_points, marker='x', cmap='hsv', color='red');
    
    ax.set_title('MVTX Single Track')
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')
    ax.set_zlabel('Z axis')
    plt.show()

def SingleTrackInLayers():
    fig = plt.figure()
    ax = plt.axes(projection="3d")
    hitIdsThisTrack = instance.Events[0].TruthHit.TruthTracks[0].HitSequenceInEvent
    hitsThisTrack = [hit for hit in hits_event0 if hit.ID.HitSequenceInEvent in hitIdsThisTrack]
    x = [hit.Coordinate[0] for hit in hitsThisTrack]
    y = [hit.Coordinate[1] for hit in hitsThisTrack]
    z = [hit.Coordinate[2] for hit in hitsThisTrack]
    ax.scatter3D(x, y, z, marker='x', color='red');

    x_all = [hit.Coordinate[0] for hit in hits_event0]
    y_all = [hit.Coordinate[1] for hit in hits_event0]
    z_all = [hit.Coordinate[2] for hit in hits_event0]
    colors_all = [layer2color[hit.ID.Layer] for hit in hits_event0]
    ax.scatter3D(x_all, y_all, z_all, c=colors_all, marker='.');

    ax.set_title('MVTX Single Track in Layers')
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')
    ax.set_zlabel('Z axis')
    ax.view_init(90, 0)
    plt.show()

def SingleTrackInLayers_old():
    fig = plt.figure()
    ax = plt.axes(projection="3d")
    #HitSequenceInEvent_1[]
    if(instance.Events[0].TruthHit.TruthTracks[0].TrackSequenceInEvent == 7):
        #print("HitSequenceInEvent: ",instance.Events[0].TruthHit.TruthTracks[0].HitSequenceInEvent) #DEBUG
        HitSequenceInEvent_1=instance.Events[0].TruthHit.TruthTracks[0].HitSequenceInEvent
        #for i in range(len(HitSequenceInEvent_1)): #DEBUG
            #print(HitSequenceInEvent_1[i]) #DEBUG
        #print("ALEX this works",instance.Events[0].RawHit.MVTXHits[0].ID.HitSequenceInEvent) #DEBUG
        for i in range(len(instance.Events[0].RawHit.MVTXHits)):
            if(instance.Events[0].RawHit.MVTXHits[i].ID.Layer == 0):
                x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
                y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
                z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
                ax.scatter3D(x_points, y_points, z_points, marker='.', edgecolor='orange', facecolor="yellow");
            if(instance.Events[0].RawHit.MVTXHits[i].ID.Layer == 1):
                x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
                y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
                z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
                ax.scatter3D(x_points, y_points, z_points, marker='.', edgecolor='blue', facecolor="blue");
            if(instance.Events[0].RawHit.MVTXHits[i].ID.Layer == 2):
                x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
                y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
                z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
                ax.scatter3D(x_points, y_points, z_points, marker='.', edgecolor='green', facecolor="yellow");
            for j in range(len(HitSequenceInEvent_1)):
                if(instance.Events[0].RawHit.MVTXHits[i].ID.HitSequenceInEvent==HitSequenceInEvent_1[j]):
                    #print(instance.Events[0].RawHit.MVTXHits[i].Coordinate) #DEBUG
                    x_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[0]
                    y_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[1]
                    z_points = instance.Events[0].RawHit.MVTXHits[i].Coordinate[2]
                    ax.scatter3D(x_points, y_points, z_points, marker='x', cmap='hsv', color='red');
    
    ax.set_title('MVTX Single Hit in Layers')
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')
    ax.set_zlabel('Z axis')
    ax.view_init(90, 0)
    plt.show()

def print_menu():       ## Your menu design here
    print (25 * "-" , "MVTX VISUALIZE MENU" , 25 * "-")
    print ("1. Single Event Layers in 2D")
    print ("2. Single Event Layers in 3D")
    print ("3. Stave ALPIDE Chips")
    print ("4. Stave ALPIDE Chips half colors")
    print ("4a. ALPIDE Pixel X")
    print ("4b. ALPIDE Pixel Y")
    print ("5. Single Hit")
    print ("6. Single Hit in Layers")
    print ("7. Exit")
    print (71 * "-")
 
loop=True       
#print(sys.argv[0], sys.argv[1]) 
while loop:          
    print_menu()   
    choice = input("Enter your choice [1-7]: ")
    print("Choice:", choice)
    if choice=='1':     
        print ("Showing Single Event Layers in 2D\n")
        Layers2D()
    elif choice=='2':
        print ("Showing Single Event Layers in 3D\n")
        Layers()
    elif choice=='3':
        print ("Showing Stave ALPIDE Chips\n")
        Chips()
    elif choice=='4':
        print ("Showing Stave ALPIDE Chips half colors\n")
        HalfChips()
    elif choice=='4a':
        print ("Showing Pixel X\n")
        PixelX()
    elif choice=='4b':
        print ("Showing Pixel Y\n")
        PixelY()
    elif choice=='5':
        print ("Showing Single Track\n")
        SingleTrack()
    elif choice=='6':
        print ("Showing Single Track in Layers\n")
        SingleTrackInLayers()
    elif choice=='7':
        print ("Have an nice day : - )\n")
        loop=False
    else:
        print("Wrong choice. Try again.\n")