############################## Forecast ##############################
source('src/datavalidator.R') 
oldwd<-getwd()
# Read in data 
setwd('data')
catch.no<-read.ices('cn.dat')
catch.mean.weight<-read.ices('cw.dat')
dis.mean.weight<-read.ices('dw.dat')
land.mean.weight<-read.ices('lw.dat')
stock.mean.weight<-read.ices('sw.dat')
prop.mature<-read.ices('mo.dat')
natural.mortality<-read.ices('nm.dat')
surveys<-read.surveys('survey.dat')
land.frac<-read.ices('lf.dat')
prop.f<-read.ices('pf.dat')
prop.m<-read.ices('pm.dat')
setwd('..')

catch.no.land<-land.frac*catch.no
catch.no.dis<-(1-land.frac)*catch.no

require(MASS, quietly=TRUE)

source('src/common.R')
fit.current<-read.fit('run/sam')

minAge<-min(fit.current$res[,3])  
maxAge<-max(fit.current$res[,3]) 

library(MASS)

conffile<-'run/model.cfg'
conflines<-readLines(conffile)
idx1<-grep('of fishing mortality STATES',conflines)
idx2<-grep('Use correlated random walks',conflines)
tc<-textConnection(conflines[idx1:idx2])
iF<-as.vector(as.matrix(read.table(tc,head=FALSE))[1,])+ncol(stock.mean.weight)
iF<-iF[iF!=0]
close(tc)

idx1<-grep('log N RW VARIANCES',conflines)
idx2<-grep('of OBSERVATION VARIANCES',conflines)
tc<-textConnection(conflines[idx1:idx2])
idxNVar<-as.vector(as.matrix(read.table(tc,head=FALSE))[1,])
close(tc)
idx1<-grep('Fbar range',conflines)
bAges<-scan(conffile,skip=idx1, n=2, quiet=TRUE)
bAges<-do.call(':',as.list(bAges))

doone<-function(fit, stock.mean.weight, catch.mean.weight, prop.mature, natural.mortality,
                fadj=c(1,1,1), last.year.used=max(fit.current$years)-1, noSim=1000, 
                idxN=1:ncol(stock.mean.weight), 
                idxF=c((ncol(stock.mean.weight)+1):fit$stateDim,fit$stateDim), 
                barAges=bAges+(1-minAge), aveYears=3, sampleRec=NA, 
                fixNextYearTAC=NA, fixNextYearTACAll=FALSE, 
                ffix=c(NA,NA,NA), ffixAll=FALSE, bwscale=c(1,1,1,1), use.next.years.N=TRUE, aveYearsLF=1){
  set.seed(12345678)
  fadj[is.na(fadj)]<-1
  fadjlocal<-c(fadj[1],fadj[2]/fadj[1],fadj[3]/fadj[2])
  idxno<-which(fit$years==last.year.used)
  dim<-fit$stateDim
  
  sel<-function(idxVec=(idxno-aveYears+1):idxno){
    Sa<-rep(0,length(idxF))
    K<-0
    for(idxno in idxVec){
      idx<-which(fit$names=='U')[((idxno-1)*dim+1):(idxno*dim)]
      state<-fit$est[idx]
      Sa<-Sa+exp(state[idxF])
      K<-K+mean(exp(state[idxF[barAges]]))
    }
    return(Sa/K)
  }
  select<-sel()  
  ave.sw<-colMeans(stock.mean.weight[(idxno-aveYears+1):idxno,,drop=FALSE])
  ave.cw<-colMeans(catch.mean.weight[(idxno-aveYears+1):idxno,,drop=FALSE])
  ave.pm<-colMeans(prop.mature[(idxno-aveYears+1):idxno,,drop=FALSE])
  ave.nm<-colMeans(natural.mortality[(idxno-aveYears+1):idxno,,drop=FALSE])
  ave.lf<-colMeans((catch.no.land/catch.no)[(idxno-aveYearsLF+1):idxno,,drop=FALSE])
  ave.cw.land<-colMeans(land.mean.weight[(idxno-aveYears+1):idxno,,drop=FALSE]) 


  idx<-which(fit$names=='U')[((idxno-1)*dim+1):(idxno*dim)]
  last.state<-fit$est[idx]
  last.state.var<-fit$cov[idx,idx]

  if(use.next.years.N){
    idxno.next<-which(fit$years==last.year.used+1)
    idx.next<-which(fit$names=='U')[((idxno.next-1)*dim+1):(idxno.next*dim)]
    last.two.state<-fit$est[c(idx,idx.next)]
    last.two.state.var<-fit$cov[c(idx,idx.next),c(idx,idx.next)]
  }

  catch.1<-function(x){
    F<-exp(x[idxF])
    Z<-F+natural.mortality[idxno,]
    N<-exp(x[idxN])
    lf<-catch.no.land[idxno,]/catch.no[idxno,]
    C<-F/Z*(1-exp(-Z))*N
    return(sum(catch.mean.weight[idxno,]*C))
  }

  catch.1.land<-function(x){
    F<-exp(x[idxF])
    Z<-F+natural.mortality[idxno,]
    N<-exp(x[idxN])
    lf<-catch.no.land[idxno,]/catch.no[idxno,]
    C<-F/Z*(1-exp(-Z))*N*lf
    return(sum(land.mean.weight[idxno,]*C))
  }

  catch<-function(x){
    F<-exp(x[idxF])
    Z<-F+ave.nm
    N<-exp(x[idxN])
    C<-F/Z*(1-exp(-Z))*N
    return(sum(ave.cw*C))
  }

  catch.land<-function(x){
    F<-exp(x[idxF])
    Z<-F+ave.nm
    N<-exp(x[idxN])
    C<-F/Z*(1-exp(-Z))*N*ave.lf
    return(sum(ave.cw.land*C))
  }

  fbar.1.land<-function(x){
    F<-exp(x[idxF])
    Z<-F+natural.mortality[idxno,]
    N<-exp(x[idxN])
    lf<-catch.no.land[idxno,]/catch.no[idxno,]
    C.l<-F/Z*(1-exp(-Z))*N*lf
    Fvec<-(C.l*Z)/((1-exp(-Z))*N)
    return(mean(Fvec[barAges]))
  }

  fbar.land<-function(x){
    F<-exp(x[idxF])
    Z<-F+ave.nm
    N<-exp(x[idxN])
    C.l<-F/Z*(1-exp(-Z))*N*ave.lf
    Fvec<-(C.l*Z)/((1-exp(-Z))*N)
    return(mean(Fvec[barAges]))
  }

  # Last year Y
  if(use.next.years.N){
    last.two.state.sim<-mvrnorm(noSim, mu=last.two.state, Sigma=last.two.state.var)
    last.state.sim<-last.two.state.sim[,(1:dim)]
    last.state.sim.next<-last.two.state.sim[,-(1:dim)]
  }else{
    last.state.sim<-mvrnorm(noSim, mu=last.state, Sigma=last.state.var)
  }

  last.ssb.sim<-exp(last.state.sim[,idxN])%*%(stock.mean.weight[idxno,]*prop.mature[idxno,])
  last.fbar.sim<-apply(last.state.sim,1,function(x)mean(exp(x[idxF[barAges]])))
  last.catch.sim<-apply(last.state.sim,1,function(x)catch.1(x))
  last.land.sim<-apply(last.state.sim,1,function(x)catch.1.land(x))
  last.fbarl.sim<-apply(last.state.sim,1,function(x)fbar.1.land(x))
 
  # Step one Year 
  step<-function(x, scale=1.0){
    F<-exp(x[idxF])
    Z<-F+ave.nm
    if(is.na(sampleRec)){
      N<-c(exp(x[idxN[1]]),
           exp(x[idxN[-length(idxN)]])*exp(-Z[-length(Z)])+
           c(rep(0,length(idxN)-2),exp(x[idxN[length(idxN)]])*exp(-Z[length(Z)]))
          )
    }else{
      N<-c(sample(fit$R[(idxno-sampleRec+1):idxno,1],1),
           exp(x[idxN[-length(idxN)]])*exp(-Z[-length(Z)])+
           c(rep(0,length(idxN)-2),exp(x[idxN[length(idxN)]])*exp(-Z[length(Z)]))
          )
    }
    fbar<-mean(exp(x[idxF[barAges]]))*scale
    F<-fbar*select
    return(log(c(N,F[-length(F)])))
  }

  addscale <- 1
  s1.state.sim<-t(apply(last.state.sim,1,step, scale=fadjlocal[1]*addscale))

  if(use.next.years.N){
    ##s1.state.sim[,idxN[-1]]<-last.state.sim.next[,idxN[-1]]
    s1.state.sim[,idxN]<-last.state.sim.next[,idxN]
  }else{
    m<-rep(0,length(idxN[-1]))
    s<-diag(exp(2.0*fit$est[fit$names=='logSdLogN'][idxNVar]))
    if(is.na(sampleRec)){
      s<-s[idxN,idxN]
      s1.state.sim[,idxN]<-s1.state.sim[,idxN]+mvrnorm(noSim,m,s)
    }else{
      s<-s[idxN[-1],idxN[-1]]
      s1.state.sim[,idxN[-1]]<-s1.state.sim[,idxN[-1]]+mvrnorm(noSim,m,s)
    }
  }

  # Step one Year 
  scaleN12<-function(x, scale=c(1.0,1.0)){
    x[idxN[1]]<-x[idxN[1]]+log(scale[1])
    x[idxN[2]]<-x[idxN[2]]+log(scale[2])
    return(x)
  }
  s1.state.sim<-t(apply(s1.state.sim,1,scaleN12, scale=bwscale[1:2]))

  if((!is.na(fixNextYearTAC))&(!fixNextYearTACAll)){# adjust median 
    TAC<-fixNextYearTAC
    obj<-function(logS){
      C<-median(apply(s1.state.sim,1,function(x){x[-idxN]<-x[-idxN]+logS;return(catch(x))}))
      return((TAC-C)^2)
    }
    logS<-optimize(obj,interval=c(-10,10))$min
    s1.state.sim<-t(apply(s1.state.sim,1,function(x){x[-idxN]<-x[-idxN]+logS;return(x)}))
  }

  if((!is.na(fixNextYearTAC))&(fixNextYearTACAll)){# adjust all individually 
    TAC<-fixNextYearTAC
    adjone<-function(x){
      obj<-function(logS){
        xx<-x
        xx[-idxN]<-xx[-idxN]+logS;
        return((TAC-catch(xx))^2)
      }
      logS<-optimize(obj,interval=c(-10,10))$min
      x[-idxN]<-x[-idxN]+logS
      return(x)
    }
    s1.state.sim<-t(apply(s1.state.sim,1,adjone))
  }

  if((!is.na(ffix[1]))&(!ffixAll)){
    fac<-ffix[1]/median(apply(s1.state.sim,1,function(x)mean(exp(x[idxF[barAges]]))))
    s1.state.sim<-t(apply(s1.state.sim,1,function(x){x[-idxN]<-x[-idxN]+log(fac);return(x)}))
    newf1<-ffix[1]/median(last.fbar.sim)
    fadjlocal<-c(newf1,fadj[2]/newf1,fadj[3]/fadj[2])
  }

  if((!is.na(ffix[1]))&(ffixAll)){
    s1.state.sim<-t(apply(s1.state.sim,1,function(x){
      x[-idxN] <- x[-idxN]+log(ffix[1]/mean(exp(x[idxF[barAges]])));return(x)
    }))
    newf1<-ffix[1]/median(last.fbar.sim)
    fadjlocal<-c(newf1,fadj[2]/newf1,fadj[3]/fadj[2])
  }

  s1.ssb.sim<-exp(s1.state.sim[,idxN])%*%(ave.sw*ave.pm)
  s1.fbar.sim<-apply(s1.state.sim,1,function(x)mean(exp(x[idxF[barAges]])))
  s1.catch.sim<-apply(s1.state.sim,1,function(x)catch(x))
  s1.land.sim<-apply(s1.state.sim,1,function(x)catch.land(x))
  s1.fbarl.sim<-apply(s1.state.sim,1,function(x)fbar.land(x))

  # Step one Year 
  s2.state.sim<-t(apply(s1.state.sim,1,step, scale=fadjlocal[2]))
  s<-diag(exp(2.0*fit$est[fit$names=='logSdLogN'][idxNVar]))
  if(is.na(sampleRec)){
    m<-rep(0,length(idxN))
    s<-s[idxN,idxN]
    s2.state.sim[,idxN]<-s2.state.sim[,idxN]+mvrnorm(noSim,m,s)
  }else{
    m<-rep(0,length(idxN[-1]))
    s<-s[idxN[-1],idxN[-1]]
    s2.state.sim[,idxN[-1]]<-s2.state.sim[,idxN[-1]]+mvrnorm(noSim,m,s)
  }

  s2.state.sim<-t(apply(s2.state.sim,1,scaleN12, scale=c(bwscale[3],1)))


  if((!is.na(ffix[2]))&(!ffixAll)){
    fac<-ffix[2]/median(apply(s2.state.sim,1,function(x)mean(exp(x[idxF[barAges]]))))
    s2.state.sim<-t(apply(s2.state.sim,1,function(x){x[-idxN]<-x[-idxN]+log(fac);return(x)}))
    newf2<-ffix[2]/median(last.fbar.sim)
    fadjlocal[2:3]<-c(fadj[2]/newf2,fadj[3]/newf2)
  }

  if((!is.na(ffix[2]))&(ffixAll)){
    s2.state.sim<-t(apply(s2.state.sim,1,function(x){x[-idxN]<-x[-idxN]+log(ffix[2]/mean(exp(x[idxF[barAges]])));return(x)}))
    newf2<-ffix[2]/median(last.fbar.sim)
    fadjlocal[2:3]<-c(fadj[2]/newf2,fadj[3]/newf2)
  }

  s2.ssb.sim<-exp(s2.state.sim[,idxN])%*%(ave.sw*ave.pm)
  s2.fbar.sim<-apply(s2.state.sim,1,function(x)mean(exp(x[idxF[barAges]])))
  s2.catch.sim<-apply(s2.state.sim,1,function(x)catch(x))
  s2.land.sim<-apply(s2.state.sim,1,function(x)catch.land(x))
  s2.fbarl.sim<-apply(s2.state.sim,1,function(x)fbar.land(x))

  # Step one Year
  s3.state.sim<-t(apply(s2.state.sim,1,step, scale=fadjlocal[3]))
  
  s<-diag(exp(2.0*fit$est[fit$names=='logSdLogN'][idxNVar]))
  if(is.na(sampleRec)){
    m<-rep(0,length(idxN))
    s<-s[idxN,idxN]
    s3.state.sim[,idxN]<-s3.state.sim[,idxN]+mvrnorm(noSim,m,s)
  }else{
    m<-rep(0,length(idxN[-1]))
    s<-s[idxN[-1],idxN[-1]]
    s3.state.sim[,idxN[-1]]<-s3.state.sim[,idxN[-1]]+mvrnorm(noSim,m,s)
  }

  s3.state.sim<-t(apply(s3.state.sim,1,scaleN12, scale=c(bwscale[4],1)))


  if((!is.na(ffix[3]))&(!ffixAll)){
    fac<-ffix[3]/median(apply(s3.state.sim,1,function(x)mean(exp(x[idxF[barAges]]))))
    s3.state.sim<-t(apply(s3.state.sim,1,function(x){x[-idxN]<-x[-idxN]+log(fac);return(x)}))
  }

  if((!is.na(ffix[3]))&(ffixAll)){
    s3.state.sim<-t(apply(s3.state.sim,1,function(x){x[-idxN]<-x[-idxN]+log(ffix[3]/mean(exp(x[idxF[barAges]])));return(x)}))
  }

  s3.ssb.sim<-exp(s3.state.sim[,idxN])%*%(ave.sw*ave.pm)
  s3.fbar.sim<-apply(s3.state.sim,1,function(x)mean(exp(x[idxF[barAges]])))
  s3.catch.sim<-apply(s3.state.sim,1,function(x)catch(x))
  s3.land.sim<-apply(s3.state.sim,1,function(x)catch.land(x))
  s3.fbarl.sim<-apply(s3.state.sim,1,function(x)fbar.land(x))

  TAC<-function(x){
    F<-exp(x[idxF])
    N<-exp(x[idxN])
    Z<-F+ave.nm
    W<-ave.cw
    PHI<-ave.lf
    bias<-exp(fit$est[fit$names=='logScale'])
    aveBias<-mean(bias[length(bias)-(aveYears-1):0])
    sum(F/Z*W*PHI*N*(1-exp(-Z)))/aveBias
  }

  tac<-apply(s2.state.sim,1,function(x)TAC(x))
  
  return(list(last.state.sim=last.state.sim,last.ssb.sim=last.ssb.sim, last.fbar.sim=last.fbar.sim, 
         last.catch.sim=last.catch.sim, last.land.sim=last.land.sim, last.fbarl.sim=last.fbarl.sim,
         last.state=last.state, select=select, ave.pm=ave.pm, ave.sw=ave.sw, ave.nm=ave.nm, 
         s1.state.sim=s1.state.sim, s1.ssb.sim=s1.ssb.sim, s1.fbar.sim=s1.fbar.sim, s1.catch.sim=s1.catch.sim,
         s1.land.sim=s1.land.sim, s1.fbarl.sim=s1.fbarl.sim,  
         s2.state.sim=s2.state.sim, s2.ssb.sim=s2.ssb.sim, s2.fbar.sim=s2.fbar.sim, s2.catch.sim=s2.catch.sim,
         s2.land.sim=s2.land.sim, s2.fbarl.sim=s2.fbarl.sim,
         s3.state.sim=s3.state.sim, s3.ssb.sim=s3.ssb.sim, s3.fbar.sim=s3.fbar.sim, s3.catch.sim=s3.catch.sim,
         s3.land.sim=s3.land.sim, s3.fbarl.sim=s3.fbarl.sim, 
         last.year.used=last.year.used, fadj=fadj, tac.sim=tac)
        )
}

filen<-textConnection(gsub("^[[:blank:]]*#.*","",readLines("conf/forecast.cfg")))
oldtac<-scan(filen, quiet=TRUE, n=1)
foretab<-read.table(filen, header=FALSE)


do<-function(i){
  tabline<-as.numeric(foretab[i,-ncol(foretab)])
  doone(fit.current, stock.mean.weight, catch.mean.weight, prop.mature, natural.mortality, 
        fadj=tabline[1:3], fixNextYearTAC=tabline[4], fixNextYearTACAll=ifelse(tabline[5]==1,TRUE,FALSE), 
        ffix=tabline[6:8], ffixAll=ifelse(tabline[9]==1, TRUE, FALSE), bwscale=tabline[10:13])
}

forecast<-lapply(1:nrow(foretab),do)

save(forecast, foretab, oldtac, file='run/FORECAST.RData')