# this script makes multispecies projections using the OP model given
#   a range of F targets by species

OP.simulate<-function(
          scenario="my-scenario",
          my.area=c('North Sea','Baltic Sea')[1],
          my.last.year=2061, years.in.average=50,

          HCR=NA,
          FBlim.adjust=1,

          stochastic.recruitment=0,
          recruit.adjust.CV=0,        #adjust recruitment with half of the variance (factor exp(-(CV^2)/2).  0=no adjustment, 1=do adjustment

          first.run.no=1,
          first.no.iter.stoch=1,
          no.iter.stoch=10,  # default for stochastic recruitment

          do.indicators=T,

          do.simulation=T,     # do the simulations and create data files for plots
          just.batch.file=F,   # make a batch file for a later run

          combine.scenario.run=F,    # special for combining scenarios
          scenario.dirs='a',         # name of scenarios to be combined
          scenario.dirs.out='Samlet',   # directory name for output R data set

          adjustBaltic=F,
          OP.output=14,  # 14 for combinations of F
          change.Other.pred.N=F,
          adjust.ini.N=T,     # adjust inital population size  (NS cod)

          read.condense=T,
          read.indicators=T,

          do.plot.indicators=T,
          do.plot.condense=T,
          ReduceTable=F,      # stepwise selection table of F from given criteria
          ReduceTableDetails=F,   # aoutput full (detalied results after each reduction)
          riskLevels=c(5,5,5,5,5,5,5,5), # Risk levels for SSB below Blim
          keepPro=0.95,  # keep F values which is factor keepPro of  MSY, used in stepwise reduction

          subset.out=F,       # make a subset of data set used to make plots
          justAboveBlim=F,    # make plot for Fcombinations where the risk to Blim <5%
          justAboveBpa=F,
          subset.species=matrix(c(rep(0,10),rep(2,10)),ncol=10,byrow=T),  # min and max F for inclusion in plots
          ask.me=F,
          targetFs=NA,   # combinations of F

           my.dev='png'
){

first.year.output<-my.last.year-years.in.average+1
cat("first year used in output and in average values:",first.year.output,"\nlast year:",my.last.year,"\n")

no.iter<-1  # for determenistic recruitment
if (stochastic.recruitment>0) no.iter<-no.iter.stoch

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

#  runs are made in a separate dirictory

scenario<-paste(scenario,paste("HCR",HCR[1],sep=''),FBlim.adjust[1],paste("Rec",stochastic.recruitment,sep=''),ifelse(recruit.adjust.CV[1]==1 | recruit.adjust.CV[1]==2,"Recadj",""),my.last.year,sep='_')

scenario.gem<<-scenario
scenario.dir<-file.path(data.path,scenario)

ref<-Read.reference.points.OP(dir=data.path)
Blim<-as.vector(ref[,'Blim'])
Bpa<-as.vector(ref[,'Bpa'])
Fpa<-as.vector(ref[,'Fpa'])

if (my.area=='North Sea'){
  #             Cod     Whiting     Haddock      Saithe     Herring     Sandeel   Nor. pout       Sprat      Plaice        Sole
  TSBBlim<- c(    1,          1,          1,          1,          1,     787000,     263000,      157000,          1,          1)
  TSBBpa<-  c(    1,          1,          1,          1,          1,    1098000,     440000,      213000,          1,          1)
} else {
   TSBBlim<- c(1,1,1)
   TSBBpa<-  c(1,1,1)
}

# Make an OP control object. File='OP.dat'
setwd(data.path)
source(file=file.path(prog.path,'HCR_OP_batch_common.R'))

res<-make.OP.dat(my.area=my.area,my.last.year=my.last.year,first.year.output=first.year.output,
                 do.indicators=do.indicators,
                 stochastic.recruitment=stochastic.recruitment)
OP<-res[["OP"]]
SMS<-res[["SMS"]]

sp.name<-SMS@species.names

 OP@recruit.adjust.CV[1,]<-recruit.adjust.CV

if (my.area=='Baltic Sea' & adjustBaltic) {# adjustment was done when environmental S-R were used
 OP@recruit.adjust.CV[1,1]<-0     #adjust recruitment with half of the variance (factor exp(-(CV^2)/2).  0=no adjustment, 1=do adjustment
 OP@recruit.adjust.CV[1,2]<-1     #adjust recruitment with half of the variance (factor exp(-(CV^2)/2).  0=no adjustment, 1=do adjustment
 OP@recruit.adjust.CV[1,3]<-1     #adjust recruitment with half of the variance (factor exp(-(CV^2)/2).  0=no adjustment, 1=do adjustment
}

OP@F.or.C[]<-31
OP@output<-OP.output

write.FLOP.control(OP,file='OP.dat',nice=T)

if (length(FBlim.adjust)==1) FBlim.adjust<-rep(1,length(Fpa))

### OP_trigger file
  nsp<-SMS@no.species
  n.other.pred<-sum(SMS@species.info[,'predator']==2)
  n.pred<-n.other.pred+sum(SMS@species.info[,'predator']==1)
  n.vpa<-nsp-n.other.pred
  n.vpa.pred<-sum(SMS@species.info[,'predator']==1)

OPT<-read.FLOPtrigger.control(file="OP_trigger.dat",n.VPA=n.vpa,n.other.pred=n.other.pred)
OPT@last.year<-my.last.year
OPT@first.year<-first.year.output
OPT@HCR[]<-HCR
OPT@first.run.no<-first.run.no
OPT@first.iter.no<-first.no.iter.stoch
OPT@no.iter<-no.iter

OPT@trigger['T1',]<-Blim* FBlim.adjust
OPT@trigger['T2',]<-Bpa
OPT@at.age.weighting<-1
for (s in(1:n.vpa)) if (HCR[s]==22) {
  OPT@trigger['T1',s]<-TSBBlim[s]* FBlim.adjust[s]
  OPT@trigger['T2',s]<-TSBBpa[s];
}
write.FLOPtrigger.control(OPT,file="OP_trigger.dat")

# MSFD and other options
if (do.indicators) {
 MSFD<-FLOP.MSFD.control()

 if (my.area=='Baltic Sea') {
   MSFD@do.community.biomass.demersal<-1
   MSFD@community.biomass.demersal.ages[ ,2:3]<- -1     # cod (species 1) is a demersal species

   MSFD@do.community.biomass.small<-1
   MSFD@community.biomass.small.ages['last.age',1]<- 1     # fish below 25 cm is small, that includes all herring and sprat, cod age 0-1 is small species

   MSFD@do.community.biomass.pelagic<-1
   MSFD@community.biomass.pelagic.ages[,1]<- -1     # herring and sprat (species 2, 3) are pelagic species

   MSFD@do.community.biomass.forage<-1
   MSFD@community.biomass.forage.ages[,1]<- -1     # herring and sprat (species 2, 3) are forage species

   MSFD@do.M2.bar<-1
   MSFD@M2.bar.ages['last.age',1]<- 1     # cod age 0 and 1, and all ages for herring and sprat (species 2, 3)

   MSFD@do.community.F<-1
   MSFD@community.F.sp[]<-1   #include all species
   MSFD@do.community.M<-1
   MSFD@community.M.sp[]<-1   #include all species

   MSFD@do.life.expectancy <-1
   MSFD@life.expectancy.first.age[]<-0

   MSFD@do.community.life.expectancy<-1
   MSFD@community.life.expectancy.options["first.age",]<-0
   MSFD@community.life.expectancy.options["weighting",]<-1

   MSFD@do.LFI<-1
   MSFD@LFI.sp[]<-1
   MSFD@LFI.age[]<-c(3,99,99) # only cod age 3+ is greather than 40 cm
 }

 if (my.area=='North Sea') {
    # 1 Fulmar
    # 2 Guillemot
    # 3 Her. Gull
    # 4 Kittiwake
    # 5 GBB. Gull
    # 6 Gannet
    # 7 Puffin
    # 8 Razorbill
    # 9 R. radiata
    # 10 G. gurnards
    # 11 W. mackerel
    # 12 N. mackerel
    # 13 W.horse mac
    # 14 N.horse mac
    # 15 Grey seal
    # 16 H. porpoise
    # 17 Cod
    # 18 Whiting
    # 19 Haddock
    # 20 Saithe
    # 21 Herring
    # 22 Sandeel
    # 23 Nor. pout
    # 24 Sprat
    # 25 Plaice
    # 26 Sole

   MSFD@do.community.biomass.demersal<-1
   MSFD@community.biomass.demersal.ages['first.age' ,c(1:8,11:16,21,24)]<-  99  # exclude
   MSFD@community.biomass.demersal.ages['last.age'  ,c(1:8,11:16,21,24)]<-  -1  # exclude

   MSFD@do.community.biomass.small<-0  # not ready yet!!!!!!!!!!!!!!!!!!!!!!!
   MSFD@community.biomass.small.ages['last.age',c(1:8,11:16,21,24)]<- -1     # exclude    NEEDS UPDATE

   MSFD@do.community.biomass.pelagic<-1
   MSFD@community.biomass.pelagic.ages['first.age',c(1:20,22:23,25:26)]<- 99     # exclude
   MSFD@community.biomass.pelagic.ages['last.age',c(1:20,22:23,25:26)]<- -1     # exclude

   MSFD@do.community.biomass.forage<-1
   MSFD@community.biomass.forage.ages['first.age',c(1:20,25:26)]<- 99     # exclude
   MSFD@community.biomass.forage.ages['last.age',c(1:20,25:26)]<- -1     # exclude

#Average M2      age 0  age 1  age 2  age 3  age 4  age 5  age 6  age 7  age 8  age 9  age 10
#Cod             1.541  0.986  0.553  0.076  0.000  0.000  0.000  0.000  0.000  0.000  0.000
#Whiting         1.195  1.042  0.457  0.378  0.371  0.345  0.324  0.271  0.236
#Haddock         1.104  1.116  0.355  0.166  0.111  0.062  0.034  0.018  0.007  0.003  0.000
#Saithe          0.000  0.000  0.000  0.000  0.000  0.000  0.000  0.000  0.000  0.000  0.000
#Herring         0.551  0.463  0.248  0.212  0.187  0.169  0.159  0.134
#Sandeel         0.782  0.822  0.678  0.520  0.516
#Nor. pout       1.130  1.728  1.428  1.250
#Sprat           0.647  1.165  0.891  0.625

   MSFD@do.M2.bar<-1
   MSFD@M2.bar.ages['first.age',]<- c(0,0,0,-1,0,0,0,0,-1,-1)     # all ages for forage species, and >0.2 for other species
   MSFD@M2.bar.ages['last.age',] <- c(2,8,2,-1,3,4,3,3,-1,-1)     # all ages for forage species, and >0.2 for other species

   MSFD@do.community.mean.weight <-1
   MSFD@community.mean.weight.sp[]<-1  #include all species

   MSFD@do.community.mean.weight.C <-1
   MSFD@community.mean.weight.C.sp[]<-1  #include all species

   MSFD@do.mean.weight.C <-0
   MSFD@mean.weight.C.sp[]<-1    #include all species

   MSFD@do.community.F<-1
   MSFD@community.F.sp[]<-1   #include all species
   MSFD@do.community.M<-1
   MSFD@community.M.sp[]<-1   #include all species

   MSFD@do.life.expectancy <-1
   MSFD@life.expectancy.first.age[]<-0

   MSFD@do.community.life.expectancy<-1
   MSFD@community.life.expectancy.options["first.age",]<-0
   MSFD@community.life.expectancy.options["weighting",]<-1

   MSFD@do.LFI<-1
   MSFD@LFI.sp[ ,c(1:16,18,21:24)]<- 0  # exclude
   MSFD@LFI.age[]<-c(rep(99,16),2,99,5,2,99,99,99,99,10,8)  # first age considered as large fish (greather than 40 cm)
 }
}

write.FLOP.MSFD.control(MSFD,file="OP_MSFD.dat")

if (change.Other.pred.N) {
    othN<-scan(file=file.path(data.path,'OP_Other_N.in'),comment.char = "#",quiet = T )
    dim(othN)<-c(SMS@max.age.all+1,n.other.pred,SMS@last.season)
    dimnames(othN)=list( paste('a_',(0:SMS@max.age.all),sep=''),
                     sp.names[1:n.other.pred],
                     paste('q_',(1:SMS@last.season),sep=''))
}

######### end user options #########

if (do.simulation) {

   if (file.exists(scenario.dir) & ask.me) {
       cat("Are you sure you want to make new simultaion for scenario\n ",scenario,"\n The directory exist and will be overwritten!\n Do you want to continue? (Y/N):")
       a<-readLines(n=1)
       if(a=='N' || a=='n') stop("Script stopped'")
   }

  if (file.exists(scenario.dir)) unlink(scenario.dir,recursive = T)
  dir.create(scenario.dir,showWarnings = FALSE)

  cat(first.no.iter.stoch,'\n', file=file.path(scenario.dir,"OP_seed.in"))

  n.runs<-dim(targetFs)[[1]]
  cat(n.runs,'\n', file=file.path(scenario.dir,"OP_mulTargetF.in"))
  write.table(targetFs,row.names=F,col.names=F,file=file.path(scenario.dir,"OP_mulTargetF.in"),append=T)

  OP.files<-c("area_names.in","species_names.in","OP.dat","OP.exe","OP_trigger.dat","OP_config.dat","OP_MSFD.dat","just_one.in",
      "OP_consum.in","OP_F.in","OP_M1.in","OP_M.in","OP_N.in","OP_propmat.in","OP_prop_landed.in","OP_size.in","OP_wcatch.in","OP_wsea.in",
      "OP_growth_type1.in","OP_consum_ab.in","OP_other_N.in","OP_Exploitation.in","OP_reference_points.in","covariance_rec.in","OP_price.in",
      "OP_SSB_rec_residuals.in")

  if (SMS@no.areas>1)  OP.files<-c(OP.files,"OP_Stock_distribution.in")

  for (from.file in OP.files) {
    to.file<-file.path(scenario.dir,from.file)
    file.copy(from.file, to.file, overwrite = TRUE)
  }

  sp.name<-SMS.control@species.names

  setwd(scenario.dir)

  if (adjust.ini.N) {
   ini.N<-matrix( scan(file="OP_N.in",comment.char="#",quiet = T),nrow=n.vpa,byrow=T)
   ini.N[1,]<-ini.N[1,]*2  # cod
   write.matrix(ini.N,file="OP_N.in")
  }

  tot.run<-dim(targetFs)[[1]]
  fsp<-dim(targetFs)[[2]]  # number of species with variable F target

  if (ask.me) {
    cat("You have asked for",tot.run," runs with",no.iter,"iterations.\nDo you want to continue? (Y/N):")
    a<-readLines(n=1)
    if(a=='N') stop("Script stopped'")
    cat('OK\n')
  }
  MAMsp<-c("Grey seal","H. porpoise")
  MACsp<-c("W. mackerel","N. mackerel")
  BIRsp<-c("Fulmar","Guillemot","Her. Gull","Kittiwake","GBB. Gull","Gannet","Puffin","Razorbill")

  #run OP
  comm.sms<-paste( file.path(scenario.dir,"OP.exe"),"-maxfn 0 -nohess ",sep=" ")

  if (just.batch.file) {
    cat("cd ",scenario.dir,"\n",comm.sms,' >ud.dat\n',file=file.path(scenario.dir,"run.OP.bat"))
    #stop(paste("go to directory",scenario.dir,"and run script run_op.bat\n"))
    system(command=file.path(scenario.dir,"run.OP.bat"), wait = F)

  } else  shell(comm.sms, invisible = TRUE)
}   # end do.simulations

setwd(data.path)

if (!just.batch.file) {
  if (my.area=='North Sea') spNames<-c('COD', 'WHG', 'HAD', 'POK', 'HER', 'SAN', 'NOR', 'SPR', 'PLE', 'SOL')
  if (my.area=='Baltic Sea') spNames<-c('COD', 'HER', 'SPR')

  n.run<-dim(targetFs)[[1]]

  if (n.run >1) {
    if (read.condense) {
      cat('reading condensed data from directory:  ', scenario.dir,'\n')
      setwd(scenario.dir)
      condensed<-read.table('OP_average_val.out',header=F)
      allnames<-c(c('value','yield', 'CWsum', 'Fcomb','Fbar', 'SSB', 'TSB', 'recruit','belowBlim','belowBpa', 'Species.n', 'run','iteration'),spNames)
      dimnames(condensed)[[2]]<-allnames
      save(condensed, file =file.path(scenario.dir, "condensed.RData"))
      setwd(data.path)
    } else load(file =file.path(scenario.dir, "condensed.RData"))

   if (read.condense) {     # f mult combinations
      setwd(scenario.dir)
      nr<-scan( "OP_mulTargetF.in", nmax = 1,quiet = T)
      Fcomb<-data.frame(matrix(scan( "OP_mulTargetF.in",skip=1),nrow=nr,ncol=length(spNames),byrow=T))
      dimnames(Fcomb)[[2]]<-spNames
      Fcomb$run=first.run.no:(nr+first.run.no-1)
      save(Fcomb, file =file.path(scenario.dir, "Fcomb.RData"))
      setwd(data.path)
    } else load(file =file.path(scenario.dir, "Fcomb.RData"))

   if (read.indicators ) {
      setwd(scenario.dir)
      indi<-read.table( "OP_indicator_system_avg.out",header=T)

      load(file =file.path(scenario.dir, "Fcomb.RData"))
      #cat("dim Fcomb:\n")
      #print(dim(Fcomb))
      #print(head(Fcomb))
      #cat("dim ini\n")
      #print(dim(indi))
      #print(head(indi))
      indi<-merge(indi,Fcomb)
      #cat("dim indi after merge\n")
      #print(dim(indi))

      if  (my.area=='North Sea') {
        indi$unique.id<-paste(indi$run,indi$iter,round(indi$COD,3),round(indi$WHG,3),round(indi$HAD,3),round(indi$POK,3),round(indi$HER,3),round(indi$SAN,3),round(indi$NOR,3),round(indi$SPR,3),round(indi$PLE,3),round(indi$SOL,3),sep='')
      } else if  (my.area=='Baltic Sea') {
        indi$unique.id<-paste(indi$run,indi$iter,round(indi$COD,3),round(indi$HER,3),round(indi$SPR,3),sep='')
      }

      # get SSB and yield from the condensed data set
      if  (my.area=='North Sea') {
        b<-aggregate(cbind(value,yield,CWsum,SSB,TSB)~  run+iteration+COD+WHG+HAD+POK+HER+SAN+NOR+SPR+PLE+SOL,data=condensed,sum)
        b$unique.id<-paste(b$run,b$iteration,round(b$COD,3),round(b$WHG,3),round(b$HAD,3),round(b$POK,3),round(b$HER,3),round(b$SAN,3),round(b$NOR,3),round(b$SPR,3),round(b$PLE,3),round(b$SOL,3),sep='')
        b<-subset(b,select=c(unique.id,value,yield,CWsum,SSB,TSB))
        indi<-merge(indi,b,by="unique.id")
        indi$unique.id<-NULL
      }  else if  (my.area=='Baltic Sea') {
        b<-aggregate(cbind(value,yield,CWsum,SSB,TSB)~  run+iteration+COD+HER+SPR,data=condensed,sum)
        b$unique.id<-paste(b$run,b$iteration,round(b$COD,3),round(b$HER,3),round(b$SPR,3),sep='')
        b<-subset(b,select=c(unique.id,value,yield,CWsum,SSB,TSB))
        indi<-merge(indi,b,by="unique.id")
        indi$unique.id<-NULL
      }
      save(indi, file =file.path(scenario.dir, "indicators.RData"))
      setwd(data.path)
    } else load(file =file.path(scenario.dir, "indicators.RData"))

    if (combine.scenario.run) {
      all.condensed<-NULL
      for (sc in scenario.dirs) {
        load(file =file.path(data.path,sc, "condensed.RData"))  #load condensed data set
        cat(sc,dim(condensed),"\n")
        all.condensed<-rbind(all.condensed,condensed)
        rm(condensed)
      }
      condensed<-all.condensed
      rm(all.condensed)
      save(condensed, file =file.path(data.path,scenario.dirs.out, "condensed.RData"))

      # F combinations
      all.Fcomb<-NULL
      for (sc in scenario.dirs) {
        load(file =file.path(data.path,sc, "Fcomb.RData"))  #load condensed data set
        cat(sc,dim(Fcomb),"\n")
        all.Fcomb<-rbind(all.Fcomb,Fcomb)
        rm(Fcomb)
      }
      Fcomb<-all.Fcomb
      rm(all.Fcomb)
      save(Fcomb, file =file.path(data.path,scenario.dirs.out, "Fcomb.RData"))

      # indicators
      all.indi<-NULL
      for (sc in scenario.dirs) {
        load(file =file.path(data.path,sc, "indicators.RData"))  #load indi data set
        cat(sc,dim(indi),"\n")
        all.indi<-rbind(all.indi,indi)
        rm(indi)
      }
      indi<-all.indi
      rm(all.indi)
      save(indi, file =file.path(data.path,scenario.dirs.out, "indicators.RData"))
    }

    if (read.condense | combine.scenario.run) {
      source(file.path(prog.path,"HCR_OP_condense_function.r"))
      a<-transform.condensed(a=condensed,my.area=my.area)

      save(a, file =file.path(scenario.dir, "a.RData"));
    } else load(file =file.path(scenario.dir, "a.RData"))

  }  # end combine.scenario.run

  if ((do.plot.condense | ReduceTable )& n.run >1) {

    if (my.area=='North Sea') {
        if (subset.out) {       # make a subset of data set used to make plots
           #COD WHG HAD POK HER SAN NOR SPR PLE SOL
           bb<-subset.species
           a<-droplevels(subset(a, COD>=bb[1,1] &  COD<=bb[2,1] &
                                   WHG>=bb[1,2] &  WHG<=bb[2,2] &
                                   HAD>=bb[1,3] &  HAD<=bb[2,3] &
                                   POK>=bb[1,4] &  POK<=bb[2,4] &
                                   HER>=bb[1,5] &  HER<=bb[2,5] &
                                   SAN>=bb[1,6] &  SAN<=bb[2,6] &
                                   NOR>=bb[1,7] &  NOR<=bb[2,7] &
                                   SPR>=bb[1,8] &  SPR<=bb[2,8] &
                                   PLE>=bb[1,8] &  PLE<=bb[2,9] &
                                   SOL>=bb[1,10] & SOL<=bb[2,10] ))
         }
        out.file<-'all'
        out.types<-c('Yield','SSB','Recruit','Fbar','belowBlim', 'belowBpa','riskBlim')

        source(file.path(prog.path,"FMSY_matrix_functions_plots.r"))
        if (do.plot.condense) {
          plot.matrix(tit='Yield',type=out.types[1],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)
          plot.matrix(tit='SSB',  type=out.types[2],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)
          plot.matrix(tit='Recruits', type=out.types[3],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)
          plot.matrix(tit='Average F', type=out.types[4],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)
          plot.matrix(tit='Probability (%) of SSB below Blim', type=out.types[5],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)
          plot.matrix(tit='Proportion (%) of scenarios with prob(SSB<Blim) > 5%', type=out.types[7],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)

          cat("\n\nRelative MSY:\n")
        }
        table.MSY(a,out.file=out.file,scenario.dir=scenario.dir)
        cat('\ntable 01 done\n')
        cat('justAboveBlim:');
        cat(justAboveBlim,'\n')
        if (justAboveBlim & do.plot.condense) {
           cat('\nPlots, justAboveBlim\n')
           cat("\njustAboveBlim\n")

           dropKombBlim<-subset(a,probBelowBlim>riskLevels,select=c(run,iteration))

           if (dim(dropKombBlim)[[1]]>0) {
             dropKombBlim$drop<-T
             dropKombBlim<-unique(dropKombBlim)
             MV.dropKombBlim<<-dropKombBlim

             a<-merge(x=a,y=dropKombBlim,all.x=T)
             a[is.na(a$drop),'drop']<-F
             a<-droplevels(subset(a,drop==F,select=c(-drop)))
             if (dim(a)[[1]]>1)  {
               out.file='aboveBlim'
               out.types<-c('Yield','SSB','Recruit','Fbar','belowBlim', 'belowBpa')
              source(file.path(prog.path,"FMSY_matrix_functions.r"))
              plot.matrix(tit='Yield',type=out.types[1],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)
              plot.matrix(tit='SSB',  type=out.types[2],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)
              plot.matrix(tit='Recruits', type=out.types[3],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)
              plot.matrix(tit='Average F', type=out.types[4],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)
              plot.matrix(tit='Probability (%) of SSB below Blim', type=out.types[5],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)
              cat("\n\nRelative MSY:\n")
              table.MSY(a,out.file=out.file,scenario.dir=scenario.dir)
             } else cat("\n\nAll combinations below Blim, no plotting done\n\n")
           }
         }

        if (justAboveBpa & do.plot.condense) {
           cat('\nPlots, justAboveBlpa\n')
           a$probBelowBpa<-a$belowBpa/years.in.average
           dropKomb<-subset(a,probBelowBpa>0.05,select=c(run,iteration))

           if (dim(dropKomb)[[1]]>0) {
             dropKomb$drop<-T
             dropKomb<-unique(dropKomb)

             a<-merge(x=a,y=dropKomb,all.x=T)
             a[is.na(a$drop),'drop']<-F
             a<-droplevels(subset(a,drop==F,select=c(-drop)))
            if (dim(a)[[1]]>1)  {
               out.file='aboveBpa'
               out.types<-c('Yield','SSB','Recruit','Fbar','belowBlim', 'belowBpa')

              source(file.path(prog.path,"FMSY_matrix_functions.r"))
              plot.matrix(tit='Yield',type=out.types[1],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)
              plot.matrix(tit='SSB',  type=out.types[2],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)
              plot.matrix(tit='Recruits', type=out.types[3],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)
              plot.matrix(tit='Average F', type=out.types[4],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)
              plot.matrix(tit='Probability (%) of SSB below Blim', type=out.types[5],spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=a,first.sc=22)
              cat("\n\nRelative MSY:\n")
              table.MSY(a,out.file=out.file,scenario.dir=scenario.dir)
             } else cat("\n\nAll combinations below Bpa, no plotting done\n\n")
           }
         }

    }  # end North Sea

    if (my.area=='Baltic Sea') {
       source(file.path(prog.path,"FMSY-matrix-Baltic.R"))
    } # end Baltic Sea

    if (ReduceTable) {
           source(file.path(prog.path,"FMSY_matrix_functions_tables.r"))
           cat('\nStart to reduce Table\n')

          load(file =file.path(scenario.dir, "a.RData"))  # to get a fresh, non reduced version
           aa<-a
           goOn<<-T;
           ii<-0
           while (goOn) {
             a2<-reduceCombMSYfinal(aa,out.file='MSY_redu.out',app=(ii!=0),scenario.dir=scenario.dir,keepPro=keepPro);
             if (goOn) aa<-a2
             ii<-ii+1
             if (goOn & ReduceTableDetails) table.MSY(aa,out.file='MSY_redu_detail',scenario.dir=scenario.dir,app=(ii!=1))
           }
           cat("\n",dim(aa),"\n")
           table.MSY(aa,out.file="MSY_redu_final",scenario.dir=scenario.dir)
           rm(a2)
           #rm(aa)

           # risk to Blim
           # aa<-a          Do not start with the full datat set
           goOn<<-T;
           ii<-0
           while (goOn) {
             a2<-reduceCombRiskfinal(aa,out.file='MSY_and_Blim_redu_continue.out',app=(ii!=0),scenario.dir=scenario.dir,riskLevels=riskLevels);
             if (goOn) aa<-a2
             ii<-ii+1
             if (goOn & ReduceTableDetails) table.MSY(aa,out.file='MSY_and_Blim_redu_continue_details',scenario.dir=scenario.dir,app=(ii!=1))
           }
           table.MSY(aa,out.file='MSY_and_Blim_redu_continue_final',scenario.dir=scenario.dir)
           rm(a2)
           rm(aa)

           # risk to Blim
           aa<-a          # start with the full datat set
           goOn<<-T;
           ii<-0
           while (goOn) {
             a2<-reduceCombRiskfinal(aa,out.file='Blim_redu_all.out',app=(ii!=0),scenario.dir=scenario.dir,riskLevels=riskLevels);
             if (goOn) aa<-a2
             ii<-ii+1
             if (goOn & ReduceTableDetails) table.MSY(aa,out.file='Blim_redu_all_details',scenario.dir=scenario.dir,app=(ii!=1))
           }
           table.MSY(aa,out.file='Blim_redu_all_final',scenario.dir=scenario.dir)
           rm(a2)
           rm(aa)

           # risk to Blim , using the
           aa<-a          # start with the full datat set
           goOn<<-T;
           ii<-0
           while (goOn) {
             a2<-reduceCombRiskfinal(aa,out.file='Blim_redu_v2_all.out',app=(ii!=0),scenario.dir=scenario.dir,riskLevels=riskLevels,type=2);
             if (goOn) aa<-a2
             ii<-ii+1
             if (goOn & ReduceTableDetails) table.MSY(aa,out.file='Blim_redu_v2_all_details',scenario.dir=scenario.dir,app=(ii!=1))
           }
           table.MSY(aa,out.file='Blim_redu_v2_all_final',scenario.dir=scenario.dir)
           rm(a2)
           rm(aa)
        }

   }

  if (do.plot.indicators & n.run>1) {
      if (my.area=='North Sea') {
         source(file.path(prog.path,"plot_OP_community_indicators_average_NorthSea.R"))
          n.species<-8
          #a<-droplevels(subset(a,(COD<=0.4 & POK<=0.4) & SPR<0.6))
          out.file='all_'
          plot.matrix(tit='System indicators',spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=indi,first.sc=22)

         if (justAboveBlim) {
           cat("\njustAboveBlim Indicators\n")
           if (dim(dropKombBlim)[[1]]>0) {
             dropKombBlim$iter<-dropKombBlim$iteration
             indi<-merge(x=indi,y=dropKombBlim,all.x=T)
             MV.indi1<<-indi
             indi[is.na(indi$drop),'drop']<-F
             indi<-droplevels(subset(indi,drop==F,select=c(-drop)))
             MV.indi2<<-indi
             if (dim(indi)[[1]]>1)  {
               out.file='aboveBlim_'
               plot.matrix(tit='System indicators',spNames=sp.names[17:24],my.dev='png',out.file=out.file,scenario.dir=scenario.dir,a=indi,first.sc=22)
             }  else cat("\n\nAll combinations below Blim, no plotting done\n\n")
           }
         }
      }    # end North Sea

      if (my.area=='Baltic Sea') source(file.path(prog.path,"plot_OP_community_indicators_average_Baltic.R"))
  }
 } # end just batch

  used.scenario.dir<<-scenario.dir
  scenario.dir<<-used.scenario.dir
  scenario.dir.simul<<-used.scenario.dir

}  # end function

###################################
# function for fixing POK, PLE and SOL F
POK.PL.SOL<-function(HCR=HCR1,FBlim.adjust=1,stochastic.recruitment=1,recruit.adjust.CV=0) {
    # find F-values for POK, PLE and Sole
    OP.simulate(scenario="POK",HCR=HCR,FBlim.adjust=FBlim.adjust,targetFs=expand.grid(COD=seq(0.45,0.45,0.05),WHG=seq(0.30,0.30,0.05),HAD=seq(0.35,0.35,0.05),POK=seq(0.20,0.60,0.05),HER=seq(0.30,0.30,0.05),SAN=seq(0.25,0.25,0.05),NOR=seq(0.25,0.25,0.05),SPR=seq(0.30,0.30,0.05),PLE=seq(0.50,0.50,0.05),SOL=seq(0.50,0.50,0.05)),ask.me=ask.me,my.area=my.area,my.last.year=my.last.year, years.in.average= years.in.average,stochastic.recruitment=stochastic.recruitment,no.iter.stoch=no.iter.stoch,recruit.adjust.CV=recruit.adjust.CV)
    load(file =file.path(used.scenario.dir, "condensed.RData"))
    POK<-subset(condensed,Species.n==20)
    POK2<<-POK
    OP.simulate(scenario="PLE",HCR=HCR,FBlim.adjust=FBlim.adjust, targetFs=expand.grid(COD=seq(0.45,0.45,0.05),WHG=seq(0.30,0.30,0.05),HAD=seq(0.35,0.35,0.05),POK=seq(0.30,0.30,0.05),HER=seq(0.30,0.30,0.05),SAN=seq(0.25,0.25,0.05),NOR=seq(0.25,0.25,0.05),SPR=seq(0.30,0.30,0.05),PLE=seq(0.20,0.60,0.025),SOL=seq(0.50,0.50,0.05)),ask.me=ask.me,my.area=my.area,my.last.year=my.last.year, years.in.average= years.in.average,stochastic.recruitment=stochastic.recruitment,no.iter.stoch=no.iter.stoch,recruit.adjust.CV=recruit.adjust.CV)
    load(file =file.path(used.scenario.dir, "condensed.RData"))
    PLE<-subset(condensed,Species.n==25)

    OP.simulate(scenario="SOL",HCR=HCR,FBlim.adjust=FBlim.adjust, targetFs=expand.grid(COD=seq(0.45,0.45,0.05),WHG=seq(0.30,0.30,0.05),HAD=seq(0.35,0.35,0.05),POK=seq(0.30,0.30,0.05),HER=seq(0.30,0.30,0.05),SAN=seq(0.25,0.25,0.05),NOR=seq(0.25,0.25,0.05),SPR=seq(0.30,0.30,0.05),PLE=seq(0.50,0.50,0.025),SOL=seq(0.20,0.60,0.025)),ask.me=ask.me,my.area=my.area,my.last.year=my.last.year, years.in.average= years.in.average,stochastic.recruitment=stochastic.recruitment,no.iter.stoch=no.iter.stoch,recruit.adjust.CV=recruit.adjust.CV)
    load(file =file.path(used.scenario.dir, "condensed.RData"))
    SOL<-subset(condensed,Species.n==26)

    allSp<-rbind(POK,PLE,SOL)

    if (stochastic.recruitment==0  | stochastic.recruitment==2 ) {
      allSp[allSp$belowBlim>0,'belowBlim']<-100
      allSp[allSp$belowBpa>0,'belowBpa']<-100
    }
    png(filename =file.path(used.scenario.dir,paste('POK_PLE_SOL','.png',sep='')), width = 700, height = 1100,units = "px", pointsize = 18, bg="white")
    par(mfcol=c(6,3))
    par(mar=c(2,4,3,2)) # c(bottom, left, top, right)
    a<-subset(allSp,Species.n %in% c(20,25,26))

    boxs <- defmacro(a, varsp, spname,spno,
      expr={
          boxplot(yield/1000~varsp ,data=a,show.names = T,ylab='Yield (1000t)',subset=Species.n==spno,main=spname)
          boxplot(SSB/1000~varsp ,data=a,show.names = T,ylab='SSB (1000t)',subset=Species.n==spno)
          boxplot(Fbar~varsp ,data=a,show.names = T,ylab='Realized F',xlab='F',subset=Species.n==spno)
          boxplot(belowBlim~varsp ,data=a,show.names = T,ylab='prob(SSB<Blim) %',xlab='F',subset=Species.n==spno)
          boxplot(belowBpa~varsp ,data=a,show.names = T,ylab='prob(SSB<Bpa) %',xlab='F',subset=Species.n==spno)
          boxplot(recruit/1E6~varsp ,data=a,show.names = T,ylab='recruits (10^9)',xlab='F',subset=Species.n==spno)
        })
     boxs(a, POK, 'Saithe',20)
     boxs(a, PLE,'Plaice',25)
     boxs(a, SOL,'Sole',26)

    cleanup()
}


#assert existence of
#commonfile
source(commonfile)
#pwmfile
#load(file.path(pwmdir,paste0(pwmid,'.pwmout.RData')))
#tmpdir
#outdir

#####
# make bg
datadir = paste0(tmpdir,'/',pwmid,'/')

nsites = Reduce('+',lapply(list.files(datadir,paste0('positive.tf',pwmid,'-'),full.names=T),function(i){load(i);ncol(pos.mat)}))
if(nsites < 100){fast.mode=F}

load(file.path(datadir,paste0('background.tf',pwmid,'-',seqnames(genome)[1],'.RData')))

if( two.pass & file.exists(paste0(outdir,'/params/twopass.RData')) ){
    load(paste0(outdir,'/params/twopass.RData'))
    tmat.pos = cov.premade[1:(2*wsize),]
    tmat.neg = cov.premade[(2*wsize+1):(4*wsize),]
}

#use precomputed tables if wsize is 1000
if( (!fast.mode) & (wsize <= 1000)){
    load('precalc/cov.premade.RData')
    subsel = 1000+(-(wsize-1):wsize)
    tmat.pos = cov.premade[subsel,]
    tmat.neg = cov.premade[subsel + (2000),]
}

#require at least 30k sites to do a background estimate
if((ncol(pos.mat) > 30000) & (!fast.mode) & (wsize != 1000)){
source('covfun.r')
require(Matrix)
lf=list.files(datadir,paste0('background.tf',pwmid,'-'),full.names=T)
bgwsize=2*wsize
pmc=Reduce('+',lapply(lf,function(i){
    print(i)
    load(i)
    m=rBind(pos.mat[1:bgwsize,],neg.mat[1:bgwsize,])
    m%*%t(m)
}))
alldat=do.call(cBind,lapply(lf,function(i){
    load(i)
    rBind(pos.mat[1:bgwsize,],neg.mat[1:bgwsize,])
}))
datin=t(as.matrix(rBind(pos.mat[1:bgwsize,],neg.mat[1:bgwsize,])))
datin[datin>3]=3
stabval=4
mustab=-3
sz=ncol(datin)
bind=makeblocks(1:bgwsize,2)
tabs = makeTable(datin,mustab,stabval)
matin=matrix(tabs[[1]][as.double(datin)+1],nrow(datin))
covbase = (pmc-min(pmc))/(pmc[1,1]-min(pmc))
covin = covbase
sci=solve(covin)
muin= mustab
mupri=rep(muin,sz)%*%sci
scid = solve(sci+diag(rep(tabs[[2]][1],sz)))
for(m in 1:3){
    print(m)
    musum=rep(0,ncol(covin))
    devsum=0
    muevsum=0
    covsumT=matrix(0,ncol(covin),ncol(covin))
    sqmuT=matrix(0,ncol(covin),ncol(covin))
    muT=rep(0,ncol(covin))
    tct=0
    numex=min(nrow(datin),5000)
    for(i in sample(1:nrow(datin),numex)){
        print(i)
        ff=fastFit(datin[i,],mupri,tabs[[1]],tabs[[2]],scid)
        muall= t(ff[[1]])
        covsum = ff[[2]]
        devs = 1/(1/diag(ff[[2]])+tabs[[2]][datin[i,]+1])
        muevs= (ff[[1]]/diag(ff[[2]])-tabs[[1]][datin[i,]+1]*tabs[[2]][datin[i,]+1])*devs
        #
        sqmuT=sqmuT+t(muall)%*%muall
        muT=muT+colSums(muall)
        covsumT=covsumT+covsum
        devsum=devsum+sum(devs)
        muevsum=muevsum+sum(muevs)
        tct=tct+nrow(muall)
    }
    cch=sqmuT/tct-(t(t(muT))%*%t(muT))/(tct^2)
    cpost=(covsumT/tct+cch)
    muhat=muT/tct
    print(c(stabval,muin))
    tabs = makeTable(datin,muin,stabval)
    print(tabs)
    covin=cpost
    sci=solve(covin)
    mupri=muhat%*%sci
    scid = solve(sci+diag(rep(tabs[[2]][1],sz)))
}
svpos=svd(covin[1:(2*wsize),])
svneg=svd(covin[2*wsize+(1:(2*wsize)),])

ecsel = 1:ncol(svpos$u)
ecpv = svpos$d/max(svpos$d)*100
ecnv = svneg$d/max(svneg$d)*100
ecpv[ecpv < 1]=1
ptrans = t(t(svpos$u[,ecsel])*(sqrt(ecpv[ecsel])))
ecnv[ecnv < 1]=1
ntrans = t(t(svneg$u[,ecsel])*(sqrt(ecnv[ecsel])))
tmat.pos = cbind(ptrans, matrix(0,2*wsize,length(ecsel)))
tmat.neg = cbind(matrix(0,2*wsize,length(ecsel)), ntrans)
}else{
	rot.pos = diag(wsize*2)
	rot.neg = diag(wsize*2)
}

#####
# make svlite

converter.src <- '
Rcpp::NumericMatrix triplePos(rtriplePos);
Rcpp::NumericMatrix tripleNeg(rtripleNeg);
Rcpp::NumericVector sums(rsum);
int label = as<int>(rlabel);
int negoff = as<int>(rnegoff);
Rcpp::CharacterVector cv(sums.size());
int tripleindPos = 0;
int tripleindNeg = 0;
int len = 65536 * 10;
for(int i=0; i < sums.size(); i++){
  char str[len];
  double sumct = sums[i];
  int cp = snprintf(str, 65536, "%d %d:%f", label, 1, sumct);
  while(tripleindPos < triplePos.nrow() && triplePos(tripleindPos,1)-1 == i){
    int crd = triplePos(tripleindPos,0)+1;
    double wt = triplePos(tripleindPos,2);
    cp += snprintf(str + cp, len-cp, " %d:%f", crd, wt);
    tripleindPos++;
  }
  while(tripleindNeg < tripleNeg.nrow() && tripleNeg(tripleindNeg,1)-1 == i){
    int crd = tripleNeg(tripleindNeg,0)+1+negoff;
    double wt = tripleNeg(tripleindNeg,2);
    cp += snprintf(str + cp, len-cp, " %d:%f", crd, wt);
    tripleindNeg++;
  }
  cv[i] = str;
}
return cv;
'
converter <- cxxfunction(signature(rtriplePos="numeric",rtripleNeg="numeric",rsum="numeric",rnegoff="integer",rlabel="integer"),converter.src,plugin="Rcpp",includes="#include <stdio.h>")

sumtr <-function(x){
    x
}

makesvlite <- function(filename,label,rot.pos,rot.neg,minread=5){
    print(filename)
    load(filename)
    rct=sumtr(colSums(pos.mat)+colSums(neg.mat))+1
    pms=as.matrix(summary(pos.mat))
    pns=as.matrix(summary(neg.mat))
    cv=converter(pms,pns,rct,nrow(pos.mat),label)
    cv[rct>=(sumtr(minread)+1)]
}

if(fast.mode){

validpos = list.files(datadir,paste0('positive.tf',pwmid,'-'),full.names=T)
validneg = list.files(datadir,paste0('background.tf',pwmid,'-'),full.names=T)

allpos=do.call(c,lapply(validpos,makesvlite,label=1,minread=max(10*wsize/1000,1),rot.pos=rot.pos,rot.neg=rot.neg))
allneg=do.call(c,lapply(validneg,makesvlite,label= -1,minread= 0, rot.pos=rot.pos,rot.neg=rot.neg))
writeLines(c(allpos,allneg),file.path(tmpdir,paste0(pwmid,'-svlite.txt')))

itmax=min(5e+07,500*length(allpos))

wsz = 2*nrow(pos.mat)+2
sv.fit=sofia(file.path(tmpdir,paste0(pwmid,'-svlite.txt')),verbose=T,dimensionality=wsz,random_seed=1,lambda=1*10^6/(length(allpos)*wsz),iterations=itmax,learner_type='logreg-pegasos',eta_type='basic',loop_type='balanced-stochastic')

vpos=suppressMessages(as.vector((sv.fit$weights[2+(1:(nrow(pos.mat)))])))
vneg=suppressMessages(as.vector((sv.fit$weights[2+(1:(nrow(pos.mat)))+(nrow(pos.mat))])))
sv.rotate = c(sv.fit$weights[1:2],vpos,vneg)

save(sv.fit,sv.rotate,file=file.path(tmpdir,paste0(pwmid,'.svout.RData')))

}else{

validpos = list.files(datadir,paste0('positive.tf',pwmid,'-'),full.names=T)
validneg = list.files(datadir,paste0('background.tf',pwmid,'-'),full.names=T)

allnom = c('label','rct',paste0('x',1:ncol(tmat.pos)))
alldat=rbind(
    do.call(rbind,lapply(validpos,function(i){
        print(i)
        load(i)
        rct = sumtr(colSums(pos.mat) + colSums(neg.mat))+1
        y=data.frame(label=rep(1,sum(rct>(sumtr(10)+1))),rct[rct>(sumtr(10)+1)],x=t(as.matrix(t(tmat.pos)%*%pos.mat[,rct>(sumtr(10)+1)] + t(tmat.neg)%*%neg.mat[,rct>(sumtr(10)+1)])))
        colnames(y)=allnom
        y
    })),
    do.call(rbind,lapply(validneg,function(i){
        print(i)
        load(i)
        rct = sumtr(colSums(pos.mat) + colSums(neg.mat))+1
        y=data.frame(label=-1,rct,x=t(as.matrix(t(tmat.pos)%*%pos.mat + t(tmat.neg)%*%neg.mat)))
        colnames(y)=allnom
        y
})))

write.svmlight(alldat[,1],as.matrix(alldat[,-1]),file.path(tmpdir,paste0(pwmid,'-svlite.txt')))

itmax=min(5e+07,1000*sum(alldat[,1]==1))

sv.fit=sofia(file.path(tmpdir,paste0(pwmid,'-svlite.txt')),verbose=T,dimensionality=ncol(alldat),random_seed=1,lambda=1000/sum(alldat[,1]==1),iterations=itmax,learner_type='logreg-pegasos',eta_type='basic',loop_type='balanced-stochastic')

sv.rotate=c(sv.fit$weights[1:2],tmat.pos%*%(sv.fit$weight[-(1:2)]),tmat.neg%*%(sv.fit$weight[-(1:2)]))

save(sv.fit,sv.rotate,file=file.path(tmpdir,paste0(pwmid,'.svout.RData')))

}

#
#####


library(cummeRbund)
args <- commandArgs(trailingOnly=TRUE)
error_file <- file(args[1], open="wt")
sink(error_file, type="message")
cuff <- readCufflinks()
genes.scv<-fpkmSCVPlot(genes(cuff))
png(args[2],width = 10, height = 10, units= 'in',res=600)
print(genes.scv)
dev.off()


library(RQuantLib)
library(stringr)

d0 <- as.Date("1960-01-04")
d1 <- seq(d0, as.Date("2100-01-04"), by="days")
l <- length(d1)
d0 <- rep(d0, l)

cals = c( "Canada", "Canada/Settlement", "Canada/TSX", "Germany", "Germany/FrankfurtStockExchange",
          "Germany/Settlement", "Germany/Xetra", "Germany/Eurex", "Italy", "Italy/Settlement", "Italy/Exchange",
          "Japan", "UnitedKingdom", "UnitedKingdom/Settlement", "UnitedKingdom/Exchange", "UnitedKingdom/Metals",
          "UnitedStates", "UnitedStates/Settlement", "UnitedStates/NYSE", "UnitedStates/GovernmentBond",
          "UnitedStates/NERC", "Brazil", "TARGET")

dir.create('csv')

for (cal in cals) {
  res <- isBusinessDay(cal, d1)
  df <- data.frame(d=d1, isbday=res)
  write.csv(df, file=paste0( "csv/QuantLib-isbday-",   str_replace(cal, "/", "-")   , ".csv"), row.names=FALSE)
}


library(ggplot2)
library(plyr)
theme_set(theme_bw(18))
setwd("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/20_sinking-marbles-prior-sliders/results/")
source("rscripts/helpers.r")

load("data/r.RData")
load("../../1_sinking-marbles-prior/results/data/priors.RData")
r = read.csv("data/sinking_marbles.csv", header=T)
nrow(r)
head(r)
r$trial = r$slide_number_in_experiment - 2
r = r[,c("assignmentid","workerid", "rt", "effect", "cause","language","gender.1","age","gender","other_gender", "object_level", "response", "object","slider_id","num_objects","trial","enjoyment","asses","comments","Answer.time_in_minutes","slider_id")]
r$object_level = factor(r$object_level, levels=c("object_high", "object_mid", "object_low"))
r$Item = as.factor(paste(r$effect,r$object))
table(r$Item)

## add old priors to data.frame
priorprobs = read.table(file="~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/12_sinking-marbles-prior15/results/data/smoothed_15marbles_priors_withnames.txt",sep="\t", header=T, quote="")
head(priorprobs)
row.names(priorprobs) = paste(priorprobs$effect,priorprobs$object)
mpriorprobs = melt(priorprobs, id.vars=c("effect", "object"))
head(mpriorprobs)
row.names(mpriorprobs)= paste(mpriorprobs$variable, mpriorprobs$effect,mpriorprobs$object)
head(mpriorprobs)

r$PriorProbability = mpriorprobs[paste(paste("X",r$slider_id,sep=""),as.character(r$Item), sep=" "),]$value

# compute normalized probabilities
nr = ddply(r, c("workerid","trial"), summarise,
           sumresponse=sum(response))
row.names(nr) = paste(nr$workerid, nr$trial)

r$normresponse = r$response/nr[paste(r$workerid, r$trial),]$sumresponse

# test: sums should add to 1
sums = ddply(r, .(workerid,trial), summarise, sum(normresponse))
colnames(sums) = c("workerid","trial","sum")
summary(sums)
sums[is.na(sums$sum),]

save(r, file="data/r.RData")

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

ggplot(aes(x=gender.1), data=r) +
  geom_histogram()

ggplot(aes(x=rt), data=r) +
  geom_histogram() +
  scale_x_continuous(limits=c(0,50000))

ggplot(aes(x=age), data=r) +
  geom_histogram()

ggplot(aes(x=age,fill=gender.1), data=r) +
  geom_histogram()

ggplot(aes(x=enjoyment), data=r) +
  geom_histogram()

ggplot(aes(x=asses), data=r) +
  geom_histogram()

ggplot(aes(x=Answer.time_in_minutes), data=r) +
  geom_histogram()

ggplot(aes(x=age,y=Answer.time_in_minutes,color=gender.1), data=unique(r[,c("assignmentid","age","Answer.time_in_minutes","gender.1")])) +
  geom_point() +
  geom_smooth()

unique(r$comments)


suppressMessages(library(pbdMPI))
suppressMessages(library(tasktools))

costly = function(x, waittime)
{
  Sys.sleep(waittime)
  cat(paste("iter", x, "executed on rank", comm.rank(), "\n"))

  sqrt(x)
}

ret = mpi_napply(10, costly, checkpoint_path="/tmp", preschedule=FALSE, waittime=0.5)
comm.cat("\n", quiet=TRUE)
comm.print(unlist(ret))

finalize()


library(tidyverse)
library(cowplot)
library (brms)
options (mc.cores=parallel::detectCores ()) # Run on multiple cores
# library(nnet)
# devtools::install_github('michael-franke/bayes_mixed_regression_tutorial/faintr', build_vignettes = TRUE)
library(faintr)

d = read_csv("../../data/06_sentence_select_with_training/data_raw_final.csv")

# check comments
d$comments %>% unique

# participants (gender, age...)
d %>% select(submission_id, gender) %>% unique() %>% group_by(gender) %>% summarize(count = n())
mean(d$age)
sd(d$age)

##########################################
# allocation to conditions
##########################################

message("Nr. of submissions for each co-player type:")
show(table(d$coplayer_type)/122)

##########################################
## color blindness
##########################################

color_blindness = filter(d, trial_type == "color_blindness_test") %>%
  select(submission_id, correct, response) %>%
  mutate(correct = ifelse(correct == response,1,0)) %>%
  group_by(submission_id) %>%
  summarize(color_blind = mean(correct) < 1)
color_blindness_failures = filter(color_blindness, color_blind == TRUE) %>% pull(submission_id)
message("Number of divergences on color blindness test: ", length(color_blindness_failures))
message("Inspect these to determine gravity of divergence!")

# add info to full data set
d = full_join(d, color_blindness, by = "submission_id")

# inspect colorblind failures for gravity:

filter(d, color_blind == TRUE, trial_type == "color_blindness_test") %>%
  select(submission_id, correct, response)

## we do not exlcude anybody for a severe color-blindness

##########################################
# inspect data from training with feedback
##########################################

success_threshold = 0.5

d_training = filter(d, trial_type == "sentence_completion_training")

d_training_summary = d_training %>%
  mutate(winning_move = ifelse(coplayer_type == "unstrategic", "red", "green"),
         winning_move_chosen = ifelse(winning_move == response, 1, 0),
         loosing_move = ifelse(coplayer_type == "unstrategic", "green", "red"),
         loosing_move_chosen = ifelse(loosing_move == response, 1, 0),
         distractor_move_chosen = ifelse(response == "false", 1, 0)) %>%
  group_by(submission_id, coplayer_type, condition) %>%
  summarize(win = mean(winning_move_chosen),
            loose = mean(loosing_move_chosen),
            dstrct = mean(distractor_move_chosen)) %>%
  ungroup() %>%
  gather(key = response, value = proportion, win, loose, dstrct) %>%
  mutate(new_col_names = paste0("train_", condition, "_", response)) %>%
  select(submission_id, coplayer_type, new_col_names, proportion) %>%
  spread(key = new_col_names, value = proportion) %>%
  group_by(submission_id) %>%
  mutate(
    training_score = mean(c(train_all_win, train_none_win)),
    training_successful = ifelse(training_score > success_threshold, 1, 0)
  ) %>%
  ungroup()

# show distribution of training scores (average sucess)
d_training_summary %>% ggplot(aes(x = training_score)) + geom_histogram() + facet_grid( ~coplayer_type)

d_training_summary %>%
  group_by(coplayer_type) %>%
  summarize(mean_score = mean(training_score))

d = full_join(d, select(d_training_summary, - coplayer_type), by = "submission_id")

message("Nr learned from training?")
d %>% select(submission_id, coplayer_type, training_successful) %>%
  unique() %>%
  group_by(coplayer_type, training_successful) %>%
  summarize(count = n()) %>% show()

##########################################
# inspect data from TVJ task
##########################################

control_threshold = 0.7
semprag_threshold = 0.5

d_tvj = filter(d, trial_type == "truth_value_judgements") %>%
  mutate(response = ifelse(response == "true", 1, 0))

d_tvj_summary = d_tvj %>%
  group_by(submission_id, trigger, condition) %>%
  summarize(mean_true = mean(response)) %>%
  arrange(submission_id) %>%
  ungroup() %>%
  mutate(new_col_names = paste0("tvj_", trigger, "_", condition)) %>%
  select(submission_id, new_col_names, mean_true) %>%
  spread(key = "new_col_names", value = mean_true)

d_tvj_controls = d_tvj %>%
  filter( condition != "implicature" ) %>%
  group_by(submission_id, trigger, condition) %>%
  summarize(mean_true = mean(response)) %>%
  arrange(submission_id) %>%
  ungroup() %>%
  mutate(control_score = ifelse(condition == "true", mean_true, 1-mean_true)) %>%
  group_by(submission_id) %>%
  summarize(tvj_control_mean = mean(control_score),
            tvj_control_fail = ifelse(tvj_control_mean < control_threshold, 1, 0))

message("Number of people who performed really badly in the TVJ: ",
        d_tvj_controls$tvj_control_fail %>% sum())

d_tvj_implicatures = d_tvj %>%
  filter( condition == "implicature", trigger != "number" ) %>%
  group_by(submission_id, trigger, condition) %>%
  summarize(mean_true = mean(response)) %>%
  group_by(submission_id) %>%
  summarize(tvj_implicature_mean = mean(mean_true),
            tvj_semprag_type = ifelse(tvj_implicature_mean <= semprag_threshold,
                                      "pragmatic", "semantic"))

message("Number of participants who expected co-player to interpret pragmatically in TVJ: ",
        sum(d_tvj_implicatures$tvj_semprag_type == "pragmatic"))

d = full_join(d, d_tvj_summary, by = "submission_id")
d = full_join(d, d_tvj_controls, by = "submission_id")
d = full_join(d, d_tvj_implicatures, by = "submission_id")

# exclude bad performers
d = filter(d, tvj_control_fail == 0)

##########################################
# looking at how many combinations we got
##########################################

d %>% filter(tvj_control_fail == 0) %>%
  select(submission_id,
         coplayer_type,
         training_successful,
         tvj_semprag_type,
         tvj_control_fail) %>%
  unique() %>%
  group_by(coplayer_type, training_successful, tvj_semprag_type) %>%
  summarize(count = n()) %>% show()

# e = d %>% filter(tvj_control_fail == 0) %>%
#   select(submission_id,
#          coplayer_type,
#          training_successful,
#          tvj_semprag_type,
#          tvj_control_fail) %>%
#   unique()
# View(e)

##########################################
# prepare subset of data for main analysis
##########################################

# filter and select the relevant bits for analysis
# code the responses as win/lose/opt_out
d_analysis = filter(d, trial_type == "sentence_completion") %>%
  mutate(
    rsp = relevel(factor(
      ifelse(coplayer_type == "unstrategic",
             case_when(response == "red" ~ "win",
                       response == "green" ~ "lose",
                       TRUE ~ "opt_out"),
             case_when(response == "green" ~ "win",
                       response == "red" ~ "lose",
                       TRUE ~ "opt_out"))),
      ref = "win"),
    condition = relevel(factor(condition), ref = "all"),
    item = number,
    win_move = ifelse(rsp == "win", 1, 0)
  ) %>%
  select(submission_id, item, coplayer_type, condition,
         training_score, training_successful,
         tvj_implicature_mean, tvj_semprag_type,
         response, rsp, win_move)


## * Assign number of children to each edge (from "to" node)
nchild2edges <- function(tedges, from.node = NULL) {
    if (is.null(from.node))
        from.node <- setdiff(tedges$from, tedges$to)
    tedges[, nchildren := NA_integer_]

    children <- tedges[from %chin% from.node]
    while(nrow(children)) {
        nchild <- children[, .N, from]
        tedges[nchild, nchildren := i.N, on = c('to' = 'from')]
        children <- tedges[from %chin% children$to]
    }
    return(invisible(tedges))
}

get.root <- function(tedges) {
    setdiff(tedges$from, tedges$to)
}

## * Assign the distance to root to each edge
dist2root <- function(tedges, from.node = NULL) {
    stopifnot(all(c('dist', 'from', 'to') %in% names(tedges)))

    if (is.null(from.node))
        from.node <- setdiff(tedges$from, tedges$to)

    if (is.null(key(tedges)) || key(tedges) != 'from')
        setkey(tedges, from)

    tedges[, cumdist := 0]
    i <- 1L
    children <- tedges[from %in% from.node]
    while(nrow(children)) {
        tedges[from %in% children$from, `:=`(cumdist = cumdist + dist, nthnode = i)]
        prevchildren <- tedges[from %in% children$from]
        children <- tedges[from %in% children$to]
        tedges[prevchildren, cumdist := i.cumdist, on = c('from' = 'to')]
        i <- i + 1L
    }
    return(invisible(tedges))

}

children_unique_vals <- function(tedges, var) {

    tedges[, cladevar := get(var)]
    tedges[, cladestart := F]
    ## if (is.null(key(tedges)) || key(tedges) != 'from')
    ##     setkey(tedges, from)

    naval <- fcase(
        tedges[, is.character(get(var))], NA_character_,
        tedges[, is.integer(get(var))], NA_integer_,
        tedges[, is.numeric(get(var))], NA_real_,
        default = NA
    )

    getun <- function(x) {
        vs <- unique(setdiff(x, NA))
        if (length(vs) == 1L) {
            return(vs)
        } else return('multiple')
    }

    tedges[, uniquevals_children := get(var)]

    sel <- tedges[nthnode == max(nthnode)]
    uns <- sel[, .(uniquevals_children = getun(uniquevals_children)), from]
    tedges[uns, uniquevals_children := i.uniquevals_children, on = c('to' = 'from')]

    for (nn in (max(tedges$nthnode)-1):1L) {
        sel <- tedges[nthnode == nn]
        uns <- sel[, .(uniquevals_children = getun(uniquevals_children)), from]
        sel[uns, uniquevals_children_new := i.uniquevals_children, on = 'from']
        cladestarts <- sel[uniquevals_children_new == 'multiple' & uniquevals_children != 'multiple']
        tedges[cladestarts, cladestart := T, on = c('from' = 'to')]
        tedges[uns, uniquevals_children := i.uniquevals_children, on = c('to' = 'from')]
    }

}

## * Create edges data.table from phylo tree, and join tips data
get.edges <- function(tree) {
    tedges <- as.data.table(tree$edge)
    setnames(tedges, c('from', 'to'))
    tedges[, eid := .I]
    tedges[, dist := tree$edge.length]
    ttips <- get.tips(tree)
    tedges[, to.tip := V2 %in% ttips$id]
    tedges <- merge(tedges, ttips, by.x = 'V2', by.y = 'id', all.x = T, all.y = F, sort = F)
    setnames(tedges, c('V1', 'V2'), c('from', 'to'))
    setcolorder(tedges, c('from', 'to'))
    nchild2edges(tedges)
    dist2root(tedges)
    tedges
}

## * Extract clade of same trait, given a tip id, from edges data.table
common.trait.clade <- function(tedges, tip, var) {

    i <- 0L
    anc <- tedges[to == tip]
    val <- anc[, get(var)]
    if (is.na(val))
        stop('Need a tip with a non-NA value to start from')

    anc[, lvl := i]
    res <- anc

    sibs <- get.children(tedges, anc$from)
    cont <- sibs[, all(get(var) %in% c(NA, val))]

    while(cont) {
        sibs <- sibs[!res, on = 'to']
        sibs[, lvl := i]
        res <- rbind(res, sibs)
        i <- i + 1L
        anc <- tedges[to == anc$from]
        sibs <- get.children(tedges, anc$from)
        cont <- sibs[, all(get(var) %in% c(NA, val))]
    }

    return(res)
}

update.tree <- function(tedges) {
    tedges <- nchild2edges(tedges)
    tedges <- dist2root(tedges)
    tedges[, to.tip := fifelse(is.na(nchildren), T, F)]
    return(invisible(tedges))
}

## x=tree
## type = "fan"; use.edge.length = TRUE; node.pos = NULL;
## show.tip.label = FALSE; show.node.label = FALSE; edge.color = "grey50";
## edge.width = 1; edge.lty = 1; font = 3; cex = par("cex");
## adj = NULL; srt = 0; no.margin = TRUE; root.edge = FALSE;
## label.offset = 0; underscore = FALSE; x.lim = NULL; y.lim = NULL;
## direction = "rightwards"; lab4ut = NULL; tip.color = "red";
## plot = TRUE; rotate.tree = 0; open.angle = 0; node.depth = 1;
## align.tip.label = FALSE
## * Adaptation of ape::plot.phylo() to return the coordinates
plot2 <- function (x, type = "phylogram", use.edge.length = TRUE, node.pos = NULL,
    show.tip.label = TRUE, show.node.label = FALSE, edge.color = "black",
    edge.width = 1, edge.lty = 1, font = 3, cex = par("cex"),
    adj = NULL, srt = 0, no.margin = FALSE, root.edge = FALSE,
    label.offset = 0, underscore = FALSE, x.lim = NULL, y.lim = NULL,
    direction = "rightwards", lab4ut = NULL, tip.color = "black",
    plot = TRUE, rotate.tree = 0, open.angle = 0, node.depth = 1,
    align.tip.label = FALSE, ...) {

    Ntip <- length(x$tip.label)

    if (Ntip < 2) {
        warning("found less than 2 tips in the tree")
        return(NULL)
    }
    .nodeHeight <- function(edge, Nedge, yy) .C(node_height,
        as.integer(edge[, 1]), as.integer(edge[, 2]), as.integer(Nedge),
        as.double(yy))[[4]]
    .nodeDepth <- function(Ntip, Nnode, edge, Nedge, node.depth) .C(node_depth,
        as.integer(Ntip), as.integer(edge[, 1]), as.integer(edge[,
            2]), as.integer(Nedge), double(Ntip + Nnode), as.integer(node.depth))[[5]]
    .nodeDepthEdgelength <- function(Ntip, Nnode, edge, Nedge,
        edge.length) .C(node_depth_edgelength, as.integer(edge[,
        1]), as.integer(edge[, 2]), as.integer(Nedge), as.double(edge.length),
        double(Ntip + Nnode))[[5]]
    Nedge <- dim(x$edge)[1]
    Nnode <- x$Nnode
    if (any(x$edge < 1) || any(x$edge > Ntip + Nnode))
        stop("tree badly conformed; cannot plot. Check the edge matrix.")
    ROOT <- Ntip + 1
    type <- match.arg(type, c("phylogram", "cladogram", "fan",
        "unrooted", "radial"))
    direction <- match.arg(direction, c("rightwards", "leftwards",
        "upwards", "downwards"))
    if (is.null(x$edge.length)) {
        use.edge.length <- FALSE
    } else {
        if (use.edge.length && type != "radial") {
            tmp <- sum(is.na(x$edge.length))
            if (tmp) {
                warning(paste(tmp, "branch length(s) NA(s): branch lengths ignored in the plot"))
                use.edge.length <- FALSE
            }
        }
    }
    if (is.numeric(align.tip.label)) {
        align.tip.label.lty <- align.tip.label
        align.tip.label <- TRUE
    } else {
        if (align.tip.label)
            align.tip.label.lty <- 3
    }
    if (align.tip.label) {
        if (type %in% c("unrooted", "radial") || !use.edge.length ||
            is.ultrametric(x))
            align.tip.label <- FALSE
    }
    if (type %in% c("unrooted", "radial") || !use.edge.length ||
        is.null(x$root.edge) || !x$root.edge)
        root.edge <- FALSE
    phyloORclado <- type %in% c("phylogram", "cladogram")
    horizontal <- direction %in% c("rightwards", "leftwards")
    xe <- x$edge
    if (phyloORclado) {
        phyOrder <- attr(x, "order")
        if (is.null(phyOrder) || phyOrder != "cladewise") {
            x <- reorder(x)
            if (!identical(x$edge, xe)) {
                ereorder <- match(x$edge[, 2], xe[, 2])
                if (length(edge.color) > 1) {
                  edge.color <- rep(edge.color, length.out = Nedge)
                  edge.color <- edge.color[ereorder]
                }
                if (length(edge.width) > 1) {
                  edge.width <- rep(edge.width, length.out = Nedge)
                  edge.width <- edge.width[ereorder]
                }
                if (length(edge.lty) > 1) {
                  edge.lty <- rep(edge.lty, length.out = Nedge)
                  edge.lty <- edge.lty[ereorder]
                }
            }
        }
        yy <- numeric(Ntip + Nnode)
        TIPS <- x$edge[x$edge[, 2] <= Ntip, 2]
        yy[TIPS] <- 1:Ntip
    }
    z <- reorder(x, order = "postorder")
    if (phyloORclado) {
        if (is.null(node.pos))
            node.pos <- if (type == "cladogram" && !use.edge.length)
                2
            else 1
        if (node.pos == 1)
            yy <- .nodeHeight(z$edge, Nedge, yy)
        else {
            ans <- .C(node_height_clado, as.integer(Ntip), as.integer(z$edge[,
                1]), as.integer(z$edge[, 2]), as.integer(Nedge),
                double(Ntip + Nnode), as.double(yy))
            xx <- ans[[5]] - 1
            yy <- ans[[6]]
        }
        if (!use.edge.length) {
            if (node.pos != 2)
                xx <- .nodeDepth(Ntip, Nnode, z$edge, Nedge,
                  node.depth) - 1
            xx <- max(xx) - xx
        }
        else {
            xx <- .nodeDepthEdgelength(Ntip, Nnode, z$edge, Nedge,
                z$edge.length)
        }
    } else {
        twopi <- 2 * pi
        rotate.tree <- twopi * rotate.tree/360
        if (type != "unrooted") {
            TIPS <- x$edge[which(x$edge[, 2] <= Ntip), 2]
            xx <- seq(0, twopi * (1 - 1/Ntip) - twopi * open.angle/360,
                length.out = Ntip)
            theta <- double(Ntip)
            theta[TIPS] <- xx
            theta <- c(theta, numeric(Nnode))
        }
        switch(type, fan = {
            theta <- .nodeHeight(z$edge, Nedge, theta)
            if (use.edge.length) {
                r <- .nodeDepthEdgelength(Ntip, Nnode, z$edge,
                  Nedge, z$edge.length)
            } else {
                r <- .nodeDepth(Ntip, Nnode, z$edge, Nedge, node.depth)
                r <- 1/r
            }
            theta <- theta + rotate.tree
            if (root.edge) r <- r + x$root.edge
            xx <- r * cos(theta)
            yy <- r * sin(theta)
        }, unrooted = {
            nb.sp <- .nodeDepth(Ntip, Nnode, z$edge, Nedge, node.depth)
            XY <- if (use.edge.length) unrooted.xy(Ntip, Nnode,
                z$edge, z$edge.length, nb.sp, rotate.tree) else unrooted.xy(Ntip,
                Nnode, z$edge, rep(1, Nedge), nb.sp, rotate.tree)
            xx <- XY$M[, 1] - min(XY$M[, 1])
            yy <- XY$M[, 2] - min(XY$M[, 2])
        }, radial = {
            r <- .nodeDepth(Ntip, Nnode, z$edge, Nedge, node.depth)
            r[r == 1] <- 0
            r <- 1 - r/Ntip
            theta <- .nodeHeight(z$edge, Nedge, theta) + rotate.tree
            xx <- r * cos(theta)
            yy <- r * sin(theta)
        })
    }
    if (phyloORclado) {
        if (!horizontal) {
            tmp <- yy
            yy <- xx
            xx <- tmp - min(tmp) + 1
        }
        if (root.edge) {
            if (direction == "rightwards")
                xx <- xx + x$root.edge
            if (direction == "upwards")
                yy <- yy + x$root.edge
        }
    }
    if (no.margin)
        par(mai = rep(0, 4))
    if (show.tip.label)
        nchar.tip.label <- nchar(x$tip.label)
    max.yy <- max(yy)
    getLimit <- function(x, lab, sin, cex) {
        s <- strwidth(lab, "inches", cex = cex)
        if (any(s > sin))
            return(1.5 * max(x))
        Limit <- 0
        while (any(x > Limit)) {
            i <- which.max(x)
            alp <- x[i]/(sin - s[i])
            Limit <- x[i] + alp * s[i]
            x <- x + alp * s
        }
        Limit
    }
    if (is.null(x.lim)) {
        if (phyloORclado) {
            if (horizontal) {
                xx.tips <- xx[1:Ntip]
                if (show.tip.label) {
                  pin1 <- par("pin")[1]
                  tmp <- getLimit(xx.tips, x$tip.label, pin1,
                    cex)
                  tmp <- tmp + label.offset
                }
                else tmp <- max(xx.tips)
                x.lim <- c(0, tmp)
            }
            else x.lim <- c(1, Ntip)
        } else switch(type, fan = {
            if (show.tip.label) {
                offset <- max(nchar.tip.label * 0.018 * max.yy *
                  cex)
                x.lim <- range(xx) + c(-offset, offset)
            } else x.lim <- range(xx)
        }, unrooted = {
            if (show.tip.label) {
                offset <- max(nchar.tip.label * 0.018 * max.yy *
                  cex)
                x.lim <- c(0 - offset, max(xx) + offset)
            } else x.lim <- c(0, max(xx))
        }, radial = {
            if (show.tip.label) {
                offset <- max(nchar.tip.label * 0.03 * cex)
                x.lim <- c(-1 - offset, 1 + offset)
            } else x.lim <- c(-1, 1)
        })
    } else if (length(x.lim) == 1) {
        x.lim <- c(0, x.lim)
        if (phyloORclado && !horizontal)
            x.lim[1] <- 1
        if (type %in% c("fan", "unrooted") && show.tip.label)
            x.lim[1] <- -max(nchar.tip.label * 0.018 * max.yy *
                cex)
        if (type == "radial")
            x.lim[1] <- if (show.tip.label)
                -1 - max(nchar.tip.label * 0.03 * cex)
            else -1
    }
    if (phyloORclado && direction == "leftwards")
        xx <- x.lim[2] - xx
    if (is.null(y.lim)) {
        if (phyloORclado) {
            if (horizontal)
                y.lim <- c(1, Ntip)
            else {
                pin2 <- par("pin")[2]
                yy.tips <- yy[1:Ntip]
                if (show.tip.label) {
                  tmp <- getLimit(yy.tips, x$tip.label, pin2,
                    cex)
                  tmp <- tmp + label.offset
                }
                else tmp <- max(yy.tips)
                y.lim <- c(0, tmp)
            }
        } else switch(type, fan = {
            if (show.tip.label) {
                offset <- max(nchar.tip.label * 0.018 * max.yy *
                  cex)
                y.lim <- c(min(yy) - offset, max.yy + offset)
            } else y.lim <- c(min(yy), max.yy)
        }, unrooted = {
            if (show.tip.label) {
                offset <- max(nchar.tip.label * 0.018 * max.yy *
                  cex)
                y.lim <- c(0 - offset, max.yy + offset)
            } else y.lim <- c(0, max.yy)
        }, radial = {
            if (show.tip.label) {
                offset <- max(nchar.tip.label * 0.03 * cex)
                y.lim <- c(-1 - offset, 1 + offset)
            } else y.lim <- c(-1, 1)
        })
    } else if (length(y.lim) == 1) {
        y.lim <- c(0, y.lim)
        if (phyloORclado && horizontal)
            y.lim[1] <- 1
        if (type %in% c("fan", "unrooted") && show.tip.label)
            y.lim[1] <- -max(nchar.tip.label * 0.018 * max.yy *
                             cex)
        if (type == "radial")
            y.lim[1] <- if (show.tip.label)
                            -1 - max(nchar.tip.label * 0.018 * max.yy * cex)
                        else -1
    }
    if (phyloORclado && direction == "downwards")
        yy <- y.lim[2] - yy
    if (phyloORclado && root.edge) {
        if (direction == "leftwards")
            x.lim[2] <- x.lim[2] + x$root.edge
        if (direction == "downwards")
            y.lim[2] <- y.lim[2] + x$root.edge
    }
    asp <- fifelse(type %in% c("fan", "radial", "unrooted"), 1, NA_integer_)
    plot.default(0, type = "n", xlim = x.lim, ylim = y.lim, xlab = "",
        ylab = "", axes = FALSE, asp = asp, ...)
    if (plot) {
        if (is.null(adj))
            adj <- fifelse(phyloORclado && direction == "leftwards", 1, 0)
        if (phyloORclado && show.tip.label) {
            MAXSTRING <- max(strwidth(x$tip.label, cex = cex))
            loy <- 0
            if (direction == "rightwards") {
                lox <- label.offset + MAXSTRING * 1.05 * adj
            }
            if (direction == "leftwards") {
                lox <- -label.offset - MAXSTRING * 1.05 * (1 - adj)
            }
            if (!horizontal) {
                psr <- par("usr")
                MAXSTRING <- MAXSTRING * 1.09 * (psr[4] - psr[3])/(psr[2] - psr[1])
                loy <- label.offset + MAXSTRING * 1.05 * adj
                lox <- 0
                srt <- 90 + srt
                if (direction == "downwards") {
                    loy <- -loy
                    srt <- 180 + srt
                }
            }
        }
        if (type == "phylogram") {
            phylogram.plot(x$edge, Ntip, Nnode, xx, yy, horizontal,
                           edge.color, edge.width, edge.lty)
        } else {
            if (type == "fan") {
                ereorder <- match(z$edge[, 2], x$edge[, 2])
                if (length(edge.color) > 1) {
                    edge.color <- rep(edge.color, length.out = Nedge)
                    edge.color <- edge.color[ereorder]
                }
                if (length(edge.width) > 1) {
                    edge.width <- rep(edge.width, length.out = Nedge)
                    edge.width <- edge.width[ereorder]
                }
                if (length(edge.lty) > 1) {
                    edge.lty <- rep(edge.lty, length.out = Nedge)
                    edge.lty <- edge.lty[ereorder]
                }
                circular.plot(z$edge, Ntip, Nnode, xx, yy, theta,
                              r, edge.color, edge.width, edge.lty)
            } else {
                cladogram.plot(x$edge, xx, yy, edge.color, edge.width, edge.lty)
            }
        }
        if (root.edge) {
            rootcol <- fifelse(length(edge.color) == 1, edge.color, 'black')
            rootw <- fifelse(length(edge.width) == 1, edge.width, 1)
            rootlty <- fifelse(length(edge.lty) == 1, edge.lty, 1)
            if (type == "fan") {
                tmp <- polar2rect(x$root.edge, theta[ROOT])
                segments(0, 0, tmp$x, tmp$y, col = rootcol, lwd = rootw,
                         lty = rootlty)
            } else {
                switch(direction,
                       rightwards = segments(0, yy[ROOT],
                                             x$root.edge, yy[ROOT], col = rootcol, lwd = rootw,
                                             lty = rootlty), leftwards = segments(xx[ROOT],
                                                                                yy[ROOT], xx[ROOT] + x$root.edge, yy[ROOT],
                                                                                col = rootcol, lwd = rootw, lty = rootlty),
                       upwards = segments(xx[ROOT], 0, xx[ROOT], x$root.edge,
                                          col = rootcol, lwd = rootw, lty = rootlty),
                       downwards = segments(xx[ROOT], yy[ROOT], xx[ROOT],
                                            yy[ROOT] + x$root.edge, col = rootcol, lwd = rootw,
                                            lty = rootlty))
            }
        }
        if (show.tip.label) {
            if (is.expression(x$tip.label))
                underscore <- TRUE
            if (!underscore)
                x$tip.label <- gsub("_", " ", x$tip.label)
            if (phyloORclado) {
                if (align.tip.label) {
                    xx.tmp <- switch(direction, rightwards = max(xx[1:Ntip]),
                                     leftwards = min(xx[1:Ntip]), upwards = xx[1:Ntip],
                                     downwards = xx[1:Ntip])
                    yy.tmp <- switch(direction, rightwards = yy[1:Ntip],
                                     leftwards = yy[1:Ntip], upwards = max(yy[1:Ntip]),
                                     downwards = min(yy[1:Ntip]))
                    segments(xx[1:Ntip], yy[1:Ntip], xx.tmp, yy.tmp,
                             lty = align.tip.label.lty)
                }
                else {
                    xx.tmp <- xx[1:Ntip]
                    yy.tmp <- yy[1:Ntip]
                }
                text(xx.tmp + lox, yy.tmp + loy, x$tip.label,
                     adj = adj, font = font, srt = srt, cex = cex,
                     col = tip.color)
            } else {
                angle <- if (type == "unrooted")
                             XY$axe
                         else atan2(yy[1:Ntip], xx[1:Ntip])
                lab4ut <- if (is.null(lab4ut)) {
                              if (type == "unrooted")
                                  "horizontal"
                              else "axial"
                          }
                          else match.arg(lab4ut, c("horizontal", "axial"))
                xx.tips <- xx[1:Ntip]
                yy.tips <- yy[1:Ntip]
                if (label.offset) {
                    xx.tips <- xx.tips + label.offset * cos(angle)
                    yy.tips <- yy.tips + label.offset * sin(angle)
                }
                if (lab4ut == "horizontal") {
                    y.adj <- x.adj <- numeric(Ntip)
                    sel <- abs(angle) > 0.75 * pi
                    x.adj[sel] <- -strwidth(x$tip.label)[sel] *
                        1.05
                    sel <- abs(angle) > pi/4 & abs(angle) < 0.75 *
                        pi
                    x.adj[sel] <- -strwidth(x$tip.label)[sel] *
                        (2 * abs(angle)[sel]/pi - 0.5)
                    sel <- angle > pi/4 & angle < 0.75 * pi
                    y.adj[sel] <- strheight(x$tip.label)[sel]/2
                    sel <- angle < -pi/4 & angle > -0.75 * pi
                    y.adj[sel] <- -strheight(x$tip.label)[sel] *
                        0.75
                    text(xx.tips + x.adj * cex, yy.tips + y.adj *
                                                cex, x$tip.label, adj = c(adj, 0), font = font,
                         srt = srt, cex = cex, col = tip.color)
                } else {
                    if (align.tip.label) {
                        POL <- rect2polar(xx.tips, yy.tips)
                        POL$r[] <- max(POL$r)
                        REC <- polar2rect(POL$r, POL$angle)
                        xx.tips <- REC$x
                        yy.tips <- REC$y
                        segments(xx[1:Ntip], yy[1:Ntip], xx.tips,
                                 yy.tips, lty = align.tip.label.lty)
                    }
                    if (type == "unrooted") {
                        adj <- abs(angle) > pi/2
                        angle <- angle * 180/pi
                        angle[adj] <- angle[adj] - 180
                        adj <- as.numeric(adj)
                    } else {
                        s <- xx.tips < 0
                        angle <- angle * 180/pi
                        angle[s] <- angle[s] + 180
                        adj <- as.numeric(s)
                    }
                    font <- rep(font, length.out = Ntip)
                    tip.color <- rep(tip.color, length.out = Ntip)
                    cex <- rep(cex, length.out = Ntip)
                    for (i in 1:Ntip) text(xx.tips[i], yy.tips[i],
                                           x$tip.label[i], font = font[i], cex = cex[i],
                                           srt = angle[i], adj = adj[i], col = tip.color[i])
                }
            }
        }
        if (show.node.label)
            text(xx[ROOT:length(xx)] + label.offset, yy[ROOT:length(yy)],
                 x$node.label, adj = adj, font = font, srt = srt,
                 cex = cex)
    }
    L <- list(type = type, use.edge.length = use.edge.length,
        node.pos = node.pos, node.depth = node.depth, show.tip.label = show.tip.label,
        show.node.label = show.node.label, font = font, cex = cex,
        adj = adj, srt = srt, no.margin = no.margin, label.offset = label.offset,
        x.lim = x.lim, y.lim = y.lim, direction = direction,
        tip.color = tip.color, Ntip = Ntip, Nnode = Nnode, root.time = x$root.time,
        align.tip.label = align.tip.label, coords = list(edge = xe, xx = xx, yy = yy))
    assign("last_plot.phylo", c(L, list(edge = xe, xx = xx, yy = yy)),
        envir = .PlotPhyloEnv)
    invisible(L)
}

get.children <- function(tedges, from.node = NULL) {
    if (is.null(from.node))
        from.node <- setdiff(tedges$from , tedges$to)
    res <- NULL
    i <- 0L
    if (is.null(key(tedges)) || key(tedges) != 'from')
        setkey(tedges, from)
    children <- tedges[from %in% from.node]
    while(nrow(children)) {
        children[, lvl := i]
        res <- rbind(res, children)
        children <- tedges[from %in% children$to]
        i <- i + 1L
    }
    res
}

get.parents <- function(tedges, from.node) {
    i <- 0L
    res <- vector("list", 1000)
    setkey(tedges, to)
    anc <- tedges[to %in% from.node]

    while(nrow(anc)) {
        anc[, lvl := i]
        i <- i + 1L
        res[[i]] <- anc
        anc <- tedges[to %in% anc$from]
    }
    rbindlist(res)
}

## * Convert subset of the edge data.table to phylo tree
edges2phylo <- function(tedges, idcol = 'virus_name', update.first = T) {

    e <- copy(tedges)
    if (update.first | !('to.tip' %in% names(tedges)))
        update.tree(e)

    rt <- get.root(e)

    tps <- e[to.tip == T, sort(unique(to))]
    ntips <- length(tps)
    nds <- c(rt, setdiff(sort(setdiff(unlist(e[, .(from, to)]), tps)), rt))
    nnodes <- length(nds)

    ## ** Re-name nodes and tips
    lkup <- data.table(id = c(tps, nds), is.tip = c(rep(T, length.out = ntips), rep(F, length.out = nnodes)))
    lkup[, newid := .I]
    setnames(e, c('from', 'to'), c('from.ori', 'to.ori'))

    e[lkup, from := i.newid, on = c('from.ori' = 'id')]
    e[lkup, to   := i.newid, on = c('to.ori'   = 'id')]
    setorder(e, -to.tip, to, from)

    m <- as.matrix(e[, .(from, to)])
    dimnames(m) <- NULL

    cols <- c('to', idcol)
    t <- list(edge        = m,
              edge.length = e$dist,
              Nnode       = nnodes,
              node.label  = NULL,
              tip.label   = e[to.tip == T, ..cols][order(to)][, get(idcol)]
              )
    class(t) <- 'phylo'
    attr(t, 'order') <- 'cladewise'

    list(phylo = t,
         edges = e)
}

## * Prepare data to be plotted - extract coordinates and make edges square if required
edges.plot.prep <- function(tedges, square.edges = T, id.col = 'vertex_id') {

    phy <- edges2phylo(tedges, idcol = id.col) # convert edges data.table back to phylo tree
    tt <- phy$phylo
    eds <- phy$edges

    ## checkValidPhylo(tt)

    ## ** Get coordinates without plotting
    p <- plot2(tt, type = 'phylogram', plot = F,
               show.tip.label = F, show.node.label = F,
               edge.color = 'black', tip.color = 'red', cex = 0,
               root.edge = F, use.edge.length = T,
               no.margin = T)

    ed <- data.table(p$coords$edge)
    xy <- data.table(x = p$coords$xx, y = p$coords$yy)
    xy[, id := .I]

    eds[xy, `:=`(from.x = i.x, from.y = i.y), on = c('from' = 'id')]
    eds[xy, `:=`(to.x = i.x, to.y = i.y), on = c('to' = 'id')]

    ## ** Make square segments
    if (square.edges) {
        e <- eds[from.x != to.x & from.y != to.y]
        ed <- e[, .(to.x = c(from.x, to.x), from.y = c(from.y, to.y), to.tip = c(F, first(to.tip)),
                    extra = c(T, F))
              , .(from, to, eid)]

        ed <- merge(ed, e[, c('eid', setdiff(names(e), names(ed))), with=F], by = 'eid')
        ed <- rbind(ed, eds[!(from.x != to.x & from.y != to.y)], fill = T)
    } else ed <- copy(eds)

    return(ed)
}

## * Assign common trait within clades

spread.clade.trait <- function(tedges, var, agg.var.name = 'children_common_trait') {

    if ('children_agg' %in% names(tedges))
        tedges[, children_agg := NULL]
    if (agg.var.name %in% names(tedges))
        tedges[, eval(agg.var.name) := NULL]
    tedges[, children_agg := get(var)]
    if ('curr_agg' %in% names(tedges))
        tedges[, curr_agg := NULL]

    ## aggval <- tedges[to.tip == T & !is.na(children_agg), mean(children_agg)]
    ## tedges[to.tip == T & is.na(children_agg), children_agg := aggval]

    mxnode <- tedges[, max(nthnode)]
    nod=mxnode-1
    for (nod in (mxnode-1):1L) {
        higher.lvl <- tedges[nthnode == nod+1L]
        higher.lvl[, curr_agg := fifelse(uniqueN(children_agg)==1, children_agg[1], 'multiple'), from]
        lower.lvl  <- tedges[nthnode == nod]
        ## lower.lvl[higher.lvl, children_agg := mean(i.children_agg, na.rm=T), on = c('to' = 'from')]
        lower.lvl[higher.lvl, children_agg := i.curr_agg, on = c('to' = 'from')]
        ## lower.lvl[higher.lvl, curr_agg := mean(i.children_agg, na.rm=T), on = c('to' = 'from')]
        tedges[lower.lvl, children_agg := i.children_agg, on = 'eid']
    }

    if (agg.var.name %in% names(tedges)) {
        tedges[, eval(agg.var.name) := NULL]
    }
    setnames(tedges, 'children_agg', agg.var.name)
    return(invisible(tedges))
}

spread.stat <- function(tedges, var, progress = T, agg.var.name = 'children_agg', stat.fun = mean) {

    tedges[, children_agg := get(var)]
    if ('curr_agg' %in% names(tedges))
        tedges[, curr_agg := NULL]

    aggval <- tedges[to.tip == T & !is.na(children_agg), mean(children_agg)]
    tedges[to.tip == T & is.na(children_agg), children_agg := aggval]

    mxnode <- tedges[, max(nthnode)]
    nod=mxnode-1
    for (nod in (mxnode-1):1L) {
        higher.lvl <- tedges[nthnode == nod+1L]
        higher.lvl[, curr_agg := stat.fun(children_agg, na.rm=T), from]
        lower.lvl  <- tedges[nthnode == nod]
        ## lower.lvl[higher.lvl, children_agg := mean(i.children_agg, na.rm=T), on = c('to' = 'from')]
        lower.lvl[higher.lvl, children_agg := i.curr_agg, on = c('to' = 'from')]
        ## lower.lvl[higher.lvl, curr_agg := mean(i.children_agg, na.rm=T), on = c('to' = 'from')]
        tedges[lower.lvl, children_agg := i.children_agg, on = 'eid']
    }

    setnames(tedges, 'children_agg', agg.var.name)

}

## points = NULL
## width = 55.5; height = 50; square.edges = T;
## edges.colour.by = 'children_trait'; points.colour.by = 'children_trait';
## palette.fun = NULL; z.var = NULL; z.trans.fun = I; max.z = 2e6;
## bgcol = '#EEEEEE'; tooltip = T;
## edge.vars = c('accession_id', 'vertex_id', 'children_trait');
## point.vars = c('accession_id', 'collection_date', 'continent',
##                'country', 'host', 'clade', 'lineage', 'children_trait');
## node.lab.glue = "ID: {accession_id}<br>Continent: {continent}<br>country: {country}<br>Date: {collection_date}";
## edge.lab.glue = "ID: {accession_id}<br>Vertex ID: {vertex_id}<br>{children_trait}"

tree2deck <- function(tedges, points = NULL, width = 55.5, height = 50, square.edges = T,
                      edges.colour.by = 'children_trait', points.colour.by = 'children_trait',
                      thickness.var = 'children_agg',
                      palette.fun = NULL, z.var = NULL, z.trans.fun = I, max.z = 2e6,
                      bgcol = '#EEEEEE', tooltip = T,
                      edge.vars = c('eid', 'continent_spread', 'country_spread'),
                      point.vars = c('eid', 'continent_spread', 'country_spread'),
                      node.lab.glue = '{eid}<br>{continent_spread}<br>{country_spread}',
                      edge.lab.glue = '{eid}<br>{continent_spread}<br>{country_spread}') {
    library(deckgl)

    in3D <- !is.null(z.var)
    if (in3D)
        tedges[, zval := get(z.var)]

    if (!('to.tip' %in% names(tedges)))
        tedges <- update.tree(tedges)

    e <- tedges

    if (is.null(points)) {
        t <- e[to.tip == T]
    } else t <- points

    if (!is.null(edges.colour.by) && !(edges.colour.by %in% names(e)))
        edges.colour.by <- NULL
    if (!is.null(points.colour.by) && !(points.colour.by %in% names(t)))
        points.colour.by <- NULL

    ## ** Colour
    if (!is.null(edges.colour.by)) {
        if (is.numeric(e[[edges.colour.by]])) {
            e[, colour := scales::cscale(get(edges.colour.by),
                                         scales::gradient_n_pal(colours = gplots::rich.colors(30)))]
        } else {
            if (e[!is.na(get(edges.colour.by)), all(grepl('^#', get(edges.colour.by)))]) {
                e[, colour := get(edges.colour.by)]
            } else {
                e[, colour := scales::dscale(factor(get(edges.colour.by)),
                                             pals::alphabet)]
            }
        }
    } else e[, colour := NA_character_]
    e[is.na(colour), colour := '#AAAAAA']

    if (!is.null(points.colour.by)) {
        if (points.colour.by %in% names(t)) {
            if (is.numeric(t[[points.colour.by]])) {
                t[, colour := scales::cscale(get(points.colour.by),
                                             scales::gradient_n_pal(colours = gplots::rich.colors(30)))]
            } else {
                if (t[!is.na(get(points.colour.by)), all(grepl('^#', get(points.colour.by)))]) {
                    t[, colour := get(points.colour.by)]
                } else {
                    t[, colour := scales::dscale(factor(get(points.colour.by)),
                                                 gplots::rich.colors)]
                }
            }
        } else t[, colour := points.colour.by] # if colour var not a field, assume it's a colour
    } else t[, colour := NA_character_]
    t[is.na(colour), colour := '#AAAAAA']

    maxx <- max(e$to.x)
    maxy <- max(e$to.y)

    t[, x := scales::rescale(to.x, c(0, width) , c(0, maxx))]
    t[, y := scales::rescale(to.y, c(0, height), c(0, maxy))]

    e[, start_x := scales::rescale(from.x, c(0, width), c(0, maxx))]
    e[, start_y := scales::rescale(from.y, c(0, height), c(0, maxy))]
    e[, end_x := scales::rescale(to.x, c(0, width), c(0, maxx))]
    e[, end_y := scales::rescale(to.y, c(0, height), c(0, maxy))]
    if (!is.null(thickness.var) && thickness.var %in% names(e)) {
        e[, linewidth := scales::rescale(sqrt(get(thickness.var)), c(1, 20))]
    } else e[, linewidth := 1]

    if (in3D) {
        zrange <- e[, range(zval, na.rm=T)]
        t[, z       := scales::rescale(z.trans.fun(zval)    , c(0, max.z), z.trans.fun(zrange))]
        e[, start_z := scales::rescale(z.trans.fun(curr_agg), c(0, max.z), z.trans.fun(zrange))]
        e[, end_z   := scales::rescale(z.trans.fun(zval)    , c(0, max.z), z.trans.fun(zrange))]
    } else {
        t[, z := NA]
        e[, c('start_z', 'end_z') := NA]
    }

    vars <- c('x', 'y', 'z', 'colour', point.vars)
    pp <- t[, ..vars]
    vars <- c('start_x', 'start_y', 'start_z', 'end_x', 'end_y', 'end_z', 'colour', 'linewidth', edge.vars)
    ll <- e[, ..vars]

    root <- setdiff(tedges$from, tedges$to)
    rootpt <- e[from == root][1, .(start_x, start_y, start_z, end_x, end_y, end_z, eid)]

    pp[, label := glue::glue(node.lab.glue, envir = .SD)]
    ll[, label := as.character(glue::glue(edge.lab.glue, envir = .SD))]

    leg <- tedges[!is.na(colour), .N, c(edges.colour.by, 'colour')][order(get(edges.colour.by))]
    setnames(leg, edges.colour.by, 'val')

    if (in3D) {

        if (tooltip) {

            proplines <- list(
                getSourcePosition = ~ start_x + start_y + start_z,
                getTargetPosition = ~ end_x + end_y + end_z,
                getTooltip = JS("object =>`${object.label}`"),
                getWidth = 1,
                pickable = T,
                opacity = 0.5,
                visible = T,
                getColor = get_color_to_rgb_array('colour')
            )

            proppoints <- list(
                getPosition = ~ x + y + z,
                getColor = get_color_to_rgb_array('colour'),
                getFillColor = get_color_to_rgb_array('colour'),
                getRadius = 1000,
                getTooltip = JS("object =>`${object.label}`"),
                stroked = T,
                pointSize = 2,
                radiusScale = 8,
                radiusMinPixels = 1.5,
                radiusMaxPixels = 5,
                opacity = 0.8
            )

            proproot <- list(
                getPosition = ~ start_x + start_y + start_z,
                getColor = '#000000',
                getFillColor = '#FFFFFF',
                getRadius = 1000,
                getTooltip = "Tree root<br>Epi-id: {{eid}}",
                stroked = T,
                radiusScale = 8,
                radiusMinPixels = 5,
                radiusMaxPixels = 15,
                opacity = 1
            )

            deck <- deckgl(width='100vw', height='100vh', latitude = 25, longitude = 25, zoom = 3,
                           pitch = fifelse(in3D, 60, 0), maxPitch = 89.9,
                           style = list(background = bgcol)) %>%
                add_line_layer(data = ll[, .(start_x, start_y, start_z, end_x, end_y, end_z, colour, label)],
                               properties = proplines) %>%
                add_scatterplot_layer(data = pp[, .(x, y, z, colour, label)], properties = proppoints) %>%
                add_scatterplot_layer(data = rootpt[, .(start_x, start_y, start_z, eid)], properties = proproot) %>%
                add_legend(
                    colors = leg$colour,
                    labels = leg$val,
                    title = '',
                    pos = 'top-left'
                )

        } else {

            proppoints <- list(
                getPosition = ~ x + y + z,
                getColor = get_color_to_rgb_array('colour'),
                getFillColor = get_color_to_rgb_array('colour'),
                getRadius = 1000,
                stroked = T,
                pointSize = 2,
                radiusScale = 8,
                radiusMinPixels = 1.5,
                radiusMaxPixels = 5,
                opacity = 0.8
            )

            proplines <- list(
                getSourcePosition = ~ start_x + start_y + start_z,
                getTargetPosition = ~ end_x + end_y + end_z,
                getWidth = 1,
                pickable = T,
                opacity = 0.5,
                visible = T,
                getColor = get_color_to_rgb_array('colour')
            )

            proproot <- list(
                getPosition = ~ start_x + start_y + start_z,
                getColor = '#000000',
                getFillColor = '#FFFFFF',
                getRadius = 1000,
                stroked = T,
                radiusScale = 8,
                radiusMinPixels = 5,
                radiusMaxPixels = 15,
                opacity = 1
            )

            deck <- deckgl(width='100vw', height='100vh', latitude = 25, longitude = 25, zoom = 3,
                           pitch = fifelse(in3D, 60, 0), maxPitch = 89.9,
                           style = list(background = bgcol)) %>%
                add_line_layer(data = ll[, .(start_x, start_y, start_z, end_x, end_y, end_z, colour)],
                               properties = proplines) %>%
                add_scatterplot_layer(data = pp[, .(x, y, z, colour)], properties = proppoints) %>%
                add_scatterplot_layer(data = rootpt[, .(start_x, start_y, start_z)], properties = proproot) %>%
                add_legend(
                    colors = leg$colour,
                    labels = leg$val,
                    title = '',
                    pos = 'top-left'
                )

        }

    } else {

        if (tooltip) {

            proplines <- list(
                getSourcePosition = ~ start_x + start_y,
                getTargetPosition = ~ end_x + end_y,
                getTooltip = JS("object =>`${object.label}`"),
                getWidth = ~ linewidth,
                widthScale = 1,
                pickable = T,
                opacity = 0.8,
                visible = T,
                getColor = get_color_to_rgb_array('colour')
            )

            proppoints <- list(
                getPosition = ~ x + y,
                getColor = get_color_to_rgb_array('colour'),
                getFillColor = get_color_to_rgb_array('colour'),
                getRadius = 1000,
                getTooltip = JS("object =>`${object.label}`"),
                stroked = T,
                radiusScale = 8,
                radiusMinPixels = 1.5,
                radiusMaxPixels = 5,
                opacity = 0.8
            )

            proproot <- list(
                getPosition = ~ start_x + start_y,
                getColor = '#000000',
                getFillColor = '#FFFFFF',
                getRadius = 1000,
                getTooltip = "Tree root<br>Epi-id: {{eid}}",
                stroked = T,
                radiusScale = 8,
                radiusMinPixels = 5,
                radiusMaxPixels = 15,
                opacity = 1
            )

            deck <- deckgl(width='100vw', height='100vh', latitude = 25, longitude = 25, zoom = 3,
                           pitch = fifelse(in3D, 60, 0), maxPitch = 89.9,
                           style = list(background = bgcol)) %>%
                add_line_layer(data = ll[, .(start_x, start_y, end_x, end_y, colour, label, linewidth)],
                               properties = proplines) %>%
                add_scatterplot_layer(data = pp[, .(x, y, colour, label)], properties = proppoints) %>%
                add_scatterplot_layer(data = rootpt[, .(start_x, start_y, eid)], properties = proproot) %>%
                add_legend(
                    colors = leg$colour,
                    labels = leg$val,
                    title = '',
                    pos = 'top-left'
                )

        } else {

            proppoints <- list(
                getPosition = ~ x + y,
                getColor = get_color_to_rgb_array('colour'),
                getFillColor = get_color_to_rgb_array('colour'),
                getRadius = 1000,
                stroked = T,
                radiusScale = 8,
                radiusMinPixels = 1.5,
                radiusMaxPixels = 5,
                opacity = 0.8
            )

            proproot <- list(
                getPosition = ~ start_x + start_y,
                getColor = '#000000',
                getFillColor = '#FFFFFF',
                getRadius = 1000,
                stroked = T,
                radiusScale = 8,
                radiusMinPixels = 5,
                radiusMaxPixels = 15,
                opacity = 1
            )

            proplines <- list(
                getSourcePosition = ~ start_x + start_y,
                getTargetPosition = ~ end_x + end_y,
                getWidth = 1,
                pickable = T,
                opacity = 0.8,
                visible = T,
                getColor = get_color_to_rgb_array('colour')
            )

            deck <- deckgl(width='100vw', height='100vh', latitude = 25, longitude = 25, zoom = 3,
                           pitch = fifelse(in3D, 60, 0), maxPitch = 89.9,
                           style = list(background = bgcol)) %>%
                add_line_layer(data = ll[, .(start_x, start_y, end_x, end_y, colour)],
                               properties = proplines) %>%
                add_scatterplot_layer(data = pp[, .(x, y, colour)], properties = proppoints) %>%
                add_scatterplot_layer(data = rootpt[, .(start_x, start_y)], properties = proproot) %>%
                add_legend(
                    colors = leg$colour,
                    labels = leg$val,
                    title = '',
                    pos = 'top-left'
                )

        }

    }

    return(deck)
}

tree.from.edge <- function(tedges, sel.eid=123309) {
    from.node <- tedges[eid == sel.eid, to]
    pars <- get.parents(tedges, from.node)
    childs <- get.children(tedges, from.node)
    rbind(pars, childs)
}

tree.from.attribute <- function(tedges, attr.var = 'country', attr.val = 'Australia') {
    te <- tedges[get(attr.var) == attr.val]
    tetips <- te[to.tip == T]
    parents <- get.parents(tedges, tetips$to)
    parents <- unique(parents[, -'lvl'])
    parents
}

upper1st <- function (x, prelower = T) {
    if (is.factor(x)) {
        x0 <- levels(x)
    } else x0 <- as.character(x)
    if (prelower) x0 <- tolower(x0)
    nc <- nchar(x0)
    nc[is.na(nc)] <- 0
    x0[nc > 1] <- sprintf("%s%s", toupper(substr(x0[nc > 1], 1, 1)),
                          substr(x0[nc > 1], 2, nchar(x0[nc > 1])))
    x0[nc == 1] <- toupper(x0[nc == 1])
    if (is.factor(x)) {
        levels(x) <- x0
    } else x <- x0
    return(x)
}

hier.pal <- function(category, subcategory, spread.within = 0.3,
                     chroma = 110, luminance = 70, alpha = 1) {

    cols0 <- data.table(id1 = category, id2 = subcategory)
    cols0[, i := .I]
    setkey(cols0, id1, id2)

    cols <- cols0[, .N, .(id1, id2)]

    ncats <- uniqueN(cols$id1)
    steps <- seq(0, 360, length.out = ncats + 1)
    diff <- steps[2]-steps[1]
    steps <- steps + diff / 2

    steps <- head(steps, ncats)
    names(steps) <- unique(cols$id1)
    cols[, step := steps[id1]]

    spread <- diff * spread.within / 2
    cols[, offset := seq(-spread, spread, length.out = .N), id1]
    cols[, hue := step + offset]

    cols[, c := chroma + offset]
    cols[, l := luminance + offset]

    cols[, col := hcl(h = hue, c = c, l = l, alpha = alpha)]

    setorder(cols0, i)
    return(cols[cols0, col, on = c('id1', 'id2')])
}