#############################################################################################
# Simplified random walk model for in-canopy transfer modelling of trace gases
# 2018-10-17 - simple version as support for Master 2 Clues course
# Copyright B. Loubet Benjamin.Loubet@inra.fr
# this script hasis just an example and should not be used for research or evaluation purpose
#############################################################################################


library(doParallel)
library(numDeriv)
#library("dplyr")

source('DiffusivityFunction.R')

RandomWalk_model_2018 = function(ustar = c(0.1,0.1,0.5),    #m s-1
                                 LMO=c(1e2,-1e2,-1e5),      #m
                                 hc= c(1),                  #m
                                 d = hc*2/3,                #m
                                 zsources = seq(0,hc,0.05), #m
                                 z = seq(0,10,0.1))         #m
{
  # Input data
  # canopy data
  # ustar = db$Ustar    #m s-1
  # LMO=db$LObukhov     #m
  # hc= db$hc           #m
  # d = db$d            #m
  # zsources = zsources #m
  # z = zmesh           #m
  
  # UC time check
  start.time <- Sys.time()
  
  #model parameters
  Duration = 30; #run duration in seconds of each line (should be larger than time spent in the canopy by the lowest trajectory)
  NbUstar = length(ustar) # number of lines of the dataset
  NbMaillesSource = length(zsources)-1; # number of source heights
  Nbz = length(z) # number of target heights
  dt = 0.01; # arbitrary time step (s)
  
  # calculate the turbulence parameters (see functions below)
  tol=sigmaw=K = array(NaN,c(length(z),length(ustar)));
  for (i in (1:length(z))){
    Res=Diffusivity_simple(z[i],ustar,hc,d,LMO);
    tol[i,]=Res$tol        # lagrangian time scale (s)
    sigmaw[i,]=Res$sigmaw  # vertical velocity standard deviation (m s-1)
    K[i,]=Res$K            # vertical diffusivity (m2 s-1)
  }
  summary(Res)
  
  #plot turbulent profiles
  plot(sigmaw[,1]/ustar[1],z)
  plot(tol[,1],z)
  plot(K[,1],z)
  
  
  
  #prerpare model inputs and outputs arrays
  dz=rep(c(diff(z)),NbUstar);dim(dz)=c(Nbz-1,NbUstar)
  dKdZ=diff(K)/ dz 
  dKdZ=rbind(dKdZ,dKdZ[dim(dKdZ)[1],])
  
  #coefficients of the randow walk model 
  Coeff1 = z*0;
  #Coeff1 = dKdZ*dt; -> accounting for the non homogeneity of the turbulence in the vertical
  Coeff2 = (2*K)^0.5;
  
  DIJ = array(0,c(NbMaillesSource,NbMaillesSource,NbUstar));
  NbPart = 1000;
  index_particle_in_z_mesh = index_particle_in_source_mesh = array(0,c(NbPart,NbUstar));
  Nbtimesteps = round(Duration/dt);
  indexesDIJ=array(F,c(NbPart*NbUstar,3))
  
  #main loop : for each source follow Nbnpart particles
  
  for(IndexSource in (1:NbMaillesSource)){
    Zpart = runif(NbPart*NbUstar)*(zsources[IndexSource+1]-zsources[IndexSource])+ zsources[IndexSource];
    dim(Zpart)=c(NbPart,NbUstar)
    BruitBlanc = rnorm(Nbtimesteps*NbPart*NbUstar,mean=0,sd=(dt)^0.5);
    dim(BruitBlanc)=c(Nbtimesteps,NbPart*NbUstar)
    
    indexesDIJ[,2]=IndexSource                   #index of source
    indexesDIJ[,3]=rep(1:NbUstar,each=NbPart)    #index of Nbustar
    
    #plot(rep(1,NbPart),Zpart,ylim=c(0,3),xlim=c(1,Nbtimesteps),pch=46)
    
    for(i in (1:Nbtimesteps)){
      
      index_particle_in_source_mesh = as.numeric(cut(Zpart,zsources,labels = c(1:(length(zsources)-1))))
      dim(index_particle_in_source_mesh)=c(NbPart,NbUstar)
      index_particle_in_z_mesh = as.numeric(cut(Zpart,z,labels = c(1:(length(z)-1))))
      dim(index_particle_in_z_mesh)=c(NbPart,NbUstar)
      
      #affect particle locations and increment DIJ NA are those out of source boundaries
      indexesDIJ[,1]=index_particle_in_source_mesh #indexes where particles are at time step i*dt
      filterna = is.na(indexesDIJ[,1])
      indexesDIJ_filtred = as.data.frame(indexesDIJ[!filterna,])
      
      fidij=aggregate(indexesDIJ_filtred, by=list(indexesDIJ_filtred$V1,
                                                  indexesDIJ_filtred$V2,
                                                  indexesDIJ_filtred$V3), length)
      
      DIJ[fidij$Group.1,fidij$Group.2,fidij$Group.3] = 
        DIJ[fidij$Group.1,fidij$Group.2,fidij$Group.3] + 
        fidij$V1*dt
      
      filterindex=!is.na(index_particle_in_z_mesh);
      Zpart[filterindex] = Zpart[filterindex] + 
        Coeff1[index_particle_in_z_mesh[filterindex]] +
        Coeff2[index_particle_in_z_mesh[filterindex]]*
        BruitBlanc[i,filterindex]
      Zpart[Zpart < 0] = -Zpart[Zpart < 0];#rebound
      if ((i%%100)==0){
        print(i)
      }
      
      #points(rep(i,NbPart),Zpart,ylim=c(0,3),xlim=c(1,Nbtimesteps),pch=46)
      
      #break if Zpart > z
      if(all(Zpart > z)) break; # not possible in foreach
      
    }
    #print(IndexSource)
    write.table(IndexSource,file='iteratorfollower.txt',append=T)
  }
  DIJ = DIJ/NbPart;
  
  end.time <- Sys.time()
  time.taken <- end.time - start.time
  print(time.taken)
  
  return (DIJ)
}



Fcorr=function(Zeta){
  phiw=1.25*(1+0.2*1)+Zeta*0
  phih= 1+5*1+Zeta*0
  
  phiw[Zeta < -2] = 1.25*(1 + 3*abs(-2))^(1/3)
  phih[Zeta < -2] = (1 + 16*abs(-2))^(-0.5)
  
  phiw[-2 <= Zeta && Zeta <=0] = 1.25*(1 + 3*abs(Zeta[-2 <= Zeta && Zeta <=0]))^(1/3)
  phih[-2 <= Zeta && Zeta <=0] = (1 + 16*abs(Zeta[-2 <= Zeta && Zeta <=0]))^(-0.5)
  
  phiw[0 < Zeta && Zeta <=1] = 1.25*(1+0.2*Zeta[0 < Zeta && Zeta <=1])
  phih[0 < Zeta && Zeta <=1] =  1+5*Zeta[0 < Zeta && Zeta <=1]
  
  Fcorrw = phiw/1.25
  Fcorrtol = (1.25)^2/(phiw^2*phih)
  Res = data.frame(Fcorrw=Fcorrw,Fcorrtol=Fcorrtol)
  return(Res)
}


Diffusivity=function(z,Ustar,hc,d,L){
  #voir arguments
  
  if (z==0) {z= 0.001}
  
  Zeta = hc/L
  if (z >= hc + d){
    Zeta = (z-d)/L
  }
  
  if  ( 0.8*hc <= z){
    x=z/hc
    x1=0.850*x+1.25
    x2=4*0.98*0.850*1.25*x
    y=(x1-(x1^2-x2)^(1/2))/(2*0.98)
    sigmaw=y*Ustar
    
    x=(z-d)/hc
    x1=0.98*0.41/1.25^2*x + 0.4*0.98
    x2=4*0.98*0.41/1.25^2*0.4*0.98*0.98*x
    y=(x1+(x1^2-x2)^(1/2))/(2*0.98)
    tol=y*hc/Ustar
    
  } else if (0.25*hc <= z && z < 0.8*hc){
    x=(z/hc)-0.8
    x3=0.256*x+0.40
    x4=4*0.98*0.256*0.40*x
    y=(x3+(x3^2-x4)^(1/2))/(2*0.98)
    tol=y*hc/Ustar
    sigmaw=0.2*Ustar*exp(1.5*z/hc)
    
  } else {
    (z<0.25*hc)
    x=4*z/hc
    x5=0.850*x+0.41
    x6=4*0.98*0.850*0.41*x
    y=(x5-(x5^2-x6)^(1/2))/(2*0.98)
    tol=y*hc/Ustar
    sigmaw=0.2*Ustar*exp(1.5*z/hc)
  }
  
  Fcorr=Fcorr(Zeta)
  sigmaw = sigmaw*Fcorr$Fcorrw
  tol = tol*Fcorr$Fcorrtol
  K=(sigmaw^2)*tol
  Res=data.frame(sigmaw=sigmaw,tol=tol,K=K)
  return(Res)
}


Diffusivity_simple=function(z,Ustar,hc,d,L){
  #voir arguments
  
  if (z==0) {z= 0.001}
  
  Zeta = hc/L
  if (z >= hc + d){
    Zeta = ((z-d)/L)
  }
  
  if (z < hc){
    tol=0.3*hc/Ustar
    sigmaw=exp(1.5*z/hc)/exp(1.5)*1.25*Ustar
    
  } 
  
  if  (z >= hc){
    sigmaw=1.25*Ustar
    tol=0.3*hc/Ustar
    
  } 
  if(z >= 1.2*hc){
    tol=0.3*(z/hc-0.7)*hc/((1.2-0.7)*Ustar)
  }
  
  Fcorr=Fcorr(Zeta)
  sigmaw = sigmaw*Fcorr$Fcorrw
  tol = tol*Fcorr$Fcorrtol
  K=(sigmaw^2)*tol
  Res=data.frame(sigmaw=sigmaw,tol=tol,K=K)
  return(Res)
}
