# R functions for ESN distributions
# Author: A.Azzalini (Universita` di Padova):
#
# Original S-plus code: July-October 1999,  ported to R in October 2007,
# amedend in February 2010;  public distribution in March 2011
#
# This code is freely usable under GPL license available from
# http://www.gnu.org/licenses/gpl.html
#                                 AND 
# the promise that you do not complain about the lack of documentation,
# being content to look at the documentation for the matching routines 
# (without the "e") in the "sn" package available from CRAN:
# http://cran.r-project.org/web/packages/sn/index.html
# while for the extra parameter "tau" it is your responsibility to read
# the published papers, specifically Capitanio et al. (2003), cf
# http://azzalini.stat.unipd.it/SN/abstracts.html
#---------------------------------------------------------------


rmesn <- function(n=1, xi=rep(0,k), Omega, alpha, tau=0)
{#  generates SN_k(xi,Omega,alpha,tau) variates using conditioning method
  k <- ncol(Omega)
  Z <- mesn.quantities(xi,Omega,alpha,tau)
  O.star <- rbind(c(1,Z$delta),cbind(Z$delta,Z$Omega.cor))
  z<- norm.censor(n,k,chol(O.star),tau)
  y <- t(xi+Z$omega*t(z))
  return(y)
  }

norm.censor <- function(n, k, root.cov, tau)
{# routine di supporto per rmesn (ma non si puo` inserire al suo interno)
    if(n<=0 | k<1) return()
    if(tau>0) eff <- pnorm(tau) else eff <- 2*pnorm(tau)
    n0 <- round(n/eff)
    x <- matrix(rnorm(n0*(k+1)),n0,(k+1)) %*% root.cov
    # each row of x is N_{k+1}(0,Omega*)
    id<- (x[,1]+tau>0)
    z <- x[id,-1,drop=F]
    if(tau<=0) {
      id<- (-x[,1]+tau>=0)
      z <- rbind(z,-x[id,-1,drop=F])
      }
    if(is.null(dim(z))) n.sampled<-0 else n.sampled<-nrow(z)    
    if(n.sampled<n) 
      z <- rbind(z, norm.censor(n - n.sampled, k, root.cov, tau))	
    else
      z <- z[1:n,]
    z 
}

dmesn <- function(x, xi=rep(0,k), Omega, alpha, tau=0)
{
# Density of Multivariate SN rv with parameters (xi, Omega, alpha, tau) 
# evaluated at x, which is either a k-vector or n x k matrix
  scale <- sqrt(diag(Omega))
  if(is.vector(x)) {
    n <- 1
    k <- length(x)
  }
  else {
    n <- dim(x)[1]
    k <- dim(x)[2]
  }
  X <- t(matrix(x, nrow = n, ncol = k)) - xi
  z <- X/scale
  eta <- alpha/scale
  alpha0 <- tau * sqrt(1+sum(sum(eta*(Omega%*%eta))))
  Q <- apply((solve(Omega) %*% X) * X, 2, sum)
  # diagonal of (x Omega^(-1) x^T)
  # Det <- as.numeric(det.Hermitian(as.Matrix(Omega), logarithm = F)$modulus)
  Det <- abs(prod(diag(qr(Omega)$qr)))
  pdf <- (exp(-0.5 * Q) * pnorm(alpha0 + t(z) %*% as.matrix(alpha)))/
         (sqrt((2 * pi)^k * Det)* pnorm(tau))
  pdf
}


plot.desn2 <- function(x, y, xi, Omega, alpha, tau, ...)
{
        # plot bivariate density SN_2(xi,Omega,alpha) computed at (x,y) grid
        if(any(dim(Omega) != c(2, 2))) stop("dim(Omega) != c(2,2)")
        nx <- length(x)
        ny <- length(y)
        xoy <- cbind(rep(x, ny), as.vector(matrix(y, nx, ny, byrow = T)))
        X <- as.matrix(xoy, nx * ny, 2, byrow = F)
        pdf <- dmesn(X, xi, Omega, alpha, tau)
        pdf <- matrix(pdf, nx, ny)
        contour(x, y, pdf, ...)
        invisible(list(x = x, y = y, density = pdf, xi = xi, Omega = Omega,
                alpha = alpha))
}

#----------------------------------------

mesn.mle <- function(X, y, freq,  start, fixed.tau=NA, trace=F,
                     control=list(iter.max=150, x.tol=1e-10))
{
  y <- as.matrix(y)
  if(missing(X)) X <- rep(1,nrow(y))
  if(missing(freq)) freq <- rep(1,nrow(y))
  X <- as.matrix(X) 
  k <- ncol(y)  
  n <- sum(freq)
  m <- ncol(X)
  y.names<-dimnames(y)[[2]] 
  x.names<-dimnames(X)[[2]]
  if(missing(start)){
    qrX   <- qr(X)
    beta  <- qr.coef(qrX,y)
    a     <- msn.moment.fit(qr.resid(qrX,y))
    Omega <- a$Omega
    omega <- a$omega
    alpha <- a$alpha
    if(!a$admissible) alpha<-alpha/(1+max(abs(alpha)))
    beta[1,] <- beta[1,]-omega*a$delta*sqrt(2/pi)  
    tau <- 0
    }
  else {
    beta  <- start$beta
    Omega <- start$Omega
    alpha <- start$alpha
    tau   <- start$tau
  } 
  eta<-alpha/sqrt(diag(Omega))
  Oinv<-solve(Omega)
  Oinv<-(Oinv+t(Oinv))/2
  # a<-chol2(Oinv)
  upper<-chol(Oinv)
  d<-diag(upper)
  tA<-t(upper/d)
  d <- d^2
  if(is.na(fixed.tau))
    param<-c(beta,eta,log(d),tA[lower.tri(tA)],tau) 
  else
    param<-c(beta,eta,log(d),tA[lower.tri(tA)])
  opt <- nlminb(param, mesn.dev, mesn.dev.grad, control=control,
              X=X, y=y, freq=freq, trace=trace, fixed.tau=fixed.tau)
  dev <- opt$obj
  cat("nlminb message:", opt$message,"\n")
  cat("deviance:", dev,"\n")
  param<-opt$par
  beta<- matrix(param[1:(m*k)],m,k)
  eta <- param[(m*k+1):(m*k+k)]
  d <- exp(param[(m*k+k+1):(m*k+2*k)])
  tA <- diag(k)
  tA[lower.tri(tA)] <- param[(m*k+2*k+1):(m*k+2*k+k*(k-1)/2)]
  if(is.na(fixed.tau))  
    tau <- param[m*k+2*k+k*(k-1)/2+1]
  else 
    tau<-fixed.tau
  A <- t(tA)
  Omega <- solve(tA %*% diag(d) %*% A)
  alpha<- eta*sqrt(diag(Omega))
  dimnames(beta) <- list(x.names,y.names)
  dimnames(Omega) <- list(y.names, y.names)
  names(alpha) <- y.names
  dp <- list(beta=beta, Omega=Omega,  alpha=alpha, tau=tau)
  list(call=match.call(), logL=-0.5*dev, deviance=dev, dp=dp, opt=opt)
}

mesn.logL.profile <- function(X, y, freq, tau=seq(-3,3,length=31), trace=F)
{
  y <- as.matrix(y)
  if(missing(X)) X <- rep(1,nrow(y))
  if(missing(freq)) freq <- rep(1,nrow(y))
  logL<-numeric(length(tau))
  X <- as.matrix(X)
  # fit <- lm.fit.qr(X,y)
  # beta  <- coef(fit)
  qrX  <- qr(X) 
  beta <- qr.coef(qrX, y)
  a     <- msn.moment.fit(qr.resid(qrX,y))
  Omega <- a$Omega
  omega <- a$omega
  alpha <- a$alpha
  if(!a$admissible) alpha<-alpha/(1+max(abs(alpha)))
  beta[1,] <- beta[1,]-omega*a$delta*sqrt(2/pi)  
  start<-list(beta=beta,Omega=Omega,alpha=alpha)
  for(ip in 1:length(tau)){
    cat("tau: ",tau[ip],"\n")
    mle.tau <- mesn.mle(X=X, y=y, freq=freq, start= start, fixed.tau=tau[ip], 
                 trace=trace)
    logL[ip] <- mle.tau$logL
    start<- mle.tau$dp
    }
  plot(tau,logL,type="l")
  invisible(list(tau=tau,logL=logL))
}

mesn.dev <- function(param, X, y, freq, trace=F, fixed.tau)
{
  k <- ncol(y)
  # if(missing(freq)) freq<-rep(1,nrow(y))
  n <- sum(freq)
  m <- ncol(X)
  beta<-matrix(param[1:(m*k)],m,k)
  eta <-param[(m*k+1):(m*k+k)]
  d <- exp(param[(m*k+k+1):(m*k+2*k)])
  tA <- diag(k)
  tA[lower.tri(tA)] <- param[(m*k+2*k+1):(m*k+2*k+k*(k-1)/2)]
  if(is.na(fixed.tau))  tau <-param[m*k+2*k+k*(k-1)/2+1]
  else tau<-fixed.tau
  A <- t(tA)
  Oinv <- tA %*% diag(d) %*% A
  Omega <- solve(Oinv)
  u <-  y - X %*% beta
  UU <- t(u) %*% (u*freq)
  trace.UU.Oinv <- sum(diag(UU %*% Oinv))
  v <-  as.vector(u %*% eta +tau*sqrt(1+sum(eta*as.vector(Omega%*%eta))))
  logDet<- sum(log(2*pi/d))
  dev <- n*logDet+trace.UU.Oinv-2*sum(zeta(0,v)*freq)+2*n*zeta(0,tau)
  if(trace) cat("mesn.dev: ",dev,"\n")
  dev
}


mesn.dev.grad <- function(param, X, y, freq, trace=F, fixed.tau=NA)
{
  k <- ncol(y)
  # if(missing(freq)) freq<-rep(1,nrow(y))
  n   <- sum(freq)
  m   <- ncol(X)
  beta<-matrix(param[1:(m*k)],m,k)
  eta <-param[(m*k+1):(m*k+k)]
  d   <- exp(param[(m*k+k+1):(m*k+2*k)])
  tA  <- diag(k)
  tA[lower.tri(tA)] <- param[(m*k+2*k+1):(m*k+2*k+k*(k-1)/2)] 
  if(is.na(fixed.tau))  tau <-param[m*k+2*k+k*(k-1)/2+1]
  else tau<-fixed.tau
  A     <- t(tA)
  Ai    <- solve(A)
  tAi   <- t(Ai)
  Oinv  <- tA %*% diag(d) %*% A
  Omega <- solve(Oinv)
  Omega.eta <- as.vector(Omega%*%eta)
  eOe   <- sum(eta*Omega.eta)
  u     <- y-X %*% beta
  UU    <- (t(u)%*%(u*freq))
  v     <- as.vector(u %*% eta +tau*sqrt(1+eOe))
  p1    <- zeta(1,v)
  Dbeta <- t(X) %*% (u*freq) %*% Oinv - outer(as.vector(t(X*freq)%*% p1), eta)
  Deta  <- as.vector((t(u)+tau*Omega.eta/sqrt(1+eOe)) %*% (p1*freq))
  Dtau  <- sum(p1*freq)*sqrt(1+eOe) - n*zeta(1,tau)
  Dzeta.eOe <- 0.5*sum(p1*freq)*tau/sqrt(1+eOe)
  DeOe.d<- -diag(diag(1/d) %*% tAi %*% outer(eta,eta) %*% Ai %*% diag(1/d))
  Dd    <- 0.5*n/d - 0.5*diag(A %*% UU %*% tA) + Dzeta.eOe * DeOe.d
  M     <- t(-2*t(Ai) %*% outer(eta,eta) %*% Omega)
  DeOe.A<- M[lower.tri(M)]
  R     <- t(diag(d) %*% A %*% UU)
  DA    <-  -R[lower.tri(R)] + Dzeta.eOe * DeOe.A
  if(is.na(fixed.tau)) grad  <- (-2)*c(Dbeta,Deta,Dd*d,DA,Dtau)
  else grad  <- (-2)*c(Dbeta,Deta,Dd*d,DA)
  if(trace) {
     cat("mesn.dev.grad: norm is ",sqrt(sum(grad^2)),"\n")  
     # browser()
     }
  return(grad)
}

# -----------------------------------------------------

msn.moment.fit <- function(y)
{
  # 31-12-1997: simple fit of MSN distribution usign moments
  #  cat("MLE con metodo momenti\n")
  Diag<-function(d) if(length(d)>1) diag(d) else d
  y <- as.matrix(y)
  k <- ncol(y)
  m.y <- apply(y, 2, mean)
  var.y <- var(y)
  y0<- (t(y) - outer(m.y, rep(1,nrow(y))))/sqrt(diag(var.y))
  gamma1 <- apply(y0^3, 1, mean)
  out <- (abs(gamma1) > 0.99527)
  gamma1[out] <- sign(gamma1[out]) * 0.995
  a <- sign(gamma1) * ((2 * abs(gamma1))/(4 - pi))^0.33333
  delta <- (sqrt(pi/2) * a)/sqrt(1 + a^2)
  m.z <- delta * sqrt(2/pi)
  omega <- sqrt(diag(var.y)/(1 - m.z^2))
  Omega <- var.y + outer(omega * m.z, omega * m.z)
  xi <- m.y - omega * m.z
  O.cor <- Diag(1/omega) %*% Omega %*% Diag(1/omega)
  O.cor <- (t(O.cor) + O.cor)/2
  O.inv <- solve(O.cor)
  tmp <- as.vector(1 - t(delta) %*% O.inv %*% delta)
  if(tmp <= 0) {
          tmp <- 0.0001
          admissible <- F
  }
  else admissible <- T
  alpha <- as.vector(O.inv %*% delta)/sqrt(tmp)
  list(xi = xi, Omega = Omega, alpha = alpha, Omega.cor = O.cor, omega =
          omega, delta = delta, skewness = gamma1, admissible = 
          admissible)
}
 

mesn.quantities <- function(xi, Omega, alpha, tau=0)
{
  # 21-12/1997; computes various quantieies related to SN_k(xi,Omega,alpha)
  k <- length(alpha)
  Diag<-function(d) if(length(d)>1) diag(d) else d
  if(length(xi) != k | any(dim(Omega) != c(k, k)))
          stop("dimensions of arguments do not match")
  omega <- sqrt(diag(Omega))
  O.cor <- diag(1/omega) %*% Omega %*% diag(1/omega)
  tmp <- as.vector(sqrt(1 + t(as.matrix(alpha)) %*% O.cor %*% alpha))
  delta <- as.vector(O.cor %*% alpha)/tmp
  lambda <- delta/sqrt(1 - delta^2)
  D <- Diag(sqrt(1 + lambda^2))
  if(tau==0) {
    Psi <- D %*% (O.cor - outer(delta, delta)) %*% D
    Psi <- (Psi + t(Psi))/2 }
  else
    Psi <- NULL
  O.inv <- solve(Omega)
  oi <- sqrt(diag(O.inv))
  O.conc <- Diag(1/oi) %*% (-O.inv) %*% Diag(1/oi)
  diag(O.conc) <- rep(1, k)
  muZ <- delta * zeta(1,tau)
  muY <- xi + omega * muZ
  # Sigma <- diag(omega) %*% (O.cor - outer(muZ, muZ)) %*% diag(omega)
  Sigma <- Omega + zeta(2,tau) * outer(omega*delta, omega*delta)
  Sigma <- (Sigma + t(Sigma))/2
  cv <- muZ/sqrt(1 - muZ^2)
  # gamma1 <- 0.5 * (4 - pi) * cv^3
  gamma1 <- NULL
  list(xi = xi, Omega = Omega, alpha = alpha, tau=tau,omega = omega,
         mean = muY, variance = Sigma, Omega.cor = O.cor, Omega.conc = O.conc, 
         Psi = Psi, lambda = lambda, delta = delta, skewness = gamma1)
}


zeta <- function(k, x)
{
        # k integer \in (0,4)
        k <- as.integer(k)
        na <- is.na(x)
        x <- replace(x, na, 0)
        if(any(abs(x) == Inf))
                stop("Inf not allowed")
        # funzionerebbe per k=0 e 1, ma non per k>1
        ok <- (-30 < x)
        if(k == 0) {
                ax <- ( - x[!ok])
                ay <- (-0.918938533204673) - 0.5 * ax^2 - log(ax)
                y <- rep(NA, length(x))
                y[ok] <- log(2 * pnorm(x[ok]))
                y[!ok] <- ay
        }
        else {
                if(k == 1) {
                        y <- ( - x) * (1 + 1/x^2)
                        y[ok] <- dnorm(x[ok])/pnorm(x[ok])
                }
                else {
                        if(k == 2)
                                y <- ( - zeta(1, x) * (x + zeta(1, x)))
                        else {
                                if(k == 3)
                                        y <- ( - zeta(2, x) * (x + zeta(1,
                                                x)) - zeta(1, x) * (1 + zeta(
                                                2, x)))
                                else {
                                        if(k == 4)
                                                y <- ( - zeta(3, x) * (x + 2 *
                                                        zeta(1, x)) - 2 * zeta(
                                                        2, x) * (1 + zeta(
                                                        2, x)))
                                        else {
                                                warning("k>4")
                                                y <- rep(NA, length(x))
                                        }
                                }
                        }
                }
        }
        replace(y, na, NA)
}