#various optimization functions

"blkdiag" <-
function (m1, m2, p.tr = 0, p.ll = 0)
{
    ## p.tr and p.ll are padding values
    topleft <- m1
    topright <- matrix(p.tr, nrow(m1), ncol(m2))
    colnames(topright) <- colnames(m2)
    lowleft <- matrix(p.ll, nrow(m2), ncol(m1))
    lowright <- m2
    rbind(cbind(topleft, topright), cbind(lowleft, lowright))
}

linfit<-function(A, q, gamma){
    require(Rglpk)
   #minimize ||Ax-q||_1 + gamma||DAx||_1
   #1^Tx=1
   #x>=0

    n<-nrow(P)
    p<-ncol(P)
    b<-rep(0,n)
    r<-length(q)
    obj<-c(rep(0,p), rep(1,r), rep(gamma,r-1))
    D<-diff(diag(r))
    ineqconst<-rbind(cbind(A, -diag(r), matrix(0,r,r-1)),
                     cbind(A, diag(r), matrix(0,r,r-1)),
                     cbind(D%*%A, matrix(0,r-1,r), -diag(r-1)),
                     cbind(D%*%A, matrix(0,r-1,r), diag(r-1)))
    const.mat <- ineqconst
    const.dir <- c(rep("<=",r), rep(">=",r), rep("<=",r), rep(">=",r))
    const.rhs <- c(1,q,q,rep(0,r),rep(0,r))
    model <- Rglpk_solve_LP(obj,const.mat,const.dir,const.rhs)
    list(weight=model$solution[1:p], model=model)
}

linfit1<-function(P, b, A, q, gamma){
   #minimize ||Ax-q||_1 + gamma||DAx||_1
   # s.t. Px = b
   #1^Tx=1
   #x>=0
    require(Rglpk)
    n<-nrow(P)
    p<-ncol(P)
    r<-length(q)
    obj<-c(rep(0,p),rep(1,r), rep(gamma,r-1))
    D<-diff(diag(r))
    linconst<-rbind(cbind(P, matrix(0,n,2*r-1)),
                    c(rep(1,p), rep(0,2*r-1)))
    ineqconst<-rbind(cbind(A,-diag(r), matrix(0,r,r-1)),
                     cbind(A,diag(r), matrix(0,r,r-1)),
                     cbind(D%*%A, matrix(0,r-1,r), -diag(r-1)),
                     cbind(D%*%A, matrix(0,r-1,r), diag(r-1)))
    const.mat<-rbind(linconst, ineqconst)
    const.dir<-c(rep("==",n+1), rep("<=",r), rep(">=",r), rep("<=",r), rep(">=",r))
    const.rhs<-c(b, 1, q, q, rep(0,r), rep(0,r))
    model <- Rglpk_solve_LP(obj, const.mat, const.dir, const.rhs)
    list(weight=model$solution[1:p], model=model)
}

linfit2<-function(P, A, q, gamma){
   #minimize ||Ax-p||_1 + gamma||DAx||_1
   # s.t. Px=0
   #1^Tx=1
   #x>=0
    require(Rglpk)
    n <- nrow(P)
    p <- ncol(P)
    b <- rep(0,n)
    r <- length(q)
    obj <- c(rep(0,p),rep(1,r), rep(gamma,r-2))
    D <- diff(diff(diag(r)))
    linconst <- rbind(cbind(P,matrix(0,n,2*r-2)),
                    c(rep(1,p),rep(0,2*r-2)))
    ineqconst <- rbind(cbind(A,-diag(r),matrix(0,r,r-2)),
                     cbind(A,diag(r),matrix(0,r,r-2)),
                     cbind(D%*%A,matrix(0,r-2,r),-diag(r-2)),
                     cbind(D%*%A,matrix(0,r-2,r),diag(r-2)))
    const.mat <- rbind(linconst,ineqconst)
    const.dir <- c(rep("==",n+1), rep("<=",r), rep(">=",r), rep("<=",r), rep(">=",r))#"
    const.rhs <- c(b,1,q,q,rep(0,r),rep(0,r))
    model <- Rglpk_solve_LP(obj, const.mat, const.dir, const.rhs)
    list(weight=model$solution[1:p], model=model)
}


minfit<-function(P){
	#minimize ||Px||_1
	#subject to x>=0
	#		1^Tx=1

	require(Rglpk)
	n<-nrow(P)
	p<-ncol(P)
	b<-rep(0,n)
	obj<-c(rep(0,p),1)
	linconst<-t(c(rep(1,p),0))
	ineqconst<-rbind(cbind(P,-rep(1,n)),
					 cbind(P,rep(1,n)))
	const.mat<-rbind(linconst,ineqconst)
	const.dir<-c("==",rep("<=",n),rep(">=",n))#"
	const.rhs<-c(1,rep(0,2*n))
	model <- Rglpk_solve_LP(obj,const.mat,const.dir,const.rhs)
	list(weight=model$solution[1:p],model=model)
}

quadfit<-function(A, P, q, b=rep(0,nrow(P))){
   #solves for x, min||Ax-q||^2
   #        	s.t. Px=b
   #		      x>=0
   #		      sum(x)=1
    alpha<-0.4
    beta<-0.8
    t<-1/alpha
    b<-c(b,1)
    P<-rbind(P,rep(1,ncol(P)))
    attach(svd(A))
    #init x and nu
    x<-rep(1,ncol(P))/ncol(P)
    nu<-rep(1,nrow(P))
    for( i in 1:10){
        t <- alpha*t
        decr <- Inf
        while(decr>1e-6){
            #solve the KKT equation
            rdual <- t(A)%*%A%*%x-2*t(t(q)%*%A)-t/x+t(P)%*%nu
            rprimal <- P%*%x-b
            #inverse of H=A^TA-t*diag(1/w^2)  via woodbury
            vainv <- sweep(t(v),2,x^2,"*")#"
            Hinv <- 1/t*diag(as.numeric(x)^2)-1/t^2*sweep(v,1,x^2,"*")%*%solve(diag(1/d)^2+1/t*vainv%*%v,vainv)#"
            S <- -P%*%Hinv%*%t(P)
            Dnu_pd <- solve(S,P%*%Hinv%*%rdual-rprimal)
            Dx_pd <- -Hinv%*%(t(P)%*%Dnu_pd+rdual)
            #backtracking search
            phi <- 1
            newx <- x+phi*Dx_pd
            newnu <- nu+phi*Dnu_pd
            newrdual <- t(A)%*%A%*%newx-2*t(t(q)%*%A)-t/newx+t(P)%*%newnu
            newrprimal <- P%*%newx-b

            while((sqrt(sum(newrdual^2)+sum(newrprimal^2))>(1-alpha*phi)*sqrt(sum(rdual^2)+sum(rprimal^2))) || (sum(newx<0)>0)){
                phi <- phi*beta
                newx <- x+phi*Dx_pd
                newnu <- nu+phi*Dnu_pd
                newrdual <- t(A)%*%A%*%newx-2*t(t(q)%*%A)-t/newx+t(P)%*%newnu
                newrprimal <- P%*%newx-b
            }
            x <- newx
            nu <- newnu
            decr <- sqrt(sum(newrdual^2)+sum(newrprimal^2))
            cat(decr,"\n")
        }
    }
    return(x)
}


KLfit<-function(P, q, b=rep(0,nrow(P))){
   #solves for x min KL(x,q)
   #        	s.t. Px=b
   #		      x>=0
   #		      sum(x)=1
    alpha <- 0.4
    beta <- 0.8
    b <- c(b,1)
    P <- rbind(P,rep(1,ncol(P)))
    #init x and nu
    x <- rep(1,ncol(P))/ncol(P)
    nu <- rep(1,nrow(P))
    decr <- Inf
    eps <- 1e-12
    niter <- 1
    q <- (q+eps)/sum(q+eps)
    while(decr>eps & niter<500){
        #solve the KKT equation
        rdual <- 1+log(x)-log(q)+t(P)%*%nu
        rprimal <- P%*%x-b
        S <- -P%*%diag(as.numeric(x))%*%t(P)
        Dnu_pd <- solve(S,P%*%diag(as.numeric(x))%*%rdual-rprimal)
        Dx_pd <- -diag(as.numeric(x))%*%(t(P)%*%Dnu_pd+rdual)
        #backtracking search
        phi <- 1
        newx <- x+phi*Dx_pd
        newnu <- nu+phi*Dnu_pd

        while(sum(newx<0)>0){
            phi <- phi*beta
            newx <- x+phi*Dx_pd
            newnu <- nu+phi*Dnu_pd
        }
        newrdual <- 1+log(newx)-log(q)+t(P)%*%newnu
        newrprimal <- P%*%newx-b
        while((sqrt(sum(newrdual^2)+sum(newrprimal^2))>(1-alpha*phi)*sqrt(sum(rdual^2)+sum(rprimal^2)))){
            phi <- phi*beta
            newx <- x+phi*Dx_pd
            newnu <- nu+phi*Dnu_pd
            newrdual <- 1+log(newx)-log(q)+t(P)%*%newnu
            newrprimal <- P%*%newx-b
        }
        x <- newx
        nu <- newnu
        decr <- sqrt(sum(newrdual^2)+sum(newrprimal^2))
        niter <- niter+1
        #cat(decr,"\n")
    }
    return(list(obj=sum(x*log(x/q)), weight=x, status=decr))
}

KLfitmod<-function(P,q,b=rep(0,nrow(P)),lambda,H,x=rep(1,ncol(P))/ncol(P)){
	#solves for x min KL(x,q)+lambda*<x,H^2>
	#        	s.t. Px=b
	#		      x>=0
	#		      sum(x)=1
    alpha<-0.4
    beta<-0.8
    b <- c(b,1)
    P <- rbind(P,rep(1,ncol(P)))
    #init x and nu
    nu <- rep(1,nrow(P))
    decr <- Inf
    t <- 1
    while(decr>1e-6){
        #solve the KKT equation
        rdual <- 1+log(x)-log(q)+lambda*H^2+t(P)%*%nu
        rprimal<-P%*%x-b
        S <- -P%*%diag(as.numeric(x))%*%t(P)
        Dnu_pd <- solve(S,P%*%diag(as.numeric(x))%*%rdual-rprimal)
        Dx_pd <- -diag(as.numeric(x))%*%(t(P)%*%Dnu_pd+rdual)
        #backtracking search
        phi<-1
        newx<-x+phi*Dx_pd
        newnu<-nu+phi*Dnu_pd

        while(sum(newx<0)>0){
            phi<-phi*beta
            newx<-x+phi*Dx_pd
            newnu<-nu+phi*Dnu_pd
        }
        newrdual<-1+log(newx)-log(q)+lambda*H^2+t(P)%*%newnu
        newrprimal<-P%*%newx-b
        while((sqrt(sum(newrdual^2)+sum(newrprimal^2))>(1-alpha*phi)*sqrt(sum(rdual^2)+sum(rprimal^2)))){
            phi<-phi*beta
            newx<-x+phi*Dx_pd
            newnu<-nu+phi*Dnu_pd
            newrdual<-1+log(newx)-log(q)+lambda*H^2+t(P)%*%newnu
            newrprimal<-P%*%newx-b
        }
        x<-newx
        nu<-newnu
        decr<-sqrt(sum(newrdual^2)+sum(newrprimal^2))
        decr<-sum(newrdual^2)
        #cat(decr,"\n")
    }
    return(x)
}

KLfit2<-function(P,q,b=rep(0,nrow(P))){
	#solves for x min KL(q,x)
	#        	s.t. Px=b
	#		      x>=0
	#		      sum(x)=1
	alpha<-0.4
	beta<-0.8
	b<-c(b,1)
	P<-rbind(P,rep(1,ncol(P)))
	#init x and nu
	x<-rep(1,ncol(P))/ncol(P)
	nu<-rep(1,nrow(P))
	decr<-Inf
	t<-1
	while(decr>1e-6){
		#solve the KKT equation
		rdual<--q/x+t(P)%*%nu
		rprimal<-P%*%x-b
		S <- -P%*%diag(as.numeric(x^2/q))%*%t(P)
		Dnu_pd <- solve(S,P%*%diag(as.numeric(x^2/q))%*%rdual-rprimal)
		Dx_pd<--diag(as.numeric(x^2/q))%*%(t(P)%*%Dnu_pd+rdual)
		#backtracking search
		phi<-1
		newx<-x+phi*Dx_pd
		newnu<-nu+phi*Dnu_pd

		while(sum(newx<0)>0){
			phi<-phi*beta
			newx<-x+phi*Dx_pd
			newnu<-nu+phi*Dnu_pd
		}
		newrdual<--q/newx+t(P)%*%newnu
		newrprimal<-P%*%newx-b
		while((sqrt(sum(newrdual^2)+sum(newrprimal^2))>(1-alpha*phi)*sqrt(sum(rdual^2)+sum(rprimal^2)))){
			phi<-phi*beta
			newx<-x+phi*Dx_pd
			newnu<-nu+phi*Dnu_pd
			newrdual<--q/newx+t(P)%*%newnu
			newrprimal<-P%*%newx-b
		}
		x<-newx
		nu<-newnu
		decr<-sqrt(sum(newrdual^2)+sum(newrprimal^2))
		decr<-sum(newrdual^2)
		cat(decr,"\n")
	}
	return(x)
}

solveKKT<-function(H,A,g,h,inv){
    ## solves the KKT system
    ## inv(H,v) should solve H^-1g
    S <- -t(A)%*%inv(H,A)
    if(ncol(A)==1){
        Dy <- 1/S*(h-t(A)%*%inv(H,g))
    }else{
        Dy <- solve(S,h-t(A)%*%inv(H,g))
    }
    Dx <- inv(H,g-A%*%Dy)
    return(list(Dx=as.vector(Dx),Dy=as.vector(Dy)))
}

lmconst<-function(Y,X,A,b,pos){
    ##solves for  \beta min ||Y-X\beta||
    ##        	s.t. A\beta = b
    ##               \beta>=0 (if pos = TRUE)
    alpha <- 0.4
    beta <- 0.8
    flag <- 1
    ##init x and nu
    x <- rep(1,ncol(A))
    nu <- rep(1,nrow(A))

    t <- pos
    XTX <- crossprod(X,X)
	while(t >1e-10 || flag){
            t <- alpha*t
            flag <- 0
            decr<-Inf
            niter<-1
            while(decr>1e-6 & niter<100){
                ## solve the KKT equation
                rdual <- -crossprod(X,Y-X%*%x) + crossprod(A,nu) - t/x
                rprimal <- A%*%x-b
                H <- XTX+t*diag(1/x^2)
                KKT <- solveKKT(H,t(A),-rdual,-rprimal,invchol)
                Dnu_pd <- KKT$Dy
                Dx_pd <- KKT$Dx
                ## backtracking search
                phi<-1
                newx<-x+phi*Dx_pd
                newnu<-nu+phi*Dnu_pd
                if( pos==1 ){
                    while(min(newx)<0){
                        phi<-phi*beta
                        newx<-x+phi*Dx_pd
                        newnu<-nu+phi*Dnu_pd
                    }
                }
                newrdual<--crossprod(X,Y-X%*%newx)+crossprod(A,newnu)-t/newx
                newrprimal<-A%*%newx-b
                while((sqrt(sum(newrdual^2)+sum(newrprimal^2))>(1-alpha*phi)*sqrt(sum(rdual^2)+sum(rprimal^2)))){
                    phi<-phi*beta
                    newx<-x+phi*Dx_pd
                    newnu<-nu+phi*Dnu_pd
                    newrdual<--crossprod(X,Y-X%*%newx)+crossprod(A,newnu)-t/newx
                    newrprimal<-A%*%newx-b
                }
                x<-newx
                nu<-newnu
                decr<-sqrt(sum(newrdual^2)+sum(newrprimal^2))
                ##cat(decr,"\n")
                niter<-niter+1
            }
	}
    return(list(obj=norm(Y-X%*%x),beta=x))
}

#cleanup cdf
cleancdf<-function(x,lambda){
    require(Rglpk)
    ##minimize for a ||a-x||_1+lambda*||diff(x)||_1
    ##subject to diff(x)>=0
    r<-length(x)
    obj<-c(rep(0,r),rep(1,r), rep(lambda,r-1))
    D<-diff(diag(r))
    const.mat<-rbind(cbind(diag(r),diag(r),matrix(0,r,r-1)),
                     cbind(D,matrix(0,r-1,r),matrix(0,r-1,r-1)),
                     cbind(D,matrix(0,r-1,r),-diag(r-1)),
                     cbind(diag(r),-diag(r),matrix(0,r,r-1) ))

    const.dir<-c(rep(">=",2*r-1),rep("<=",2*r-1))#"
    const.rhs<-c(x,rep(0,r-1),rep(0,r-1),x)
    model <- Rglpk_solve_LP(obj,const.mat,const.dir,const.rhs)
    model$solution[1:r]
}

posls<-function(Y,X,w=rep(1,nrow(X))){
    ##minimize ||Y-X\beta||_w
    ##subject to beta>=0
    beta <- rep(1,ncol(X))
    t <- 1
    phi <- 0.6
    for(i in 1:30){
        dbeta<-rep(Inf,ncol(X))
        while(norm(dbeta)>=1e-10){
            g<--crossprod(X,w*(Y-X%*%beta))-t/beta
            H<-crossprod(X,sweep(X,1,w,"*"))+diag(as.vector(t/beta^2))
            dbeta<-solve(H,g)
            while(min(beta-dbeta)<0){
                dbeta<-phi*dbeta
            }
            beta<-beta-dbeta
            cat(norm(dbeta),"\n")
		}
        t<-phi*t
    }
    return( beta )
}

invchol <- function(H,b){
    R <- chol(H)
    z <- forwardsolve(t(R),b)
    x <- backsolve(R,z)
    return(x)
}

#solve the KKT equation by block elimination
solveswm <- function(A,B,C,b){
    ##solve Dx=b for D=A+BC with A diagonal and B and C have low ranks
    z <- 1/A*b
    E <- diag(nrow(C))+C%*%sweep(B,1,A,"/")#"
    w <- solve(E,C%*%z)
    x <- z-sweep(B%*%w,1,A,"/")
    return( x )
}

norm <- function(x){
  sqrt(sum(x^2))
}