library(crone)
library(deSolve)

average_death <- read.table("death_rate_avg_2.csv", header=TRUE, sep=",")


beta_a = 0.8125
t_bar=8

t_max=24
end.time = t_max

initial_param = c(10, 10, 0.068,0.8125)
xstart.values = initial_param[3:4]


###########################
#  least squares function #
###########################


least.squares.fn <- function(par,data){
  
  
  delta_xa_val <- par[1]
  delta_ax_val <- par[2]
  xstart.values[1] <- par[3]
  xstart.values[2] <- par[4]
  
  
  run <- run_ODEs_death_rate(delta_xa_val,delta_ax_val,beta_a,t_bar,ODEs_death_rate,xstart.values)
  
  run_data_points <- run$'1'[c(1,3,5,7,24)]
  
  
  if (delta_xa_val < 0 || delta_ax_val < 0 || xstart.values[1] < 0 || xstart.values[2] < 0 || xstart.values[2] > 0.8125){
    error <- 1e8 
  } else{
    error <- sum((data-run_data_points)^2)/length(data)
  }
  
  
  return(error)
}

##############
#  ODE model #
##############

ODEs_death_rate <- function(t, x, parameters ) { 
  ex <- x[1]
  A <- x[2]
  
  with( as.list(parameters), { 
    dex <- -delta_xa*A*ex
    dA <- beta_a*heaviside(t_bar-t) - delta_ax*A*ex
    res <- c(dex, dA)
    list(res)
  }
  )
}


run_ODEs_death_rate <- function(delta_xa_val,delta_ax_val,beta_a_val,t_bar_val,ODEs_death_rate,xstart.values){
  times <- 1:24
  parameters <- c(delta_xa=delta_xa_val, delta_ax=delta_ax_val, beta_a, t_bar) 
  xstart <- xstart.values
  output <- as.data.frame(lsoda( xstart, times, ODEs_death_rate, parameters) )
  return(output)
}




###########################
#  OPTIMISING PARAMETERS  #
###########################

model_fitting_death_rate <- function(start){
  xstart.values = start[3:4]
  fn <- function(par) least.squares.fn(par,data=average_death$live.bacteria.fraction)
  fit <- optim(start,fn)
  while(identical(signif(fit$par,digits=5),signif(start,digits=5))==FALSE) {
    start <- fit$par
    fit <- optim(start,fn)
  }
  delta_xa_fit <- fit$par[1]
  delta_ax_fit <- fit$par[2]
  xstart.values[1] <- fit$par[3]
  xstart.values[2] <- fit$par[4]
  runfit <- run_ODEs_death_rate(delta_xa_val=delta_xa_fit, delta_ax_val=delta_ax_fit, beta_a, t_bar, ODEs_death_rate, xstart.values=xstart.values)
  plot(average_death$time, average_death$live.bacteria.fraction, col = "red", pch=18, ylim=c(0,0.08),
       xlab = "time (hours)",
       ylab = "ex")
  arrows(x0=average_death$time, y0=average_death$live.bacteria.fraction-average_death$sd, x1=average_death$time, y1=average_death$live.bacteria.fraction+average_death$sd, code=3, angle=90, length=0.05)
  lines(runfit$time, runfit$'1', col="blue")
  legend("topright", legend=c("data","model"),
         col=c("red", "blue"), pch= c(18,NA), lty=c(NA,1))
  plot(runfit$time, runfit$'2', type="l",
       xlab= "time (hours)",
       ylab = "A")
  outputs <- list(fit$par, runfit)
  #return(outputs)
  return(fit$par)
}
model_fitting_death_rate(initial_param)







###############################################
#  PARAMETER ESTIMATION USING THE ABC METHOD  #
###############################################
N = 50
epsilon = 7.5e-5
epsilon = 1e-4
accepted_values = matrix(0,N,2)
i = 0
while(i<N){
  #prior.val = sqrt((c(runif(1, 0, 2e3), runif(1, 0, 5e5),0.068,0.8125)^2))
  #prior.val = sqrt((c(runif(1, 0, 2e3),  runif(1, 0, 5e5),0.068,1.3e-5)^2))
  prior.val = sqrt((c(0.64, runif(1, 0, 5e5),0.068,0.8125)^2))
  #prior.val = sqrt((c(13.54, runif(1, 0, 5e5),0.068,1.3e-5)^2))
  #prior.val = sqrt((c(runif(1, 0, 2e3), 9.1e-7,0.068,0.8125)^2))
  #prior.val = sqrt((c(runif(1, 0, 2e3), 3.2e2,0.068,1.3e-5)^2))
  least.squares.run = least.squares.fn(prior.val, data=average_death$live.bacteria.fraction)
  if (!is.na(least.squares.run) && least.squares.run < epsilon){
    print(least.squares.run)
    i = i + 1
    accepted_values[i,] = prior.val[1:2]
  }
  cat("\r",round(i/N*100),'% done')
}
par(mfrow=c(2,1))
param_means = c()
parameter = c(expression(delta[xa]), expression(delta[ax]))
for (i in 1:2) {
  hist(accepted_values[,i], main = "", xlab = parameter[i], ylab = "Frequency")
  param_means = c(param_means, mean(accepted_values[,i]))
}
print(param_means)