#FIRST SET WD AND SOURCE THE MODEL FILE

#source it here or this won't work :) 
setwd()
source()
#

library(ggplot2)
library(reshape2)

times0 = seq(0,200, by = 0.1)


#import empirical data for comparison to model.
rat_sum = read.csv('ratio_summary.csv')

#prepare dataframe
rat_sum$cyt = rep(c(0,1,10,100),4)
rat_sum$genotype[rat_sum$dgenotype == 'clv1ab'] = 'clv1'
rat_sum$genotype[rat_sum$dgenotype == 'clv1ab;rpk2'] = 'clv1abrpk2'
rat_sum$genotype[rat_sum$dgenotype == 'rpk2'] = 'rpk2'
rat_sum$genotype[rat_sum$dgenotype == 'wt'] = 'wt'
rat_sum$norm_eb = rat_sum$SEM/rat_sum$rat_total[rat_sum$dgenotype=='wt' & rat_sum$cytokinin_status == 'bcd']

#optional subset to drop chk
#rat_sum = subset(rat_sum, rat_sum$cytokinin_status != 'chk')

#normalize data
dat_sum_norm =  rat_sum$rat_total/ rat_sum$rat_total[rat_sum$dgenotype == 'wt' &rat_sum$cytokinin_status == 'bcd']

rat_sum$initNorm = dat_sum_norm
rat_sum$upperN = rat_sum$initNorm+rat_sum$norm_eb
rat_sum$lowerN = rat_sum$initNorm-rat_sum$norm_eb

rat_sum$genotype = factor(rat_sum$genotype, levels = c('wt','clv1','rpk2','clv1abrpk2'))


#Put in your starting parameters here; in this case for model 6. \
#interaction parameters
kb = c(cyt_x = 1, # kb1
       clv_x = 0,  #kb2
       rpk2_x = 1,  #kb3
       clv_y = 1,  #kb4
       rpk2_y = 0, #kb5
       x_init = .1, #kb6
       y_init = .1, #kb7
       cyt_z = .5, #kb8
       z_init = .5 #kb9
       ) 
#synthesis parameters 
pb = c(.01, 
       .03,
       0.01,
       0.01)

#initial values
base0 = c(x = 1,
          y = 1,
          init = 1,
          z = 1)
#degredation values, set to 5 percent
db = c(0.05,
       0.05,
       0.05,
       0.05)
#make list with the set of conditions 

conset = list(kb, pb, db, base0)


# TEST RUN ####
#make sure your model converges

#starting values for variables
parmsb = c(cyt = 0.1, clv = 3, rpk2 = 3)

#run the model
#FOR MODEL 7: change model to 'endfb3'
testrun = lsoda(func = endfb4, y = base0, times = times0, parms = parmsb)
testrunD = data.frame(testrun) 

#view model output
testrunD = melt(testrunD, id.var = 'time')
ggplot(data = testrunD, aes(x = time, y = value, color = variable, group = variable))+
  geom_line(size = 1)+
  theme_bw()

#How good is the fit?####

#In this section test how good the fit is for the base set of conditions across mutants and cytokinin values 
#want

#record initial variables
ogkb = kb
pbo = pb
base0o = base0
#make dataframe to log values as model runs through different conditions
ssvals_dto = data.frame()

#run the model through different conditions
for(v in c(3,0)){
  #parmsb[3] is rpk2; rpk2 = 3 is 'wt', rpk2 =0 is rpk2 mutant
  parmsb[3] = v  
  for(q in c(3,0)){
    #parmsb[2] is clv1; clv = 3 is 'wt', clv =0 is clv mutant
    parmsb[2] = q
    for(i in c(0,1,10,100)){
      #parmsb[1] is cytokinin; set to 0, 1, 10, 100 for chk mutant, mock, 10nM, 100nM
      parmsb[1] = i
      baserun = lsoda(func = endfb4, y = base0, times = times0, parms = parmsb)
      #pull out a vector to fill the data frame with, taking the steady state value
      fill_vector = c(baserun[2000, 2:5], 'Cyt' = parmsb[1], 'Clv' = parmsb[2], 'RPK2' = parmsb[3])
      ssvals_dto = rbind(ssvals_dto, fill_vector)
    }
  }
}
kb = ogkb
pb = pbo
base0 = base0o
#clean up the data frame, adding genotype values ####
colnames(ssvals_dto) = c('x','y','br_init', 'z', 'cyt', 'clv','rpk2')

#fill genotype column with the custom function
ssvals_dto = name_genotypes(ssvals_dto)

#normalize the scores
ssvals_dto$norm = ssvals_dto$br_init/ssvals_dto$br_init[ssvals_dto$genotype=='wt' & ssvals_dto$cyt == 1]

#then get each individual fit score for each mutant/cytokinin treatment (16 of them) by dividing the normalized model score over the corresponding normalized empirical data. Since each should ideally be one, find the divergence then. Since being half the expected value is just as bad as twice the expected value, take the square of the log of each fold change sum it.

log_divergence = log(ssvals_dto$norm/dat_sum_norm)
mod_tot_score = sum(log_divergence^2)


#run the optimizer ####
#make empty data frame for optimizer
optimizerDat = data.frame()

#number of rounds to run the optimizer
for (iteration in 1:300){
  
  #1er kb, pb, and db, but store an original value that we can reset to if our new parameter isn't better. 
  #set the potency of the mutation per round (standard dev)
  potency = 0.1
  
  #how long is the vector to mutate the parameters (the length of all the parameters)
  mlength = length(kb) + length(pb) + length(db) + length(base0)
  mutator = rnorm(mean = 1, n = mlength, sd = potency)
  
  #in three steps, for interaction, synthesis etc. 
  #recorded parameters
  #set current parameters to be the mutated version
  #mutate recorded parameters
  
  ogkb = conset[[1]]
  tkb = mutator[1:9]*ogkb
  kb = tkb
  
  ogpb = conset[[2]]
  tpb = mutator[10:13]*ogpb
  pb = tpb
  
  ogdb = conset[[3]]
  tdb = mutator[14:17]*ogdb
  db = tdb
  
  ogbase = conset[[4]]
  tbase0 = mutator[18:21]*ogbase
  base0 = tbase0
  
  # run the model
  # make an empty data frame
  #will pass through cytokinin 0, 0.5, and 10, and 100, which are what I'm interested in here
  #next loop will pass through clv1 values 0 (clv1 mutant) and 3 (wt)

  ssvals_dto2 = data.frame()
  for(v in c(3,0)){
    parmsb[3] = v
    for(q in c(3,0)){
      parmsb[2] = q
      for(i in c(0,1,10,100)){
        parmsb[1] = i
        baserun = lsoda(func = endfb4, y = base0, times = times0, parms = parmsb)
        #pull out a vector to fill the data frame with, taking the steady state value
        fill_vector = c(baserun[2000, 2:5], 'Cyt' = parmsb[1], 'Clv' = parmsb[2], 'RPK2' = parmsb[3])
        ssvals_dto2 = rbind(ssvals_dto2, fill_vector)
      }
    }
  }
  kb = ogkb
  pb = pbo
  base0 = base0o
  #clean up the data frame, adding genotype values ####
  colnames(ssvals_dto2) = c('x','y','br_init', 'z', 'cyt', 'clv','rpk2')
  
  #fill genotype column
  #at which cells is CLV1 intact, ie wt ie 4? Same for RPK2
  
  ssvals_dto2 = name_genotypes(ssvals_dto2)
  
  #normalize the scores
  ssvals_dto2$norm = ssvals_dto2$br_init/ssvals_dto2$br_init[ssvals_dto2$genotype=='wt' & ssvals_dto2$cyt == 1]
  
  
  #then get each individual fit score (16 of them) by dividing the normalized model score over the corresponding normalized empirical data. Since each should ideally be one, find the divergence then. Since being half the expected value is just as bad as twice the expected value, take the square of the log of each fold change sum it.
  
  log_divergence = log(ssvals_dto2$norm/dat_sum_norm)
  totalScore = sum(log_divergence^2)
  
  
  
  print(paste('Total score is ', totalScore))
  print(paste("Previous Best is " , mod_tot_score))
  
  #We want to reduce each score, or at least get them in an acceptable range
  if(totalScore < mod_tot_score){
    print('found a better set!')
    
    #update the benchmark to the new score
    mod_tot_score = totalScore
    
    #update the conset. These new values will be grabbed and mutated again for the next iteration.
    
    conset[[1]] = tkb; conset[[2]] = tpb; conset[[3]] = tdb; conset[[4]] = tbase0; 
    
  } #else don't update the conset
  
  #record kb values and the additivity
  storevec = c(tkb, tpb, tdb, tbase0, score = totalScore, iteration)
  optimizerDat = rbind(optimizerDat, storevec)
  
  #also want to store the normalized values, the score diff, 
  
  
  print(paste('iteration', iteration))
} # iterations loop

#apply names of parameters to the data frame
colnames(optimizerDat) = c('kb1','kb2','kb3','kb4','kb5','kb6','kb7', 'kb8','kb9',
                           'pb1','pb2','pb3', 'pb4',
                           'db1','db2','db3', 'pb5',
                           'init1','init2','init3', 'init4',
                           'totalScore', 'iteration')


#If need to read in optimizer dat
optd = optimizerDat

#what is the best iteration?
bestScore = min(optd$totalScore)
bestIter = which(optd$totalScore == bestScore)


#Import the coefficient values from that best iteration
kb = as.numeric(optd[bestIter,1:9])
kbs = kb
pb = as.numeric(optd[bestIter, 10:13])

db = as.numeric(optd[bestIter, 14:17])

base0 = as.numeric(optd[bestIter,18:21])

#plot the fit over iterations

total_fit_plot = ggplot(data = optimizerDat, aes(x = iteration, y = totalScore))+
  geom_line()+
  ggtitle('base-1, means')


total_fit_plot


#now that the coefs are imported, run the model through a cytokinin gradient for displaying on figures

ogkb = kb
pbo = pb
base0o = base0
ssvals_dto = data.frame()
for(v in c(3,0)){
  parmsb[3] = v
  for(q in c(3,0)){
    parmsb[2] = q
    for(i in c(0,1,2,4,7.5,10,20,35,60,100)){
      parmsb[1] = i
      baserun = lsoda(func = endfb4, y = base0, times = times0, parms = parmsb)
      #pull out a vector to fill the data frame with, taking the steady state value
      fill_vector = c(baserun[2000, 2:4], 'Cyt' = parmsb[1], 'Clv' = parmsb[2], 'RPK2' = parmsb[3])
      ssvals_dto = rbind(ssvals_dto, fill_vector)
    }
  }
}
kb = ogkb
pb = pbo
base0 = base0o
#clean up the data frame, adding genotype values ####
colnames(ssvals_dto) = c('x','y','br_init', 'cyt', 'clv','rpk2')

#fill genotype column
#at which cells is CLV1 intact, ie wt ie 4? Same for RPK2

dat = name_genotypes(ssvals_dto)

#normalize the scores
dat$initNorm = dat$br_init/dat$br_init[dat$genotype=='wt' & dat$cyt == 1]


#plot the data with the empirical data

plot1 = ggplot(data = NULL, aes(x = log(cyt+1), y = initNorm, color = genotype))+
  geom_line(data = dat, aes(group = genotype), size = 1)+
  theme_bw()+
  ggtitle(paste("Model Test 2.1.3 initiation"))+
  labs(y = 'Branch initiation')+
  geom_point(data = rat_sum, size = 3)+
  geom_errorbar(data = rat_sum, aes(ymin = lowerN, ymax = upperN), width = 0.2)
#geom_vline(xintercept = log10(0.5), color = 'cyan', alpha = 0.4, lwd = 5)+

plot1

plot2 = plot1 + facet_grid(cols = vars(genotype))

plot2