######################################## for both known and unknown mean #####################################

library(RandomFields)
library(ggplot2)

#cat("\014")
#set.seed(999)


subcase = 1        ### choosing the covariance functions
l = 1                  ### for cosine covariance only

mu = 0                ### Mean 
Var0=1                ### variance

################################# the covariance functions ###################################################


Model_case<- switch(subcase,
                    model<- RMtrend(mean=mu) +RMexp(var=Var0,scale=1),
                    model<- RMtrend(mean=mu) +RMgauss(var=Var0, scale=1/sqrt(0.1)),
                    model<- RMtrend(mean=mu) +RMdampedcos(lambda=0, var=Var0, scale=0.6),
                    model<- RMtrend(mean=mu) +RMwave(var=Var0, scale=1),
                    model<-RMtrend(mean=mu) +RMbessel(nu=0, var=Var0, scale=1),
)


Model_case
plot(Model_case,xlim=c(0,20))

####################################### user defined variables ################################################



T = 100              ### upper limit of time points on [0,T]
N = 50               ### discretization of time points on [0,T]
Delta = T/N            ### step size of discretization of time points
observations = 10  ### number of observations on [0,T]
n = observations  

delta = T/n            ### step size of extrapolation 
dimension = 2


#dimension=1
step1=round(N/n)
x = seq(0,T,Delta)
N=length(x)
cord=expand.grid(x)   # points of extrapolations
t_obs=which((1:N)%% step1==0,arr.ind = TRUE) #indecies of observations
cord_0 = cord[t_obs,] # points of observations

if  (dimension == 2){
  
  cord = expand.grid(x,x)
  tcord = expand.grid(1:N,1:N)
  t_obs=which((tcord[,1]%% step1==0)&(tcord[,2]%% step1==0),arr.ind = TRUE)
  cord_0 = cord[t_obs,]
  names(cord)<-c("x","y")
} else {
  if (dimension == 3){
    cord = expand.grid(x,x,x)
  tcord = expand.grid(1:N,1:N,1:N)
  t_obs=which((tcord[,1]%% step1==0)&(tcord[,2]%% step1==0)&(tcord[,3]%% step1==0),arr.ind = TRUE)
  cord_0 = cord[t_obs,]
  names(cord)<-c("x","y","z")
  }
}

n=length(t_obs)        # number of observations
N1=length(cord[,1])    # number of extrapolation points






mean_vector  = rep(mu,N1)
mean_vector_reg  = rep(mu,n)


######################################## sigma_ext is matrix \Sigma in the paper ###############################################################################


sigma_ext=RFcovmatrix(model=Model_case,x=cord_0)
QR = qr(sigma_ext)

########################################## cc_t is a matrix C_{t} in the paper ########################################################################################

cc_t=RFcovmatrix(model=Model_case,x=cord)
cc_t=cc_t[,t_obs]


###################################### Unknown Mean Case #####################################################

############# ee matrix is matrix e in the paper ########

cat("Creating matrices for known and unknownn mean case.\n")
ee = matrix(data = replicate(n,1), ncol = 1);

##############################################################

############## bb_0 and bb_1 are matrix b_0  and b_1 in theorem 4.2 ##########

bb_0 = rep(0,N1)
bb_1 = rep(0,N1)
Sigma.inv.c=matrix(rep(0,N1*n),nrow=N1)
for (i3 in 1:N1){
  Sigma.inv.c[i3,]= solve.qr(QR, cc_t[i3,])
  bb_0[i3] = t(cc_t[i3,])%*%Sigma.inv.c[i3,]
  bb_1[i3] =  t(ee)%*%Sigma.inv.c[i3,]
}

Sigma.inv.e=t(solve.qr(QR, ee))
Sigma.inv.e=Sigma.inv.e[rep(1,N1),]

###############################################################

############# bb_2 is matrix b_2 in theorem 4.2 ############

bb_2 = as.vector(t(ee) %*% solve.qr(QR, ee))

cat("\t\t\t\t\t\t Task complete\n")
###################### Matrices Created ######################
cat("Computing weights (lambdas) for known and unknown mean case, simple and ordinary kriging case and Gaussian linear regression case.")

lambda_unknown_mean = matrix(rep(0,n*N1), nrow = N1, ncol = n)
lambda_known_mean = lambda_unknown_mean
lambda_ord_krig = lambda_unknown_mean
lambda_sim_krig = lambda_unknown_mean



##---- Unknown mean ---------------
first_term_unknown = matrix(sqrt((Var0*bb_2 - 1) / (bb_0*bb_2 - (bb_1)^2)), nrow = N1)
second_term_unknown = matrix(bb_1/bb_2,nrow=N1)

first_term_unknown=first_term_unknown[,rep(1,n)]
second_term_unknown=second_term_unknown[,rep(1,n)]

lambda_unknown_mean=first_term_unknown*Sigma.inv.c-first_term_unknown*second_term_unknown*Sigma.inv.e+Sigma.inv.e/bb_2

##---- Known mean ---------------
first_term_known=matrix(bb_0,nrow=N1)
first_term_known=sqrt(Var0/first_term_known[,rep(1,n)])
lambda_known_mean=first_term_known* Sigma.inv.c 


##---- Ordinary Kriging ---------------
delta_KR=matrix((1-bb_1)/bb_2,nrow=N1)
delta_KR=delta_KR[,rep(1,n)]
lambda_ord_krig=Sigma.inv.c+delta_KR*Sigma.inv.e

##---- Simple Kriging ---------------
lambda_sim_krig =Sigma.inv.c

#------Extrapolation---------------------------

X = RFsimulate(model = Model_case, x=cord,dim=dimension,n=1,spConform=FALSE)

X_hat_known_mean = lambda_known_mean %*% X[t_obs]
X_hat_unknown_mean = lambda_unknown_mean %*% X[t_obs]
X_hat_ok = lambda_ord_krig %*% X[t_obs]
X_hat_sk = lambda_sim_krig %*% X[t_obs]


cat("X_hat for all above mentioned cases are calculated. Now, poltting the results for comparison.")
cat("\n\n")

if (dimension == 1) {

# pdf(paste("comparison", "observation = ", n, ".pdf", sep = ""),
#     width = 5.52, height = 3.47)

par(mar = c(7, 4, 4, 3), xpd = TRUE)
plot(x, X, type = "l", col = "red", main = "Gaussian Realisation & Extrapolated Realisation",
     ylab = "X_hats & X", cex.main = 0.9)
points(x, X_hat_known_mean, type = "l", col = "green")
points(x, X_hat_unknown_mean, type = "l", col = "blue")
points(x, X_hat_ok, type = "l", col = "black")
points(x, X_hat_sk, type = "l", col = "orange")
legend("bottom", inset=c(0,-1), legend=c("X","X_hat_known_mean","X_hat_unknown_mean","X_hat_ok","X_hat_sk"),
       lty = c(1, 1, 1, 1, 1), col = c("red","green","blue", "black", "orange"),
       title="Extrapolation", cex = 0.5 , ncol = 3)

# legend("bottom", inset=c(0,-1), legend=c("X","X_hat_unknown_mean"), lty = c(1, 1),
#        col = c("red","blue"), title="Extrapolation", cex = 0.5 , ncol = 2)  ## use this for any two comparison
#dev.off()

#################################### Level Set Comparison #########################

#-----number of repetitions
Nr=1200

XX = RFsimulate(model = Model_case, x=x,dim=1,n=Nr,spConform=FALSE)

simmdiff<-function(u,Xhat,Xtrue)
{
  return(sum(Xhat>u)+sum(Xtrue>u)-2*sum((Xtrue>u)*(Xhat>u)))
}

#-----
u=-0
DeltaM=matrix(rep(0,4*Nr),nrow=4)
XH=array(rep(0,4*Nr*N1),dim=c(4,N1,Nr))
for (ip in 1:Nr)
{
  X_1 = lambda_known_mean %*% XX[t_obs,ip]
  X_2 = lambda_unknown_mean %*% XX[t_obs,ip]
  X_3 = lambda_ord_krig %*% XX[t_obs,ip]
  X_4 = lambda_sim_krig %*% XX[t_obs,ip]
  XH[1,,ip]=X_1
  XH[2,,ip]=X_2
  XH[3,,ip]=X_3
  XH[4,,ip]=X_4
  DeltaM[1,ip]=simmdiff(u,as.vector(X_1),XX[,ip])
  DeltaM[2,ip]=simmdiff(u,as.vector(X_2),XX[,ip])
  DeltaM[3,ip]=simmdiff(u,as.vector(X_3),XX[,ip])
  DeltaM[4,ip]=simmdiff(u,as.vector(X_4),XX[,ip])
}
DeltaM=DeltaM*Delta



#png(paste("boxplot_comparison", "observation = ", n, ".pdf", sep = ""),    width = 7.52, height = 3.47)

boxplot(DeltaM[1,], DeltaM[2,], DeltaM[3,], DeltaM[4,],
        main = paste("Comaprison of new extrapolation method \n with simple and ordinary kriging and Gaussian linear regression"),
        names = c("X_hat_known","X_hat_unknown","X_hat_ok", "X_hat_sk"),
        col = c(2,3,4,5), cex.main =0.8)
#dev.off()


boxplot(apply(XX, 1, var),apply(XH[1,,], 1, var),apply(XH[2,,], 1, var),apply(XH[3,,], 1, var),apply(XH[4,,], 1, var), main="Variance")

boxplot(apply(XX, 1, mean),apply(XH[1,,], 1, mean),apply(XH[2,,], 1, mean),apply(XH[3,,], 1, mean),apply(XH[4,,], 1, mean), main="Mean")

}

if (dimension ==2)
{
  
  Rt=cbind(cord,X,rep("Real",N1))
  names(Rt)=c(names(cord),"V","Type")
  Rall=Rt
  Rt=cbind(cord,X_hat_known_mean,rep("Known mean",N1))
  names(Rt)=c(names(cord),"V","Type")
  Rall=rbind(Rall,Rt)
  Rt=cbind(cord,X_hat_unknown_mean,rep("Unknown mean",N1))
  names(Rt)=c(names(cord),"V","Type")
  Rall=rbind(Rall,Rt)
  Rt=cbind(cord,X_hat_ok,rep("Ordinary kriging",N1))
  names(Rt)=c(names(cord),"V","Type")
  Rall=rbind(Rall,Rt)
  Rt=cbind(cord,X_hat_ok,rep("Ordinary kriging",N1))
  names(Rt)=c(names(cord),"V","Type")
  Rall=rbind(Rall,Rt)
  Rt=cbind(cord,X_hat_sk,rep("Simple kriging",N1))
  names(Rt)=c(names(cord),"V","Type")
  Rall=rbind(Rall,Rt)
  
  
  ggplot(data=Rall,aes(x=x, y=y))+geom_raster(aes(fill = V),alpha = .8,interpolate = TRUE)+facet_grid(.~Type)+scale_y_continuous(breaks = cord_0[,2])+scale_x_continuous(breaks = cord_0[,1])
  
}

cat("\n\t\t\t\t\t\t Simulation Complete. \n\t\t\t\t\t\t Thank You ('..')")