library(stockassessment)
source("src/common.R")
# load what has been saved
setwd("run")
for(f in dir(pattern="RData"))load(f) 
setwd("..")

basefit<-NULL
if(file.exists("baserun/model.RData")){
  local({load("baserun/model.RData"); basefit<<-fit})
}else{
  basefit <- fit
}
fits <- c(base=basefit,current=fit)

exfitname <- scan("conf/viewextra.cfg", what="", comment.char="#", quiet=TRUE)
for(nam in exfitname){
  local({
    fit<-urlLoadFit(paste0("https://www.stockassessment.org/datadisk/stockassessment/userdirs/user3/",nam,"/run/model.RData"))
    if(!is.null(fit)){
      i <- length(fits)
      fits[[i+1]] <<- fit
      names(fits)[i+1] <<- nam
    }else{
      warning(paste0("View extra stock ", nam, " not found or of incompatible format (skipped)"))
    }
  })
}

plotcounter<-1
tit.list<-list()

setcap<-function(title="", caption=""){   
 tit.list[length(tit.list)+1]<<-paste("# Title",plotcounter,"\n")
 tit.list[length(tit.list)+1]<<-paste(title,"\n")
 tit.list[length(tit.list)+1]<<-paste("# Caption",plotcounter,"\n")
 tit.list[length(tit.list)+1]<<-paste(caption,"\n")
 plotcounter<<-plotcounter+1 
}


############################## plots ##############################
plots<-function(){
  par(cex.lab=1, cex.axis=1, mar=c(5,5,1,1))
    
  if(exists("fits")){
    ssbplot(fits, addCI=TRUE)
    stampit(fit)
    setcap("Spawning stock biomass", "Spawning stock biomass. 
            Estimates from the current run and point wise 95% confidence 
            intervals are shown by line and shaded area.")
    
    fbarplot(fits, addCI=TRUE,partial=FALSE,pcol='blue')
    stampit(fit)
    setcap("Average fishing mortality", "Average fishing mortality for the shown age range. 
            Estimates from the current run and point wise 95% confidence 
            intervals are shown by line and shaded area.")

    recplot(fits, addCI=TRUE, las=0)
    stampit(fit)
    setcap("Recruitment", "Yearly resruitment. 
        Estimates from the current run and point wise 95% confidence 
        intervals are shown by line and shaded area.")

    catchplot(fits, addCI=TRUE)
    stampit(fit)
    setcap("Catch", "Total catch in weight. 
        Prediction from the current run and point wise 95% confidence 
        intervals are shown by line and shaded area. The yearly
        observed total catch weight (crosses) are calculated as Cy=sum(WayCay).")

# HG added
matplot(fit$data$years, faytable(fit)/rowSums(faytable(fit)), type="b", lwd=2, pch=1:7,col=1:7, xlab="Year", ylab="Selectivity-at-age")
legend("topleft", legend=colnames(faytable(fit)), lwd=2, col=1:7, pch=1:7, lty=1:5, ncol=3)
stampit(fit)
setcap("Catch selectivity", "F at age divided by total F")

# HG added
Q<-fit$pl$logFpar
Qsd<-fit$plsd$logFpar
key<-fit$conf$keyLogFpar
fun<-function(x)if(x<0){NA}else{Q[x+1]}
FF<-Vectorize(fun)
ages<-fit$conf$minAge:fit$conf$maxAge
matplot(ages, exp(t(matrix(FF(key), nrow=nrow(key)))), type="l", lwd=5, lty="solid", xlab="Ages", ylab="Q")
legend("topright", lwd=5, col=2:nrow(key), legend=attr(fit$data, "fleetNames")[2:nrow(key)])
stampit(fit)
setcap("Survey catchability", "Survey catchability")      
        

    srplot(fit)
    stampit(fit)
    setcap("Spawner-resruits", "Estimated recruitment as a function of spawning stock biomass.")

    plot(ypr(fit))
    stampit(fit)
    setcap("Yield per Recruit", "Yield per recruit (solid line) and spawning stock biomass plotted against different levels of fishing")
   
    if(!all(fit$conf$obsCorStruct=="ID")){ 
      corplot(fit)			  
      setcap("Estimated correlations", "Estimates correlations between age groups for each fleet")
      stampit(fit)
    }
    
    for(f in 1:fit$data$noFleets){
      fitplot(fit, fleets=f)
      setcap("Fit to data", "Predicted line and observed points (log scale)")
      stampit(fit)
    }
    
       
    # how much biomass is there at each age
    barplot(prop.table(t(fit$data$stockMeanWeight*ntable(fit)),margin=2),space=0,legend=T,col=terrain.colors(dim(ntable(fit))[2]))
    setcap("Biomass at age", "Biomass at age")
    stampit(fit)
 
	# observation error??    
    ptab <- partable(fit)
    par(mar=c(8,4,1,1))
    barplot(ptab[,'sd(par)'],las=2,ylab='sd(par)')
    setcap("Observation error", "Observation error (I think)")
    stampit(fit)


  }  
  
  if(exists("RES")){  
    plot(RES)
    setcap("One-observation-ahead residuals", "Standardized one-observation-ahead residuals.")
    stampit(fit)
    par(mfrow=c(1,1))
    empirobscorrplot(RES)
    setcap("OOA residual correlations", "Empirical correlations between ages in one-observation-ahead residuals.")
    stampit(fit)
  }
  
   
  if(exists("RESP")){  
    plot(RESP)
    setcap("Process residuals", "Standardized single-joint-sample residuals of process increments")
    stampit(fit)
    par(mfrow=c(1,1))
  } 
 
  
  if(exists("LO")){  
    ssbplot(LO)
    setcap("Leaveout (SSB)", "")
    stampit(fit)
    
    fbarplot(LO)
    setcap("Leaveout (Average F)", "")
    stampit(fit)

    recplot(LO)
    setcap("Leaveout (Recruitment)", "")
    stampit(fit)

    catchplot(LO)
    setcap("Leaveout (Catch)", "")
    stampit(fit)
    
  } 
  
  if(exists("RETRO")){  
    rho <- mohn(RETRO)
    ssbplot(RETRO, las=0, drop=0)  # drop =1 original
    legend("topright", legend=paste0("Rho ",round(rho[2]*100),"%"))
    setcap("Retrospective (SSB)", "")
    stampit(fit)
    
    fbarplot(RETRO, las=0, drop=0)  # drop =1 original
    legend("topright", legend=paste0("Rho ",round(rho[3]*100),"%"))
    setcap("Retrospective (Average F)", "")
    stampit(fit)
    
    recplot(RETRO, las=0, drop=0)  # drop =1 original
    legend("topright", legend=paste0("Rho ",round(rho[1]*100),"%"))
    setcap("Retrospective (Recruitment)", "")
    stampit(fit)

    catchplot(RETRO)
    setcap("Retrospective (Catch)", "")
    stampit(fit)
    
    parplot(RETRO)
    setcap("Retrospective (parameters)", "")
    stampit(fit)
  } 
  
  if(exists("FC")){  
    #lapply(FC), function(f){plot(f); title(attr(f,"label"), outer=TRUE, line=-1); stampit(fit)})
    lapply(list(FC[[1]]), function(f){plot(f); title(attr(f,"label"), outer=TRUE, line=-1); stampit(fit)}) #hg edit: only plot first one
  }  
  
    if(exists("FC")){  

      fc <- FC[[1]] #assume that the first forecast is the headline
      catchoption <- names(FC)[1]

      final.yr <- max(fit$data$years)
      
      ## default to take average last 3 years in forecast
      ave.years <- final.yr+(-2:0)
      cw <- apply(tail(fit$data$catchMeanWeight[,,1,drop=T],length(ave.years)),2,mean) # mean catch weights
      sw <- apply(tail(fit$data$stockMeanWeight,length(ave.years)),2,mean) # mean stock weights
      mo <- apply(tail(fit$data$propMat,length(ave.years)),2,mean) # maturity
      
      ## catch-at-age in advice year (last data year = 1; intermerdiate year = 2; and advice year = 3)
      ## this does not sum up to the catch because the sum of median of catch-at-age is not the same as 
      ## the median of the sum of the catch-at-age. So rescale the catch at age to sum to catch
      cage <- apply(fc[[3]]$catchatage*cw,1,median)
      ca2 <- median(apply(fc[[3]]$catchatage*cw,2,sum))
      catch <- cage*ca2/sum(cage)
      sum(catch)
      attr(fc,'shorttab')[4,3]
      
      ## same for ssb
      ages <- min(fit$data$minAgePerFleet):max(fit$data$maxAgePerFleet)
      ## f and m before spawning both zero, so dont have to worry about them
      ssbage <-  apply(exp(fc[[4]]$sim[,1:length(ages)]), 2, median) * sw * mo
      s2 <-  median(apply(t(exp(fc[[4]]$sim[,1:length(ages)]))*sw*mo,2,sum))
      ssb <- ssbage*s2/sum(ssbage)
      sum(ssb)
      attr(fc,'shorttab')[3,4]
      
      plot.catch <-data.frame(recruitment=final.yr+2-ages, val=catch, type=paste(final.yr+2,"Catch"))
      plot.catch$Prop<- plot.catch$val/sum(plot.catch$val)*100
      plot.catch$Age<- ages
      
      plot.SSB <-  data.frame(recruitment=final.yr+3-ages, val=ssb, type=paste(final.yr+3,"SSB"))
      plot.SSB$Prop<- plot.SSB$val/sum(plot.SSB$val)*100
      plot.SSB$Age<- ages
      
      plotfun <- function(x,...) {
        b <- barplot(x$val,names=x$recruitment,horiz=F,border=NA,
                     xlab='Cohort',ylab='Tonnes',ylim=c(0,max(x$val)*1.1),
                     las=2,...)
        text(b,x$val,paste0(round(x$Prop),'%'),pos=3)
      }
      par(mfrow=c(1,2))
      plotfun(plot.catch,main=paste(final.yr+2,'Catch'))
      plotfun(plot.SSB,main=paste(final.yr+3,'SSB'))
      par(mfrow=c(1,1))
      mtext(catchoption,3,0)
      
      setcap(paste("Contribution to forecast:",catchoption), "")
      stampit(fit)
    
  }
  
}


setwd('res')
file.remove(dir(pattern='png$'))
stamp<-gsub('-[[:digit:]]{4}$','',gsub(':','.',gsub(' ','-',gsub('^[[:alpha:]]{3} ','',date()))))
png(filename = paste(stamp,"_%03d.png", sep=''), width = 480, height = 480,
    units = "px", pointsize = 10, bg = "white")
  plots()    
dev.off()

writeLines(unlist(tit.list),'titles.cfg') 

png(filename = paste("big_",stamp,"_%03d.png", sep=''), width = 1200, height = 1200, 
    units = "px", pointsize = 20, bg = "white")
  plots()    
dev.off()

#pdf(onefile=FALSE, width = 8, height = 8)
#  plots()    
#dev.off()

file.remove(dir(pattern='html$'))

tsb<-tsbtable(fit)
colnames(tsb)<-c("TSB","Low", "High")
tab.summary <- cbind(summary(fit), tsb)
xtab(tab.summary, caption=paste('Table 1. Estimated recruitment, spawning stock biomass (SSB), 
     and average fishing mortality','.',sep=''), cornername='Year', 
     file=paste(stamp,'_tab1.html',sep=''), dec=c(0,0,0,0,0,0,3,3,3,0,0,0))

ftab <- faytable(fit)
xtab(ftab, caption=paste('Table 2. Estimated fishing mortality at age','.',sep=''), cornername='Year \ Age', 
     file=paste(stamp,'_tab2.html',sep=''), dec=rep(3,ncol(ftab)))

ntab <- ntable(fit)
xtab(ntab, caption=paste('Table 3. Estimated stock numbers at age','.',sep=''), cornername='Year \ Age', 
     file=paste(stamp,'_tab3.html',sep=''), dec=rep(0,ncol(ntab)))

ptab <- partable(fit)
xtab(ptab, caption=paste('Table 4. Table of model parameters','.',sep=''), cornername='Parameter name', 
     file=paste(stamp,'_tab4.html',sep=''), dec=rep(3,ncol(ptab)))

mtab <- modeltable(c(Current=fit, base=basefit))
mdec <- c(2,0,2,6)[1:ncol(mtab)]
xtab(mtab, caption=paste('Table 5. Model fitting','.',sep=''), cornername='Model', 
     file=paste(stamp,'_tab5.html',sep=''), dec=mdec)
     
sdState<-function(fit, y=max(fit$data$years)-1:0){
  idx <- names(fit$sdrep$value) == "logR"
  sdLogR<-fit$sdrep$sd[idx][fit$data$years%in%y]
  idx <- names(fit$sdrep$value) == "logssb"
  sdLogSSB<-fit$sdrep$sd[idx][fit$data$years%in%y]
  idx <- names(fit$sdrep$value) == "logfbar"
  sdLogF<-fit$sdrep$sd[idx][fit$data$years%in%y]
  ret<-cbind(sdLogR, sdLogSSB, sdLogF)
  rownames(ret)<-y
  colnames(ret)<-c("sd(log(R))", "sd(log(SSB))", "sd(log(Fbar))")
  return(ret)
}

sdtab <- sdState(fit)
xtab(sdtab, caption=paste('Table 6. Table of selected sd','.',sep=''), cornername='Year', 
     file=paste(stamp,'_tab6.html',sep=''), dec=rep(3,ncol(sdtab)))


if(exists("FC")){  
    ii<-0
    lapply(FC, function(f){
       ii<<-ii+1;
       tf<-attr(f,"tab");
       dec <- c(3,3,3,  0,0,0,0,0,0,0,0,0,  3,3,3,3,3,3,  0,0,0,0,0,0)  #     rep(3,ncol(tf)-3));    # dec<-c(3,3,3,rep(0,ncol(tf)-3));
       xtab(tf, caption=paste0('Forecast table ',ii,'. ', attr(f,"label"),'.'), 
       cornername='Year', file=paste(stamp,'_tabX',sprintf("%03d", ii),'.html',sep=''), dec=dec);      
       })
}  


if(exists("iy")) {
  ii<<-ii+1;
  xtab(as.matrix(t(iy)),caption='Advice table 1', width="25%",
                      cornername='Variable',file=paste(stamp,'_tabX',sprintf("%03d", ii),'.html',sep=''), dec=2)
}

if(exists("cs")) {
  ii<<-ii+1;
  dec <- c(0,0,0,3,3,3,0,0,0,2)
  xtab(as.matrix(cs),caption='Advice table 2', width="50%",
       cornername='Basis', file=paste(stamp,'_tabX',sprintf("%03d", ii),'.html',sep=''), dec=dec)
}

if(exists("rp")) {
  ii<<-ii+1;
  xtab(as.matrix(t(rp)),caption='Advice table 4 (just here to make sure the reference points match)', width="25%",
       cornername='Variable',file=paste(stamp,'_tabX',sprintf("%03d", ii),'.html',sep=''), dec=3)
}


setwd("..")