# Load libraries

Sys.getpid()

rm(list=ls())
gc(TRUE)
library(mgcv)
library(Matrix)
library(MASS)

# Load function to estimate spatio-temporal expectile regression
source("function_expectreg_mgcv_bam11.R")

# Define charactersitics of the simulation study
n      <- 2000        # Number of observations
n_pred <- 10000       # Number of observations for PMWSE
N <- 100              # Number of replications

Sim_type <- "2"       # Number of Simulation styles, Check paper for details on the functions
h_type <- "homosc"  # Heteroscedasticity or homoscedasticity

fit_type <- "ti"      # Should iterations be used in the estimation (alternatively only the marginal effects are used)
delta  <- 1e-4        # Parameter for the Schall algorithm
step_max <- 10        # Parameter for the Schall algorithm
reltol <- 1e-5        # Parameter for the optim-gcv algorithm
bounce_off <- TRUE    # Parameter for the optim-gcv algorithm
initial_w <- TRUE     # Should initial weights be used in the estimation procedures

# Define formula of the model
if(fit_type == "ti") {
    f_exp <- y ~ ti(x1,k=15,bs="ps",np=F) + ti(x2,k=15,bs="ps",np=F) + ti(x1,x2,k=c(15,15),bs=c("ps","ps"),np=F)
} else { 
    f_exp <- y ~  s(x1,k=15,bs="ps")      +  s(x2,k=15,bs="ps")
}

# Define names for saving outputs
text_iniw <- "T"
if(!initial_w) text_iniw <- "F"
text_bounce <- "T"
if(!bounce_off) text_bounce <- "F"
Name1 <- paste("Sim4_",n,"_",N,"_",fit_type,"_",text_iniw,"_",h_type,"_",text_bounce,"_",Sim_type,sep="")

Name_Sav <- paste(Name1,".RData",sep="")
Name_Img <- paste(Name1,".pdf",sep="")


# Define simulation design
f1 <- function(x1,x2,sx1,sx2){
    10*pi*sx1*sx2*(1.2*exp(-(x1-0.2)^2/sx1^2 - (x2-0.3)^2/sx2^2) + 0.8*exp(-(x1-0.7)^2/sx1^2 - (x2-0.8)^2/sx2^2))
}

f2 <- function(x1,x2){
    2*sin(pi*x1) + exp(2*x2)
}

f3 <- function(x1,x2){
    1.9*(1.45 + exp(x1) * sin(13*(x1-0.6)^2))*exp(-x2)*sin(7*x2)
}

f4 <- function(x1,x2,f1,f2,sx1,sx2){
    (f1(x1,x2)+f2(x1,x2))/2
}

# Function to simulate the data
Sim_Data_fun <- function(n, Sim_type = c("1","2","3","4"), h_type = c("homosc","heterosc"),f1,f2,f3,f4) {
    
    # draw covariates
    x1 <- runif(n)
    x2 <- runif(n)
    
    sx1 <- 0.3
    sx2 <- 0.4
    
    # Build predictor
    if(Sim_type == "1") {
        eta <- f1(x1=x1,x2=x2,sx1=sx1,sx2=sx2)
    }
    if(Sim_type == "2") {
        eta <- f2(x1=x1,x2=x2)
    }
    if(Sim_type == "3") {
        eta <- f3(x1=x1,x2=x2)
    }
    if(Sim_type == "4") {
        eta <- f4(x1=x1,x2=x2,f1=f1,f2=f2,sx1=sx1,sx2=sx2)
    }
    
    # Build error term
    if(h_type == "homosc") {
        eps <- rnorm(n, mean=0, sd=2)
    }
    if(h_type == "heterosc") {
        eps <- rnorm(n, mean=0, sd=(1/1.5*(x1+1) + 1/1.5*(x2+1)))
    }
    
    # Calculate response
    y <- eta + eps
    
    data.frame(y,x1,x2)
}

# Function to evaluate the output (We call the MWSE here deviance)
DEV <- function(y,fitted,expectile=0.5) {
    weights <- rep(expectile,times=length(y))
    weights[y >= fitted] <- expectile
    weights[y < fitted] <- 1 - expectile
    mean(weights*(y-fitted)^2)
}

# Initialize output summaries
MWSE_matrix <- matrix(NA,ncol=22,nrow=N)
colnames(MWSE_matrix) <- c("01_bam_gcv","01_bam_sch_d",
                           "02_bam_gcv","02_bam_sch_d",
                           "05_bam_gcv","05_bam_sch_d",
                           "10_bam_gcv","10_bam_sch_d",
                           "20_bam_gcv","20_bam_sch_d",
                           "50_bam_gcv","50_bam_sch_d",
                           "80_bam_gcv","80_bam_sch_d",
                           "90_bam_gcv","90_bam_sch_d",
                           "95_bam_gcv","95_bam_sch_d",
                           "98_bam_gcv","98_bam_sch_d",
                           "99_bam_gcv","99_bam_sch_d")

iterations <- sm_par_vec_x1 <- sm_par_vec_x2 <- sm_par_vec_x1x2_1 <- sm_par_vec_x1x2_2 <- PMWSE_matrix <- MWSE_matrix


# Simulate data for all runs jointly, to be able to rerun one replication and check details
data_list <- list()
set.seed(1)

for(i in 1:N) {
    data_list[[i]] <- Sim_Data_fun(n,Sim_type=Sim_type, h_type=h_type,f1=f1,f2=f2,f3=f3,f4=f4)        
}

# Simulate the data for the predictive MWSE
data_pred <- Sim_Data_fun(n_pred,Sim_type=Sim_type, h_type=h_type,f1=f1,f2=f2,f3=f3,f4=f4)        


# Define which asymmetry levels should be used
expectiles_vec <- c(0.01,0.02,0.05,0.10,0.20,0.50,0.80,0.90,0.95,0.98,0.99)

# Execute prelicated estimation
for(i in 1:N) {
    cat("\n",i," ")
    data1 <- data_list[[i]]
    
    # Initial model
    set.seed(111)
    rainmodel1 <- bam(f_exp, data = data1, family = gaussian(), weight=rep(0.5,times=n),method="GCV.Cp")
    
    # Initial smoothing parameter
    sm_par_vec2 <- rainmodel1$sp
    
    quietly <- FALSE
    if(i>1) quietly <- TRUE
    
    # Run model estimation for different asymmetry levels
    for(jj1 in expectiles_vec) {
        cat(jj1, " ")
        # Run model estimation for different smoothing parameter optimization
        for(jj2 in c("bam_sch_d","bam_gcv")){#"bam_sch_n",
            cat(jj2, " ")
            
            # Define column of output savings
            text1 <- as.character(jj1*100)
            if(nchar(text1) == 1) text1 <- paste("0",text1,sep="")
            jj <- paste(text1,jj2,sep="_")
            model23 <- NULL
            # Calculate the model
            if(jj2 == "bam_gcv"){
                set.seed(111)
                model_23 <- expectreg_bam_smooth(formula=f_exp, data=data1, expectiles=jj1, sm_par_vec = sm_par_vec2, 
                                                fixed=FALSE,delta = delta, step_max = step_max, reltol = reltol, 
                                                initial_w = initial_w, opt_type = "gcv", dev_schall = "non",
                                                quietly=quietly,bounce_off=bounce_off)
                
            } else {
                set.seed(111)
                model_23 <- expectreg_bam_smooth(formula=f_exp, data=data1, expectiles=jj1, sm_par_vec = sm_par_vec2, 
                                                fixed=FALSE,delta = delta, step_max = step_max, reltol = reltol, 
                                                initial_w = initial_w, opt_type = "schall", dev_schall = "dev",
                                                quietly=quietly,bounce_off=bounce_off)
                
            }
            
            # Build the predicted valued for the PMWSE
            fitted_23 <- predict(model_23$model,newdata=data_pred,"response")
            
            # Calculate the MWSE and PMWSE
            MWSE_matrix[i,jj]  <- DEV(data1$y    , model_23$fitted, as.numeric(jj1))
            PMWSE_matrix[i,jj] <- DEV(data_pred$y, fitted_23      , as.numeric(jj1))
            
            # Save other details of the estimation
            temp_iterations <- NULL
            if(jj2=="bam_gcv") { 
                temp_iterations <- model_23$model_optim$count[1]
            } else {
                temp_iterations <- model_23$model_optim$iter
            }
            iterations[i,jj] <- temp_iterations
            
            # Save estimated smoothing parameters 
            if(fit_type == "ti"){
                sm_par_vec_x1[i,jj] <- model_23$sm_par_vec["ti(x1)"]
                sm_par_vec_x2[i,jj] <- model_23$sm_par_vec["ti(x2)"]
                
                
                sm_par_vec_x1x2_1[i,jj] <- model_23$sm_par_vec["ti(x1,x2)1"]
                sm_par_vec_x1x2_2[i,jj] <- model_23$sm_par_vec["ti(x1,x2)2"]
            }
            if(fit_type == "s"){
                sm_par_vec_x1[i,jj] <- model_23$sm_par_vec["s(x1)"]
                sm_par_vec_x2[i,jj] <- model_23$sm_par_vec["s(x2)"]
            }
            
        }
    }
}

# Save workspace and quit
save.image(Name_Sav)


if(.Platform$OS.type == "unix"){
    q("no")
}

