## R function to read .sfs file

readBjoint <- function(file=NULL,nind=10,nsites=10){
  ff <- gzfile(file,"rb")
  m<-matrix(readBin(ff,double(),(2*nind+1)*nsites),ncol=(2*nind+1),byrow=TRUE)
  close(ff)
  return(m)
}

##normalizer
norm <- function(x) x/sum(x)

##we have 2 function to estimate the SFS from the .sfs file
##bfgs using the unconstraint 2*nChr+1 (doesn't work well, in the ccode we are shifting to dim-1 space)

##em algorithm

##pick random start in constraint space  
pick <- function(nchr)  return(0.01/(1+0:(nchr-1)))

## go from full to dim-1 space
swap <- function(x){
  tmp <- 1+sum(x)
  return(c(1/tmp,x/tmp))
}



## unconstraind bfgs
myOptim <- function(ind=12,theLikes=dat){
  like2 <- function(x){
    if(length(x)!=ncol(theLikes))
      stop("Dimension mismatch")
    res <- exp(theLikes) %*% x
    res <- sum(log(res))
    return(-res)
  }

  neg_fn <- function(x,calcLike=T){
    return(like2(swap(x)))
  }
  ti <- system.time(opt <- optim(pick(ind*2),neg_fn,method="L-BFGS-B",lower=rep(0.0000001,ind*2),upper=rep(10,ind*2)))
  return(list(time=ti,opt=opt,ml=swap(opt$par)))
}

  
## em old
myOptim2 <- function(start,nIter=10,tol=0.1,theLikes){
  ind<-ncol(theLikes)
  if(missing(start))
    start<-norm(rep(1,ind))
  x<-start
  like<--Inf
  for(i in 1:nIter){
    post<-t(t(exp(theLikes))*x)
    norma<-rowSums(post)
    likeNew<-sum(log(norma))
    pp<-post/norma
    x<-colSums(pp)/nrow(theLikes)
    if((likeNew-like)<tol)
      break;
    like<-likeNew
    }
  cat("iter ",i,"\n")
  x
}

if(FALSE){
  ##these files are generated by
  ## cd data; ./supersim.g++ -npop 1 -nind 12 -outfiles data/pop1 -pvar 1 -nsites 100000;
  ## cd ..;./dirty.g++ -sim1 misc/data/pop1.glf.gz -nInd 12 -realSFS 1 -outfiles misc/data/res
  ## cd data;./optimSFS.gcc -binput data/res.sfs -ncat 24 -nThread 12
  
  dat <- readBjoint(file="../misc/data/res.sfs",nind=12,nsites=100000)
  ##  bfgs <- myOptim(ind=12,theLikes=dat)

  #fairly fast, but now always correct (bfgs is always correct)
  emR <- myOptim2(theLikes=dat)

  myTrue <- scan("../misc/data/pop1.frq")
  bfgsC <- scan("../misc/data/res.sfs.ml")
  opts <- rbind(myTrue,emR,bfgsC)
  optsLik <- apply(opts,1,function(x) sum(log(rowSums(t(exp(dat)) *x))))
  par(mfcol=c(1,2))
  barplot(opts,col=1:3)
  legend("topright",paste(names(optsLik),round(optsLik,3),sep="="),fill=1:3)
  barplot(opts[,-1],col=1:3)
  legend("topright",paste(names(optsLik),round(optsLik,3),sep="="),fill=1:3)
  
}

projectOld <- function(counts,newDim=10) {##newDim is newNumChr
  oldDim <- length(counts)-1 ##numChr

  fsTsk2 <- c()

  for(t in 0:newDim){
    tmp <- c()
    for(tal in seq_along(counts)-1)
      tmp[tal+1] <- counts[tal+1] * choose(tal,t) * choose(oldDim-tal,newDim-t)
    fsTsk2[t+1] <- sum(tmp)
  }

  return(fsTsk2/choose(oldDim,newDim))
  
}




if(FALSE){
  ##try projecting
  project<-function(newNchr,theLikes){
    oldNchr <- ncol(theLikes)-1
    mat <- sapply(0:newNchr,function(i) choose(0:oldNchr,i)*choose(oldNchr-0:oldNchr,newNchr-i)/choose(oldNchr,newNchr))
    ##    return(mat)
    return(exp(theLikes)%*%mat)
  }

  pdf("temp.pdf",width=24,height=24)
  par(mfrow=c(4,2))
  for(i in c(23,20,15,10)){
    newdat <- project(i,theLikes=dat)
    newres <- myOptim2(theLikes=log(newdat))
    barplot(newres,main=paste("downsampling to: ",i,": keeping p0" ))
    barplot(newres[-1],main=paste("downsampling to: ",i,": removing p0" ))
  }
}