import SmartPhantom as SmrtPhntm

import numpy as np
import matplotlib.pyplot as plt
from prettytable import PrettyTable

"""  --------    --------    --------    -------- 

  Script:

     --------    --------    --------    -------- """

# Set debug flag:
Debug = False
print(" Starting: debug = ", Debug)

# Define water box:
Edges = SmrtPhntm.defineWaterBox(300., 300., 300.)
Vox   = SmrtPhntm.defineVoxel(1., 1., 1.)
print(" Water box width (+/-):", Edges[0], " mm, height (+/-):", Edges[1], " mm, length (+/-):", Edges[2], " mm")
print(" Voxel size (x, y, z):", Vox, " mm")

# Define time of flight max and time step:
ToFMax = 100.
TStep  = 1.
print(" Max time of flight and time step:", ToFMax, TStep)

# Define array to store energy sum:
NBinsX, NBinsY, NBinsZ = SmrtPhntm.getNumberOfVoxels(Edges, Vox)
if Debug:
     print(" NBinsX, NBinsY, NBinsZ:", NBinsX, NBinsY, NBinsZ)
NBinsT = SmrtPhntm.getNumberOfTimeSteps(ToFMax, TStep)
if Debug:
     print(" NBinsT:", NBinsT)
     
EnergySum = np.zeros(shape=(NBinsT, NBinsX, NBinsY, NBinsZ))

# Read data, create instance of pandas data frame and NumPyArray:
PDsData, NumPyArray = SmrtPhntm.readData('11-Data/waterhits.dat')
if Debug:
     print(" Energy deposit data")
     print(" Event, Edep, RealX, RealY, RealZ, Time")

#.. initialise counters:
i                = 0  # counts rows, used for print out only
Evt              = -1 # proton beam particle ("event") counter
ToFMaxChk        = 0. # Find max time of flight
TUnderOrOverFlow = 0
XUnderOrOverFlow = 0
YUnderOrOverFlow = 0
ZUnderOrOverFlow = 0

for row in NumPyArray:
     i += 1
     if i < 10 and Debug:
          print(int(row[0]), row[1], row[2], row[3], row[4], row[5])
     if Evt < int(row[0]):
          Evt +=1
          if Evt < 10:
               print("Another proton: Evt=", Evt)

     ToF = row[5]
     if ToF > ToFMaxChk:
          ToFMaxChk = ToF

     X   = row[2]
     Y   = row[3]
     Z   = row[4]
     
     TBin = SmrtPhntm.getBin(NBinsT, 0., ToFMax, ToF)
     XBin = SmrtPhntm.getBin(NBinsX, -Edges[0], Edges[0], X)
     YBin = SmrtPhntm.getBin(NBinsY, -Edges[1], Edges[1], Y)
     ZBin = SmrtPhntm.getBin(NBinsZ, -Edges[2], Edges[2], Z)
     if i < 10 and Debug:
          print(" XBin, YBin, ZBin, TBin:", XBin, YBin, ZBin, TBin)
     
     #.. Check for overflows:
     if TBin < 0:
          TUnderOrOverFlow += 1
     elif XBin < 0:
          XUnderOrOverFlow += 1
     elif YBin < 0:
          YUnderOrOverFlow += 1
     elif ZBin < 0:
          ZUnderOrOverFlow += 1

     #.. Add energy if not an overflow
     if TBin >= 0 and XBin >= 0 and YBin >= 0 and ZBin >= 0:
          EnergySum[TBin, XBin, YBin, ZBin] += row[1] # This is adding energy; need to scale to get power later
          if i < 10 and Debug:
               print("  EnergySum[TBin, XBin, YBin, ZBin]:", EnergySum[TBin, XBin, YBin, ZBin]) 

#.. At the end, print some summary information
print(" Total number of rows:", i)
print(" Max time of flight:", ToFMaxChk, " ToFMax set:", ToFMax, ToFMaxChk/ToFMax)
print(" Number of time bin under/over flows:", TUnderOrOverFlow)
print(" Number of x bin under/over flows:", XUnderOrOverFlow)
print(" Number of y bin under/over flows:", YUnderOrOverFlow)
print(" Number of z bin under/over flows:", ZUnderOrOverFlow)


#.. Trial: plot Bragg peak: need to add energy in slices of z:
Test = np.zeros(shape=(10, 10, 10, 10))
for TBin in range(10):
     for XBin in range(10):
          for YBin in range(10):
               for ZBin in range(10):
                    if TBin == 4 and XBin == 4 and YBin == 4:
                         Test[TBin, XBin, YBin, ZBin] = float(ZBin)
                    if TBin == 5 and XBin == 5 and YBin == 5:
                         Test[TBin, XBin, YBin, ZBin] = float(ZBin)
EvsZ = np.sum(Test, axis=(0,1,2))
print(EvsZ, len(EvsZ))

"""
EvsZ = np.sum(EnergySum, axis=(0,1,2))
Bins = np.zeros(shape=(NBinsZ))
for i in range(NBinsZ):
     Bins[i] = -Edges[2] + float(i)*( 2.*Edges[2] / float(NBinsZ))
plt.plot(Bins, EvsZ, '-', color='r')
plt.xlabel('z (mm)')
plt.ylabel('Summed energy')
plt.title('Summed energy distribution vs z')
plt.show()
"""

EvsZ = np.sum(EnergySum, axis=(0,2,3))
Bins = np.zeros(shape=(NBinsZ))
for i in range(NBinsZ):
     Bins[i] = -Edges[0] + float(i)*( 2.*Edges[0] / float(NBinsX))
plt.plot(Bins, EvsZ, '-', color='r')
plt.xlabel('z (mm)')
plt.ylabel('Summed energy')
plt.title('Summed energy distribution vs z')
plt.show()
