#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Module: SmartPhantom:
=====================

Collection of methods used to process the information generated using Geant 4
for the passage of protons through the SmartPhantom.

  Methods:
  --------

  readData: Input  : filename (string): filename with path for data file
            Returns: dataframe : instance of pandas dataframe
                     NumpyArray: numpy array containing columns in dataframe

  defineWaterBox: Input  : height, width, and length (in mm) of water phantom
                  Returns: Vox[]: numpy array containing distance from centre
                           to edge of water phantom in x ([0]), y ([1]), and
                           z ([2]).

  defineVoxel: Input  : Height, width, and length (in mm) of voxel
               Returns: Vox: numpy array with voxel size in x ([0]), y ([1]),
               and z ([2]).

  getNumberOfVoxels: Input: Edges, Vox: as defined above:
                     Returns: number of bins in x, y, and z as three integers.

  getNumberOfTimeSteps: Input: Period (ns) maximum period to be consider
                               Step (ns) time step
                        Returns: number of time bins

  getBin: Input:   NBins, LowEdge, HighEdge, X of bin
          Returns: Bin number

"""

import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

#--------  --------  --------  --------  readData  --------  --------  
def readData(filename):
     """
     Each row of data frame:
          Row[0] = "Event" number--energy deposits resulting from proton "Event"
          Row[1] = Energy deposited (MeV)
          Row[2] = X (mm) of enegy deposit relative to centre of box
          Row[3] = Y (mm) of enegy deposit relative to centre of box
          Row[4] = Z (mm) of enegy deposit relative to centre of box
          Row[5] = T (ns) -- time of flight from t=0 when proton entered.
     """
     
     dataframe = pd.read_csv(filename, delim_whitespace=True, names=['StName','EventN','Edep','RealX','RealY','RealZ','Time'], low_memory=False, skiprows=1)
     print(" readData: Pandas data from created from file: ", filename)
     print("    Number of data entries: ", len(dataframe.index))

     # Convert to numpy array:
     EdepDataFrame  = dataframe[["EventN","Edep","RealX","RealY","RealZ","Time"]]
     EdepNumPyArray = EdepDataFrame.to_numpy()
     print(" readData: Numpy array created with length: ", len(EdepNumPyArray))
     
     return dataframe, EdepNumPyArray
#--------  --------  --------  --------  readData -- end  --------  --------  

#--------  --------  --------  --------  readDataNew  --------  --------  
def readDataNew(filename):
     """
     Each row of data frame:
          Row[0] = StationName
          Row[1] = EventNum
          Row[2] = Total energy deposit (MeV)
          Row[3] = True PosX [mm]
          Row[4] = True PosY [mm]
          Row[5] = True PosZ [mm]
          Row[6] = Global Time [ns]
          Row[7] = Delta time [ns]
          Row[8] = Particle Name

     Each row of numpy array:
          Row[0] = EventN
          Row[1] = Edep
          Row[2] = RealX
          Row[3] = RealY
          Row[4] = RealZ
          Row[5] = Time
          Row[6] = PId
     """
     
     dataframe = pd.read_csv(filename, delim_whitespace=True, names=['StName','EventN','Edep','RealX','RealY','RealZ','Time', 'Step', 'PId'], low_memory=False, skiprows=1)
     print(" readData: Pandas data from created from file: ", filename)
     print("    Number of data entries: ", len(dataframe.index))

     # Convert to numpy array:
     EdepDataFrame  = dataframe[["EventN","Edep","RealX","RealY","RealZ","Time", "PId"]]
     EdepNumPyArray = EdepDataFrame.to_numpy()
     print(" readData: Numpy array created with length: ", len(EdepNumPyArray))
     
     return dataframe, EdepNumPyArray
#--------  --------  --------  --------  readDataNew -- end  --------

#--------  --------  --------  --------  defineWaterBox  --------  --------
def defineWaterBox(Height=200., Width=200., Length=400.):
     """ Input: Height, Width, Length of box in mm.

         Return: Cartesian coordinates of extremes of box as distance
                 from the centre:
                 Edges[2]: Edges[0] -- distance of x edges from centre
                           Edges[1] -- distance of y edges from centre
                           Edges[2] -- distance of z edges from centre
     """

     Edges = np.array([Width/2., Height/2., Length/2.])
     print(" defineWaterBox: edges of water box at +/- x, y, z (mm): ", Edges)

     return Edges
#--------  --------  --------  --------  defineWaterBox  -- end  --------

#--------  --------  --------  --------  defineVoxel  --------  --------
def defineVoxel(Height=2., Width=2., Length=4.):
     """ Input: Height, Width, Length of voxel in mm.

         Return: Numpy array with voxel size in mm
                 Vox[2]: Vox[0] --  x length of voxel
                         Vox[1] --  y length of voxel
                         Vox[2] --  z length of voxel
     """

     Vox = np.array([Width, Height, Length])
     print(" defineVoxel: voxel size (mm): ", Vox)

     return Vox
#--------  --------  --------  --------  defineVoxel  -- end  --------

#--------  --------  --------  --------  getNumberOfBins  --------  --------
def getNumberOfVoxels(Edges, Vox):
    """ Input:  Edges (3 element numpy array) of box (as defined by 
                defineWaterBox) and voxel size (3 element numpy array,
                as defined by defineVox).
        Return: Number of bins in x, y, and z directions (integers)
    """

    NBinx = int(2.*Edges[0]/Vox[0])
    NBiny = int(2.*Edges[1]/Vox[1])
    NBinz = int(2.*Edges[2]/Vox[2])
    print(" GetNumberOfVoxels: x, y, and z:", NBinx, NBiny, NBinz)

    return NBinx, NBiny, NBinz
#--------  --------  --------  --------  getNumberOfBins  -- end  --------

#--------  --------  --------  --------  getNumberOfTimeSteps  --------
def getNumberOfTimeSteps(Period=100., Step=2.):
    """ Input:  
        - Period: "Typical" time period (float) for travel of proton in
                  water phantom until it comes to rest (ns).
        - Step  : Time step (time bin) (ns).
                as defined by defineVox).
        Return: Number of bins in t (integer)
    """

    NBint = int(Period/Step)
    print(" GetNumberOfTimeSteps: ", NBint)

    return NBint
#--------  --------  --------  --------  getNumberOfTimeSteps  -- end  --------
        
#--------  --------  --------  --------  getBin  --------  --------
def getBin(NBins=10, LowEdge=0., HighEdge=10., X=5.):
    """ Return bin number.  
        Note: if Bin < 0, then this is an underflow;
              if Bin > NBins, then this is an overflow.
    """

    BinWidth = (HighEdge - LowEdge) / float(NBins)
    aBin = (X-LowEdge) / BinWidth
    if aBin < 0.:
         Bin = -1
    elif aBin > float(NBins - 1):
         Bin = -2
    else:
         Bin      = int( aBin )

    return Bin
#--------  --------  --------  --------  getBin  -- end  --------  --------

def plotEnergySum(NBinsX, NBinsY, NBinsZ, Edges, EnergySum):
     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 (MeV)')
     plt.title('Summed energy distribution vs z')
     plt.show()

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

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

     return

def plotPower(NBinsX, NBinsY, NBinsZ, Edges, Power):
     EvsZ = np.sum(Power, 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('Power (MeV/ns)')
     plt.title('Power distribution vs z')
     plt.show()

     EvsX = np.sum(Power, axis=(0,2,3))
     Bins = np.zeros(shape=(NBinsX))
     for i in range(NBinsX):
          Bins[i] = -Edges[0] + float(i)*( 2.*Edges[0] / float(NBinsX))
     plt.plot(Bins, EvsX, '-', color='r')
     plt.xlabel('x (mm)')
     plt.ylabel('Power (MeV/ns)')
     plt.title('Power distribution vs x')
     plt.show()

     EvsY = np.sum(Power, axis=(0,1,3))
     Bins = np.zeros(shape=(NBinsY))
     for i in range(NBinsY):
          Bins[i] = -Edges[1] + float(i)*( 2.*Edges[1] / float(NBinsY))
     plt.plot(Bins, EvsY, '-', color='r')
     plt.xlabel('y (mm)')
     plt.ylabel('Power (MeV/ns)')
     plt.title('Power distribution vs y')
     plt.show()

     return

def plotEnergySum3Coord(NBinsX, NBinsY, NBinsZ, Edges, EnergySum):
     fig, (ax1, ax2, ax3) = plt.subplots(1, 3)
     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))
     ax1.plot(Bins, EvsZ, '-', color='r')
     ax1.set_xlabel('z (mm)')
     ax1.set_ylabel('Summed energy (MeV)')
     ax1.set_title('Summed energy distribution vs z')

     EvsX = np.sum(EnergySum, axis=(0,2,3))
     Bins = np.zeros(shape=(NBinsX))
     for i in range(NBinsX):
          Bins[i] = -Edges[0] + float(i)*( 2.*Edges[0] / float(NBinsX))
     ax2.plot(Bins, EvsX, '-', color='r')
     ax2.set_xlabel('x (mm)')
     ax2.set_ylabel('Summed energy (MeV)')
     ax3.set_title('Summed energy distribution vs x')

     EvsY = np.sum(EnergySum, axis=(0,1,3))
     Bins = np.zeros(shape=(NBinsY))
     for i in range(NBinsY):
          Bins[i] = -Edges[1] + float(i)*( 2.*Edges[1] / float(NBinsY))
     ax3.plot(Bins, EvsY, '-', color='r')
     ax3.set_xlabel('y (mm)')
     ax3.set_ylabel('Summed energy (MeV)')
     ax3.set_title('Summed energy distribution vs y')
     plt.show()

     return

def plotPower3Coord(NBinsX, NBinsY, NBinsZ, Edges, Power):
     fig, (ax1, ax2, ax3) = plt.subplots(1, 3)
     EvsZ = np.sum(Power, 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))
     #ax1.plot(Bins, EvsZ, '-', color='r')                                                     
     ax1.bar(Bins, EvsZ, width = 1)
     ax1.set_xlabel('z (mm)')
     ax1.set_ylabel('Power (MeV/ns)')
     ax1.set_title('Power distribution vs z')

     EvsX = np.sum(Power, axis=(0,2,3))
     Bins = np.zeros(shape=(NBinsX))
     for i in range(NBinsX):
          Bins[i] = -Edges[0] + float(i)*( 2.*Edges[0] / float(NBinsX))
     #ax2.plot(Bins, EvsX, '-', color='r')                                                     
     ax2.bar(Bins, EvsX, width = 1)
     ax2.set_xlabel('x (mm)')
     ax2.set_xlim(-20, 20)
     ax2.set_ylabel('Power (MeV/ns)')
     ax2.set_title('Power distribution vs x')

     EvsY = np.sum(Power, axis=(0,1,3))
     Bins = np.zeros(shape=(NBinsY))
     for i in range(NBinsY):
          Bins[i] = -Edges[1] + float(i)*( 2.*Edges[1] / float(NBinsY))
     #ax3.plot(Bins, EvsY, '-', color='r')                                                     
     ax3.bar(Bins, EvsY, width = 1)
     ax3.set_xlabel('y (mm)')
     ax3.set_xlim(-20, 20)
     ax3.set_ylabel('Power (MeV/ns)')
     ax3.set_title('Power distribution vs y')
     plt.show()

     return
