##^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^Lance Simms, Stanford University 2009
##RadialIonMotion_1d.py
##
##PURPOSE:
##    To take a set of parameters (Temp, Width, Bias voltages, etc.) and 
##    evolve 2 Gaussians using a solution to the radial diffusion and 
##    drift equations equation
##
##KEYWORDS:
##  Mode: int
##    0 - use an explicit method
##    1 - use a Newton Iteration methodd
##    2 - use an implicit method
##  TvFlag: int
##    0 - Don't plot
##    1 - Plot according to TvPPT
##  NSteps: int
##    The number of steps to execute
##  SaveFig: int
##    0 - Don't save any figures
##    1 - Save the figures at the end
##  Simple: int
##    0 - Approximate diode entirely with lightly doped n material
##    1 - Use p+ n n+ regions with abrupt junctions
##  UseEField: int
##    0 - Don't calculate J_drift
##    1 - Include J_drift in the FD
##  UseDiffusion: int
##    0 - Don't use diffusion
##    1 - Use diffusion in FD
##  CalcDataCube: int
##    0 - Don't fill up a datacube with signal values
##    1 - fill up a datacube with signal values
##  IncDarkCurrent: int
##    0 - Don't include dark current in the datacube
##    1 - Include dark current in the datacube
##  TP: int
##    The total number of points in the radial grid
##  RMax: float
##    The maximum value for the radial grid
##  dt: float
##    The time step used in the explicit method
##  CompProfiles: datacube
##    A set of radial profiles to compare the diffusion simulation to 
##  CompRadii: float arr
##    A set of radii to use in order to plot the Profile
##  PlotMode: int
##    0 - Plot the simulations on top of real data to see if it aligns
##    1 - Only plot simulations to speed things up and test things out
##
###################################################################
import matplotlib.pylab as mplot
from numpy import *
from pylab import *
from Return_Persistence_Quantities import *
from LU_decomposition import *
from RadGrid_1d import * 
import pdb

execfile('/afs/slac/u/ki/lances/python/python_HxRG/HxRG_Setup.py')
PlotDir = '/nfs/slac/g/ki/ki04/lances/PNSimulations/'

def RadialDiffusion_1d(Mode = 0, V_bias=0, NSteps = 20, TP=2000, RMax=20,\
                      dt=0.1, t_final=40, \
                      TvFlag = 1, TvPPT = 0, Verbosity = 0, \
                      CompProfiles=0, CompRadii=0, CompReadNum = 0, \
                      UseEField=1, UseDiffusion=1, CalcDataCube = 1, \
                      SaveFig = 0, SaveMultFigs = 0, IncDarkCurrent = 0,\
                      PlotMode = 1):

  ##Determine whether or not to plot Comparison Data from a set of real persistence curves
  if size(CompProfiles) == 0:
    PlotComp = 0
  else: 
    PlotComp = 1
    if CompReadNum == 0: CompReadNum = CompProfiles.shape[2]-1

  ##MODE STRING AND VARIABLES ##################################################
  if Mode == 0:
    AlphaInt   = 1.
    ModeStr    = 'Explicit'
    ImpBool    = 0
  elif Mode == 1:
    AlphaInt   = 1.
    ModeStr    = 'Newton'
    ImpBool    = 0
  elif Mode == 1:
    ModeStr = 'Implicit'
    AlphaInt   = 1.
    ImpBool    = 1
  elif Mode == 2:
    ModeStr = 'CrankNicholson'
    AlphaInt   = .5 
    ImpBool    = 1

  ##############################################################################
  ##CONSTANTS:
  e_0    = 8.85419e-12          # F * m^-1 - 8.85419 e-14 F*cm^-1
  e_si   = 11.8                 # Relative Permittivity for Si	
  e_sio2 = 2.5
  q      = double(1.6022e-19)   # Coulombs per electron
  q      = 1
  k      = double(1.38066e-23)  # Joules per Kelvin
  T      = double(150)           # Kelvin
  Theta  = double(q/(k*T))      # Factor relating energies

  ##DETECTOR CHARACTERISTICS
  ############################################################################
  ##Parameters for the PIN didoe - 100 um thick
  n_i         = 1.e16
  n_bulk      = 1.3e18
  N_d_max     = 1.e24
  N_a_max     = 1.e24
  PixRatio_um = 1./18e-6               # Pixels per m
  PixRatio_cm = 1./18e-4               # Pixels per cm

  ##Effective Masses
  m_e         = 9.1e-31                # Kg
  m_ee        = 0.26                   # effective mass of electrons/mass of electron
  m_eh        = 0.386                  # effective mass of holes/mass of electron

  ##Mobilities
  mu_n        = double(1200.e-4)       # m^2/V s 
  mu_p        = double(400.e-4)        # m^2/V s
  mu_h        = 0.0001                    # Unitless Mobilities
  mu_e        = 0.0001                    # Unitless Mobilities

  ##Lifetimes
  tau_n       = 1.e-5                  # Seconds
  tau_p       = 1.e-5                  # Seconds 
  tau_rise    = 500.0                   # Time constant for rise of persistence curve

  ##Diffusion constants
  D_n         = k*T*mu_n/q             # For Holes
  D_p         = k*T*mu_p/q             # For Electrons
  D_n         = 0.085                  # Unitless diffusion constants for ions
  D_p         = 0.050                  # Unitless diffusion constants for ions
  #D_n         = 0.85                  # Unitless diffusion constants for ions
  #D_p         = 0.50                  # Unitless diffusion constants for ions
  D_p_phys    = D_p/(PixRatio_cm**2)   # Diffusion Coefficient in cm^2/s
  D_n_phys    = D_n/(PixRatio_cm**2)   # Diffusion Coefficient in cm^2/s

  ##Generation-Recombination terms -- Kurata p. 14
  Alpha_no = 3.80e8         # m^-1
  Alpha_po = 2.25e9         # m^-1
  E_no     = 1.75e8         # V/m
  E_po     = 3.26e8         # V/m

  ############################################################################
  ##Parameters for the PIN didoe - 100 um thick
  n_i     = 1.e16
  n_bulk  = 1.3e18
  N_d_max = 1.e24
  N_a_max = 1.e24

  ##Assume Phi at x=-inf is 0, calculate Phi_max according to 
  ##Ashcroft Mermin p. 575
  E_g = 1.12*q               # Joules = eV*q 
  N_c  = 2.5*(m_ee)**1.5*(T/300)**1.5*10.**18
  P_v  = 2.5*(m_eh)**1.5*(T/300)**1.5*10.**18
 
  #Initial Grid Setup and Initial Radial Profile 
  ############################################################################
  GR       = RadGrid_1d(TP, RMax=RMax) #The class that represents the grid
  PixRatio = 18e-6/1.0                 #The ratio pixels/distance & simulation time/real time
  T        = 100
  N_e      = 3200
  N_h      = 10100
  D_h      = 30.0
  D_e      = 38.0
  H_mu     = 0
  E_mu     = 6.8

  #Time Values and CFL number
  #############################################################################
  t_init      = 0.1                           #The initial time to use in the calculatoin
  dt          = 0.0001                        #The time step in simulation time
  t_samp_plot = 1                             #The time interval to stop and plot quantities
  t_samp      = 1                             #The times at which the curves are sampled  
  NumTimes    = ceil((t_final-t_init)/t_samp)

  #Radial profile is sum of neg and pos gaussians
  t           = 0.1 
  #Hole        = N_h/(sqrt(4*pi*D_h*t))*exp(-(GR.r-H_mu)**2/(4*D_h*t))
  #Elec        = N_e/(sqrt(4*pi*D_e*t))*exp(-(GR.r-E_mu)**2/(4*D_e*t))
  Hole        = 1800*exp(-GR.r**2/50)
  HoleOrig    = Hole.copy()
  Elec        = GR.r*250*exp(-(GR.r-4.0)**2/20)
  Elec[:]     = 0
  ADU         = Hole-Elec
  EFld        = zeros(GR.r.size, dtype=double)   
  d_Hole_Drft = zeros(GR.r.size, dtype=double)
  d_Hole_Diff = zeros(GR.r.size, dtype=double)
  d_Elec_Drft = zeros(GR.r.size, dtype=double)
  d_Elec_Diff = zeros(GR.r.size, dtype=double)

  #Allow for a datacube that holds the values vs. time
  if CalcDataCube == 1:
    DataCube  = zeros([GR.r.size, NumTimes], dtype=double)
    RealTimes = zeros(NumTimes, dtype=double)

  #Genral Comparison Plot
  H     = 500
  H_sig = 3
  E     = -300
  E_sig = 6
  E_mu  = 6
  TestPlot = H*exp(-GR.r**2/(2*H_sig**2))+E*exp(-(GR.r-E_mu)**2/(2*E_sig**2))

  colors = ['r','b','g','m']

  if TvFlag == 1:
    if PlotMode == 0:
      mplot.hold(True)
      mplot.figure(0)
      mplot.clf()
      mplot.figure(1)
      mplot.clf()
      mplot.plot(GR.r, ADU, 'co')
      if PlotComp == 1:
        colors = ['r','b','g','m']
        for i in arange(4):
          mplot.plot(CompRadii[:,:], CompProfiles[:,:,i], colors[i]+'o')
      mplot.figure(2)
      mplot.clf()
    elif PlotMode == 1:
      mplot.hold(True)
      mplot.figure(0)
      mplot.clf()
      mplot.figure(1)
      mplot.clf()

  ##BEGIN NUMERICAL SIMULATION +=======================================================
  t_samp_ind = 0.0
  for t in frange(t_init, t_final, dt):
  
    #Solution to FD equation for diffusion yields:
    #q*D_p*{(2/r)*(dp/dr)+d^2p/dr^2}

    #Calculate Electric Field: Integrate Poisson from r=inf to r=0
    if UseEField == 1:
      for i, r_ind, r in zip(arange(GR.r.size), (arange(GR.r.size))[::-1], GR.r[::-1]):
        if i == 0: 
          EFld[r_ind] = 0
        else:
          #POSSIBLE EXPRESSIONS FOR ELECTRIC FIELD INTEGRATION
          EFld[r_ind] = GR.dr[r_ind-1]*(-(1/GR.r[r_ind+1])*EFld[r_ind+1] + \
                                       (1/(GR.dr[r_ind-1]))*EFld[r_ind+1] + \
                                       (Hole[r_ind]-Elec[r_ind])/e_sio2)


      #With EField obtained, calculate drift current
      #Use Backward Difference for Holes, Forward Difference for Electrons
      EFld[0] = 0 
      J_Drft_Hole = q* mu_h * EFld * Hole
      J_Drft_Elec = q* mu_e * EFld * Elec  
      d_Hole_Drft[1:GR.I-1] = ((1/GR.r[1:GR.I-1])*J_Drft_Hole[1:GR.I-1] + \
                             (J_Drft_Hole[1:GR.I-1]-J_Drft_Hole[0:GR.I-2])/(GR.dr[1:GR.I-1])) 
      d_Elec_Drft[1:GR.I-1] = ((1/GR.r[1:GR.I-1])*J_Drft_Elec[1:GR.I-1] + \
                             (J_Drft_Elec[2:GR.I]-J_Drft_Elec[1:GR.I-1])/(GR.dr[1:GR.I-1]))
    
      #Boundary terms
      d_Hole_Drft[0]        =  (J_Drft_Hole[1]-J_Drft_Hole[0])/GR.dr[0] #+\
                               #(1/GR.r[0])*J_Drft_Hole[0] 
      d_Elec_Drft[0]        =  (J_Drft_Elec[1]-J_Drft_Elec[0])/GR.dr[0] #+\
                               #(1/GR.r[0])*J_Drft_Elec[0] 

    #Center Values Mobility Terms
    if UseDiffusion == 1:
      d_Hole_Diff[1:GR.I-1] = -D_p*(\
                  (1/GR.r[1:GR.I-1])*(Hole[2:GR.I]-Hole[0:GR.I-2])/(2*GR.dr[1:GR.I-1])+\
                  (Hole[2:GR.I]-2*Hole[1:GR.I-1]+Hole[0:GR.I-2])/(GR.dr[1:GR.I-1]**2))
      d_Elec_Diff[1:GR.I-1] = D_n*(\
                  (1/GR.r[1:GR.I-1])*(Elec[2:GR.I]-Elec[0:GR.I-2])/(2*GR.dr[1:GR.I-1])+\
                  (Elec[2:GR.I]-2*Elec[1:GR.I-1]+Elec[0:GR.I-2])/(GR.dr[1:GR.I-1]**2))

      #Boundaries dp/dr = 0 at center
      d_Hole_Diff[0] = -D_p*(Hole[1]-2*Hole[0]+Hole[1])/(GR.dr[0]**2)
      d_Elec_Diff[0] =  D_n*(Elec[1]-2*Elec[0]+Elec[1])/(GR.dr[0]**2)

    #Update the Time Step
    Hole = Hole-(d_Hole_Diff+d_Hole_Drft)*dt
    Elec = Elec+(d_Elec_Diff+d_Elec_Drft)*dt
    
    #Get the zero crossings and other quantities used to describe the donuts 
    Sig = Hole-Elec + 1 #Add 1 so that we get two zero crossing values

    #Allow for calculation of the actual Signal vs Time based on exponential rise
    if CalcDataCube == 1 and t > 0.2 and mod(t_samp_ind*dt, t_samp) == 0:
      print 'Real Time: ' + str(t)
      t_ind = int(t)/t_samp
      DataCube[:,t_ind] = DataCube[:, t_ind-1]+(t_samp/tau_rise)*(Sig[:]-DataCube[:, t_ind-1])
      RealTimes[t_ind]  = t
      if IncDarkCurrent == 1:
        DataCube[:,t_ind]=DataCube[:,t_ind]+DarkCurRate*t

    if t > 0.02 and mod(t_samp_ind, 10000) == 0 : # and mod(t_samp_ind*dt, t_samp_plot) == 0: 
      print 'Real Time:   ' + str(t)
      MinInt, MaxInt, MinRad, MaxRad, FirstZRad, SecondZRad, SumCore,\
         SumHalo, SumOut, SumInt, FitC0, FitC1, FitC2, FitC3, FitC4,\
         MinFitInt, MaxFitInt, MinFitRad, MaxFitRad, FirstFitZRad, SecondFitZRad = \
         Return_Persistence_Quantities(Sig, 0, 0, 1, FromRadProf=1, Radii=GR.r)

      if TvFlag == 1:
        #Plot the radial profiles to see if they match up
        if PlotMode == 0:
          mplot.figure(0)
          mplot.clf()
          mplot.plot(GR.r, Hole , 'bo')
          mplot.plot(GR.r, Elec, 'ro')
          mplot.figure(1)
          mplot.clf()
          mplot.plot(CompRadii[:,:], CompProfiles[:,:,1], colors[1]+'o')
          mplot.plot(GR.r, Hole-Elec, 'go')
          mplot.figure(3)
          mplot.clf()
          mplot.plot(CompRadii[:,:], CompProfiles[:,:,2], colors[1]+'o')
          mplot.plot(GR.r, Hole-Elec, 'co')
          mplot.figure(4)
          mplot.clf()
          mplot.plot(CompRadii[:,:], CompProfiles[:,:,3], colors[1]+'o')
          mplot.plot(GR.r, Hole-Elec, 'ro')
          mplot.figure(2)
          mplot.plot([t],[MinRad], 'ro')
          mplot.plot([t],[FirstZRad], 'go')
          mplot.plot([t],[SecondZRad], 'bo')
        elif PlotMode == 1:
          mplot.figure(0)
          mplot.clf()
          mplot.plot(GR.r, Hole, 'ro')
          mplot.figure(1)
          mplot.clf()
          mplot.plot(GR.r, HoleOrig-Hole, 'bo')
      if Verbosity == 1:
        print 'MinRad: ' + str(MinRad)
        print 'FirstZRad: ' + str(FirstZRad)
        print 'SecondZRad: ' + str(SecondZRad)
    
    #Increment the sampling time index
    t_samp_ind = t_samp_ind+1

  pdb.set_trace() 
  return ADU
