# Copyright (c) 2014  Takafumi Kanamori [kanamori@is.nagoya-u.ac.jp]
# All rights reserved.  See the file COPYING for license terms.

#####################################################
## robust estimation with unnormalized Gaussian model
#####################################################
## x: n(sample size) * d(dimension) data matrix
est_UnnormalizedModel <- function(x,gamma=0.5,tol=10^(-8)){
  if(is.vector(x)) x <- matrix(x)
  dp_sol    <- est_densitypower(gamma=gamma,x=x,tol=tol)
  gamma_sol <- est_sphere(gamma=gamma,x=x,tol=tol)
  par_mu <- gamma_sol$mu; par_var <- gamma_sol$var
  sx <- t(x)-par_mu
  w <- as.vector(exp(-gamma*colSums(solve(par_var,sx)*sx)/2))
  c_est <- (1+gamma)^(ncol(x)/2)*mean(w)
  if (c_est>=1){
    sx <- t(x)-dp_sol$mu
    p <- as.vector(exp(-colSums(solve(dp_sol$var,sx)*sx)/2)/sqrt((2*pi)^(ncol(x))*det(dp_sol$var)))
    sol <- list(mu=dp_sol$mu, var=dp_sol$var, c=1, gamma=gamma, obj=dp_sol$obj, p=p, empirical=list(mu=colMeans(x),var=var(x)))
  }
  else{
    sx <- t(x)-gamma_sol$mu
    p <- as.vector(exp(-colSums(solve(gamma_sol$var,sx)*sx)/2)/sqrt((2*pi)^(ncol(x))*det(gamma_sol$var)))
    sol <- list(mu=gamma_sol$mu, var=gamma_sol$var, c=c_est, gamma=gamma, obj=gamma_sol$obj, p=p, empirical=list(mu=colMeans(x),var=var(x)))
  }
  sol
}

########################################
## regression with location scale models
########################################
## x: n(sample size) * d(dimension) data matrix of covariate variable
## y: n observations of response variable
reg_unnormalizedLSmodel <- function(x,y,gamma=1,multiple_start=10,maxit=10000){
  rlm_sol <- lmrob(y~.,data=data.frame(x,y))
  par <- c(mad(as.vector(y-predict(rlm_sol))), rlm_sol$coef)
  ini_pars <- t(matrix(par))
   if(multiple_start>0)
     ini_pars <- rbind(par,cbind(par[1]*runif(multiple_start,min=0.3,max=3), rep(1,multiple_start) %o% par[-1]))
  opt_val <- Inf
  for(j in 1:nrow(ini_pars)){
    tmp_sol <- optim(ini_pars[j,], fn=loss_gamma, gr=gr_gamma, method="BFGS", control=list(maxit=maxit), gamma=gamma, x=cbind(1,x), y=y)
    if(tmp_sol$value < opt_val){
      opt_val <- tmp_sol$value
      sol_gamma <- tmp_sol
    }
  }
  residual <- as.vector(y-cbind(1,x) %*% sol_gamma$par[-1])
  c_est <- sqrt(1+gamma)*mean(exp(-(gamma*(residual)^2)/(2*(sol_gamma$par[1])^2)))
  if (c_est<1){
    sol <- list(par=sol_gamma$par[-1], sd=abs(sol_gamma$par[1]), contam_est=1-c_est, c=c_est, gamma=gamma, lmrob_MM=list(par=par[-1],sd=par[1]))
  }else{
    soldp <- optim(par, fn=loss_densitypower, gr=gr_densitypower, method="BFGS", control=list(maxit=maxit), gamma=gamma, x=cbind(1,x), y=y)
    sol <- list(par=soldp$par[-1],     sd=abs(soldp$par[1]),     contam_est=0,       c=1,     gamma=gamma, lmrob_MM=list(par=par[-1],sd=par[1]))
  }
  sol
}


########################################################################
## functions used in est_UnnormalizedModel and reg_unnormalizedLSmodel
########################################################################

## x: data matrix (n*p)
est_densitypower <- function(gamma=0.5,x,tol=10^(-8)){
  if(is.vector(x)) x <- matrix(x)
  k <- ncol(x)
  par_mu  <- apply(x,2,median)
  par_var <- matrix(0,ncol(x),ncol(x))
  diag(par_var) <- apply(x,2,mad)^2
  tolidx <- FALSE
  obj_val <- c()
  while(tolidx == FALSE){
    sx <- t(x)-par_mu
    w <- as.vector(exp(-gamma*colSums(solve(par_var,sx)*sx)/2))
    update_mu <- colSums(x*w)/sum(w)
    sx <- t(x)-update_mu
    update_var <- sx%*%(t(sx)*w)/sum(w) + gamma*(1+gamma)^(-k/2-1)/mean(w) * par_var
    diff <- mean(abs(par_mu-update_mu))+mean(abs(par_var-update_var))
    par_mu <- update_mu; par_var <- update_var
    tolidx <- (diff < tol)
    ##
    sx <- t(x)-par_mu
    p <- (1/sqrt((2*pi)^k *det(par_var)))*exp(-colSums(solve(par_var,sx)*sx)/2)
    obj <- -mean(p^gamma)/gamma+(((2*pi)^gamma * (1+gamma))^(-k/2) * (det(par_var))^(-gamma/2))/(1+gamma)
    obj_val <- c(obj_val, obj)
  }
  list(mu=par_mu, var=par_var, gamma=gamma, obj=obj_val, empirical=list(mu=colMeans(x),var=var(x)))
}

## x: data matrix (n*p)
est_sphere <- function(gamma=0.5,x,tol=10^(-8)){
  if(is.vector(x)) x <- matrix(x)
  k <- ncol(x)
  par_mu  <- apply(x,2,median)
  par_var <- matrix(0,ncol(x),ncol(x))
  diag(par_var) <- apply(x,2,mad)^2
  tolidx <- FALSE
  obj_val <- c()
  while(tolidx == FALSE){
    sx <- t(x)-par_mu
    w <- as.vector(exp(-gamma*colSums(solve(par_var,sx)*sx)/2))
    update_mu <- colSums(x*w)/sum(w)
    update_var <- (1+gamma)*((t(x)%*%(x*w))/sum(w)-outer(update_mu,update_mu))
    diff <- mean(abs(par_mu-update_mu))+mean(abs(par_var-update_var))
    par_mu <- update_mu; par_var <- update_var
    tolidx <- (diff < tol)
    ##
    sx <- t(x)-par_mu
    p <- (1/sqrt((2*pi)^k * det(par_var)))*exp(-colSums(solve(par_var,sx)*sx)/2)
    obj <- -(log(mean(p^gamma)))/gamma-(log(((2*pi)^gamma * (1+gamma))^(k/2) * (det(par_var))^(gamma/2)))/(1+gamma)
    obj_val <- c(obj_val, obj)
  }
  list(mu=par_mu, var=par_var, gamma=gamma, obj=obj_val, empirical=list(mu=colMeans(x),var=var(x)))
}

## par=(sigma, beta), beta= (ncol(x)+1)-vec
loss_densitypower <- function(par,gamma,x,y,...){
  v <- (par[1])^2; beta <- par[-1];
  pa <- (1/sqrt(2*pi*v))^gamma * exp(-(gamma*(y-x%*%beta)^2)/(2*v))
  int_pa <- ((2*pi)^gamma*(1+gamma))^(-1/2) * v^(-gamma/2)
  -mean(pa)/gamma + int_pa/(1+gamma)
}

## par=(sigma, beta), beta= (ncol(x)+1)-vec
gr_densitypower <- function(par,gamma,x,y,...){
  v <- (par[1])^2; beta <- par[-1];
  pa <- (1/sqrt(2*pi*v))^gamma * exp(-(gamma*(y-x%*%beta)^2)/(2*v))
  gr_beta <- -colMeans(x * as.vector(pa*as.vector(y-x%*%beta)))/v
  gr_v <- -(mean(pa * ((y-x%*%beta)^2-v)))/(2*v^2)-gamma/(2*(1+gamma)^(3/2)) * (2*pi)^(-gamma/2) * v^(-gamma/2-1)
  2*par[1]*c(gr_v, gr_beta)
}

## par=(sigma, beta), beta= (ncol(x)+1)-vec
loss_gamma <- function(par,gamma,x,y,...){
  v <- (par[1])^2; beta <- par[-1];
  pa <- (1/sqrt(2*pi*v))^gamma * exp(-(gamma*(y-x%*%beta)^2)/(2*v))
  int_pa <- ((2*pi)^gamma*(1+gamma))^(-1/2) * v^(-gamma/2)
  log(int_pa)/(1+gamma)-log(mean(pa))/gamma
}

## par=(sigma, beta), beta= (ncol(x)+1)-vec
gr_gamma <- function(par,gamma,x,y,...){
  v <- (par[1])^2; beta <- par[-1];
  pa <- (1/sqrt(2*pi*v))^gamma * exp(-(gamma*(y-x%*%beta)^2)/(2*v))
  int_pa  <- ((2*pi)^gamma*(1+gamma))^(-1/2) * v^(-gamma/2)
  gr_beta <- -colMeans(x * as.vector(pa*as.vector(y-x%*%beta)))/(v*mean(pa))
  gr_v    <- -gamma/(2*(1+gamma)*v) - mean(pa*((y-x%*%beta)^2-v))/(2*v^2*mean(pa))
  2*par[1]*c(gr_v, gr_beta)
}

## Geman-McClure estimator for regression 
est_GemMc <- function(x,y,maxit=10000,rlm_maxit=1000,...){
  rlm_sol <- lmrob(y~.,data=data.frame(x,y))
  sol <- optim(rlm_sol$coef, fn=loss_GemMc, gr=gr_GemMc, method="BFGS", control=list(maxit=maxit),x=cbind(1,x), y=y)
  sol$par
}
loss_GemMc <- function(beta,x,y,...){
  r <- as.vector(y-x%*%beta)
  sum(r^2/(1+r^2))  
}
gr_GemMc <- function(beta,x,y,...){
  r <- as.vector(y-x%*%beta)
  colSums(-x*(2*r/(1+r^2)^2))
}