### Libraries and functions ----
library(rstan)
library(here)
library(bayesplot)

options(mc.cores = parallel::detectCores())

completeFun <- function(data, desiredCols) {
  completeVec <- complete.cases(data[, desiredCols])
  return(data[completeVec, ])
}

### TRAIT MODEL STAN ---- 

setwd("/home/fbalao/Datos/R/Rpackages/Polyploidy_competition")
datacomp<-read.table(here("input","Datos_modelo_COMPLETO_WUE.txt"), header = T, stringsAsFactors = T)

datacomp$PLOIDY<-as.numeric(datacomp$PLOIDY)
datacomp$PLOIDY<-datacomp$PLOIDY*-1
datacomp$PLOIDY[datacomp$PLOIDY==-1]<-4
datacomp$PLOIDY[datacomp$PLOIDY==-2]<-1
datacomp$PLOIDY[datacomp$PLOIDY==-3]<-2
datacomp$PLOIDY[datacomp$PLOIDY==-4]<-3

datacomp$TREATMENT<-as.numeric(datacomp$TREATMENT)
datacomp$TREATMENT[datacomp$TREATMENT==2]<-0

datacomp2<-scale(datacomp[,-c(1,2,3,4,5,6,7,8,9,11,12)])
datacomp3<-cbind(datacomp[,c(1,2,3,4,5,6,7,8,9,11,12)],datacomp2)
#View(datacomp3)



datacomp3$PLOIDY<-as.numeric(datacomp3$PLOIDY)
datacomp3<- completeFun(datacomp3, c("Bwf","Bwc", "SLA_f", "SLA_c"))

Biomasa<-datacomp3$Bwf
tf<-datacomp3$SLA_f
tc<-datacomp3$SLA_c
species<-datacomp3$PLOIDY
C<-datacomp3$C
N_species <- 4
Bc<-datacomp3$Bwc
water_status<-datacomp3$TREATMENT


N<- nrow(datacomp3) # Número total de observaciones 152


# Crear una lista para usar con rstan
data_list <- list(N = N,
                  Biomasa = Biomasa,
                  water_status = water_status,
                  tf = tf,
                  tc = tc,
                  Bc = Bc,
                  species = species,
                  C = C,
                  N_species = N_species
                  )

# Compila el modelo Stan
stan_model <- stan_model(here("script","stan_models",
                              "modelo_finalBiomasacompetencia.stan"))

# Ajusta el modelo usando los datos generados
fitSLA <- sampling(stan_model, data = data_list,
                init = function() list(m0 = 0, m1 = 0, sigma = 1, alpha_t = 0,
                  alpha_e = 0, alpha_d = 0),iter = 10000, chains = 8)

saveRDS(fitSLA, file = here("output", "stan_results", "fitSLA.rds"))

# Ver resultados

A<-print(fitSLA, pars=c("m1","alpha_t","alpha_e","alpha_d"))

plot(fitSLA, pars=c("m1","alpha0intra_species","alpha0inter_species",
                 "alpha_t","alpha_e","alpha_d"))


### TRAIT MODEL STAN WITH WATER TREATMENT---- 

# Compila el modelo Stan
stan_model_water <- stan_model(here("script", "stan_models",
                                    "modelo_Biomasacompetencia_waterstatus.stan"))

# Ajusta el modelo usando los datos generados
fitSLA_water <- sampling(stan_model_water, data = data_list,
                init = function() list(m0 = 0, m1 = 0,
                sigma = 1, alpha_t = 0, alpha_e = 0, alpha_d = 0),
                iter = 10000, chains = 8)

saveRDS(fitSLA_water, file = here("output", "stan_results", "fitSLA_water.rds"))

# Ver resultados
A1<-print(fitSLA_water, pars=c("m1_dry","m1_wet",
                  "alpha_t_dry", "alpha_t_wet",
                  "alpha_e_dry","alpha_e_wet",
                  "alpha_d_dry","alpha_d_wet" ))

plot(fitSLA_water, pars=c("m1_dry","m1_wet",
                 "alpha_t_dry", "alpha_t_wet",
                 "alpha_e_dry","alpha_e_wet",
                 "alpha_d_dry","alpha_d_wet" ),  ci_level = 0.95, include = T)