import SmartPhantom as SmrtPhntm

import numpy as np

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

  Script:

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

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

# Set plot flag:
Plot = False
print("          plot = ", Plot)

# 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.readDataNew('11-G4data/waterhits_new.dat')
# PDsData, NumPyArray = SmrtPhntm.readDataNew('waterhits_new.dat')
if Debug:
     print(" Energy deposit data")
     print(" Event, Edep, RealX, RealY, RealZ, Time, PId")

#.. 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

#--------  --------  Initialisation -- end  --------  -------- 

#--------  --------  Loop over rows in waterhits --------  -------- 

for row in NumPyArray:
     if (row[6] == 'proton'):
          if abs(row[2]) <= 20 and abs(row[3]) <= 20:
               i += 1
               if i < 10 and Debug:
                    print(int(row[0]), row[1], row[2], row[3], row[4], row[5], row[6])
               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 and plot energy distributions:
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)

if Plot:
     Dummy = SmrtPhntm.plotEnergySum3Coord(NBinsX, NBinsY, NBinsZ, Edges, EnergySum)

#--------  --------  Loop over rows in waterhits -- end  --------  --------

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

  Calculate power and time:

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

# Calculate power:
TBinWidth  = ( ToFMax / float(NBinsT))   # Bin width -- time range divided by number of bins
Power = EnergySum / TBinWidth            # This power is in MeV/ns
if Debug:
     print(" Scale EnergySum to power using TBinWidth =", TBinWidth, " ns.  Power: \n", Power)


# Create time bins:
TBins = np.zeros(shape=(NBinsT))
LowTBinCtr = 0. + TBinWidth/2.                # Assign time at the centre of the first bin
for i in range(NBinsT):
     TBins[i] = LowTBinCtr + float(i)*TBinWidth
if Debug:
     print(" Filled time bin centres: TBins: \n", TBins)

# Complete, plot power spectra:
print(" Power spectrum created from EnergySum")
if Plot:
     Dummy = SmrtPhntm.plotPower3Coord(NBinsX, NBinsY, NBinsZ, Edges, Power)

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

  Write data to file numpy file:

     --------    --------    --------    -------- """
np.save('power-spectrum-protons_only', Power)
print("Power spectrum saved in power-spectrum-protons_only.npy")


# Done!
