#' The bivariate Wrapped Normal distribution
#' @inheritParams rvmsin
#' @inheritParams rwnorm
#' @param int.displ integer displacement. If \code{int.displ =} M, then each infinite sum in the
#' density is approximated by a finite sum over 2*M + 1 elements. (See Details.) The allowed values are 1, 2, 3, 4 and 5. Default is 3.
#' @param mu1,mu2 vectors of mean parameters.
#' @param kappa1,kappa2,kappa3 vectors of concentration parameters; \code{kappa1, kappa2 > 0},
#' and \code{kappa3^2 < kappa1*kappa2}.
#' @param ... additional arguments passed to \link[mvtnorm]{rmvnorm} from package \code{mvtnorm}
#' @importFrom mvtnorm rmvnorm
#'
#' @details
#' The bivariate wrapped normal density at the point \eqn{x = (x_1, x_2)} is given by,
#' \deqn{f(x) = \sqrt((\kappa_1 \kappa_2 - (\kappa_3)^2)) / (2\pi) \sum \exp(-1/2 * (\kappa_1 (T_1)^2 + \kappa_2 (T_2)^2 + 2 \kappa_3 (T_1) (T_2)) )}
#' where
#' \deqn{T_1 = T_1(x, \mu, \omega) = (x_1 - \mu_1(2\pi\omega_1))}
#' \deqn{T_2 = T_2(x, \mu, \omega) = (x_2 - \mu_1(2\pi\omega_2))}
#' the sum extends over all pairs of integers \eqn{\omega = (\omega_1, \omega_2)},
#' and is approximated by a sum over \eqn{(\omega_1, \omega_2)} in \eqn{\{-M, -M+1, ..., M-1, M \}^2} if \code{int.displ = } \eqn{M}.
#'
#' Note that above density is essentially the "wrapped" version of a bivariate normal density with mean
#' \deqn{\mu = (\mu_1, \mu_2)}
#' and dispersion matrix  \eqn{\Sigma = \Delta^{-1}}, with \eqn{\Delta} being a \eqn{2 \times 2}
#' symmetric, positive definite matrix
#' with diagonal entries \eqn{\kappa_1} and \eqn{\kappa_2} and the off-diagonal entries
#' \eqn{\kappa_3}.
#'
#'
#' @return \code{dwnorm2} gives the density  and \code{rwnorm2} generates random deviates.
#'
#' @examples
#' kappa1 <- c(1, 2, 3)
#' kappa2 <- c(1, 6, 5)
#' kappa3 <- c(0, 1, 2)
#' mu1 <- c(1, 2, 5)
#' mu2 <- c(0, 1, 3)
#' x <- diag(2, 2)
#' n <- 10
#'
#' # when x is a bivariate vector and parameters are all scalars,
#' # dwnorm2 returns single density
#' dwnorm2(x[1, ], kappa1[1], kappa2[1], kappa3[1], mu1[1], mu2[1])
#'
#' # when x is a two column matrix and parameters are all scalars,
#' # dmvsin returns a vector of densities calculated at the rows of
#' # x with the same parameters
#' dwnorm2(x, kappa1[1], kappa2[1], kappa3[1], mu1[1], mu2[1])
#'
#' # if x is a bivariate vector and at least one of the parameters is
#' # a vector, all parameters are recycled to the same length, and
#' # dwnorm2 returns a vector of with ith element being the density
#' # evaluated at x with parameter values kappa1[i], kappa2[i],
#' # kappa3[i], mu1[i] and mu2[i]
#' dwnorm2(x[1, ], kappa1, kappa2, kappa3, mu1, mu2)
#'
#' # if x is a two column matrix and at least one of the parameters is
#' # a vector, rows of x and the parameters are recycled to the same
#' # length, and dwnorm2 returns a vector of with ith element being the
#' # density evaluated at ith row of x with parameter values kappa1[i],
#' # kappa2[i], # kappa3[i], mu1[i] and mu2[i]
#' dwnorm2(x, kappa1, kappa2, kappa3, mu1, mu2)
#'
#' # when parameters are all scalars, number of observations generated
#' # by rwnorm2 is n
#' rwnorm2(n, kappa1[1], kappa2[1], kappa3[1], mu1[1], mu2[1])
#'
#' # when at least one of the parameters is a vector, all parameters are
#' # recycled to the same length, n is ignored, and the number of
#' # observations generated by rwnorm2 is the same as the length of the
#' # recycled vectors
#' rwnorm2(n, kappa1, kappa2, kappa3, mu1, mu2)
#'
#' @export

rwnorm2 <- function(n, kappa1=1, kappa2=1,
                    kappa3=0, mu1=0, mu2=0, ...)
{
  if(any(c(kappa1, kappa2) < 0))
    stop("kappa1 and kappa2 must be non-negative")
  if(any(mu1 < 0 | mu1 >= 2*pi)) mu1 <- prncp_reg(mu1)
  if(any(mu2 < 0 | mu2 >= 2*pi)) mu2 <- prncp_reg(mu2)

  if(max(length(kappa1), length(kappa2), length(kappa3),
         length(mu1), length(mu2)) > 1) {
    expanded <- expand_args(kappa1, kappa2, kappa3, mu1, mu2)
    kappa1 <- expanded[[1]]; kappa2 <- expanded[[2]]; kappa3 <- expanded[[3]]
    mu1 <- expanded[[4]]; mu2 <- expanded[[5]]
    if(any(kappa1*kappa2 - kappa3*kappa3 < 0))
      stop("abs(kappa3) must be less than or equal to sqrt(kappa1*kappa2) in wnorm2")
    m <- length(kappa1)

    samp <- matrix(0, m, 2)

    for (j in 1:m) {
      if (kappa1[j] < 1e-10 & kappa2[j] < 1e-10) {
        samp[j, ] <- runif(2, 0, 2*pi)
      }

      else {
        mu_curr <-  c(mu1[j], mu2[j])
        sigma_curr <-  matrix(c(kappa2[j], -kappa3[j], -kappa3[j],
                                kappa1[j])/(kappa1[j]*kappa2[j] - kappa3[j]^2), 2)
        # samp <- t(sapply(1:m, function(j) rnorm2(1, mu_list[[j]], sigma_list[[j]])))

        samp[j, ] <-
          rmvnorm(1, mean=mu_curr, sigma=sigma_curr, ...)
      }
    }
  }

  else {
    if(kappa1*kappa2 - kappa3*kappa3 < 0)
      stop("abs(kappa3) must be less than or equal to sqrt(kappa1*kappa2) in wnorm2")

    if (kappa1*kappa2 < 1e-10) {
      samp <- matrix(runif(2*n, 0, 2*pi), ncol=2)
    }

    else {
      mu <- c(mu1, mu2)
      sigma <- matrix(c(kappa2, -kappa3, -kappa3, kappa1)/(kappa1*kappa2 - kappa3*kappa3), 2)
      samp <- rmvnorm(n, mean=mu, sigma=sigma, ...)
    }
  }
  prncp_reg(samp)
}

#' @rdname rwnorm2
#' @export

dwnorm2 <- function(x, kappa1=1, kappa2=1, kappa3=0, mu1=0,
                    mu2=0, int.displ, log=FALSE)
{
  if(missing(int.displ)) int.displ <- 3
  else if(int.displ >= 5) int.displ <- 5
  else if(int.displ <= 1) int.displ <- 1
  displ <- floor(int.displ)
  omega.2pi.all <- expand.grid(-displ:displ,-displ:displ) * (2*pi) # 2pi * integer displacements
  omega.2pi <- as.matrix(omega.2pi.all)

  if(any(c(kappa1, kappa2) < 0))
    stop("kappa1 and kappa2 must be non-negative")
  if(any(mu1 < 0 | mu1 >= 2*pi)) mu1 <- prncp_reg(mu1)
  if(any(mu2 < 0 | mu2 >= 2*pi)) mu2 <- prncp_reg(mu2)
  if((length(dim(x)) < 2 && length(x) != 2) || (length(dim(x)) == 2 && tail(dim(x), 1) != 2)
     || (length(dim(x)) > 2)) stop("x must either be a bivariate vector or a two-column matrix")


  if(max(length(kappa1), length(kappa2), length(kappa3), length(mu1), length(mu2)) > 1) {
    expanded <- expand_args(kappa1, kappa2, kappa3, mu1, mu2)
    kappa1 <- expanded[[1]]
    kappa2 <- expanded[[2]]
    kappa3 <- expanded[[3]]
    mu1 <- expanded[[4]]
    mu2 <- expanded[[5]]
  }

  par.mat <- rbind(kappa1, kappa2, kappa3, mu1, mu2)
  n_par <- ncol(par.mat)
  if (length(x) == 2) x <- matrix(x, nrow = 1)
  n_x <- nrow(x)

  if (all (kappa1 > 1e-10 | kappa2 > 1e-10)) {
    # regular wnorm2 density
    if (n_par == 1) {
      den <- c(dwnorm2_manyx_onepar(x, kappa1, kappa2, kappa3,
                                    mu1, mu2, omega.2pi))
    } else if (n_x == 1) {
      den <- c(dwnorm2_onex_manypar(c(x), kappa1, kappa2, kappa3,
                                    mu1, mu2, omega.2pi))
    } else {
      x_set <- 1:n_x
      par_set <- 1:n_par
      expndn_set <- expand_args(x_set, par_set)
      x_set <- expndn_set[[1]]
      par_set <- expndn_set[[2]]
      den <- c(dwnorm2_manyx_manypar(x[x_set, ], kappa1[par_set],
                                     kappa2[par_set], kappa3[par_set],
                                     mu1[par_set], mu2[par_set], omega.2pi))
    }
  }
  else {
    # some can be uniform
    x_set <- 1:n_x
    par_set <- 1:n_par
    expndn_set <- expand_args(x_set, par_set)
    x_set <- expndn_set[[1]]
    par_set <- expndn_set[[2]]
    x_long <- x[x_set, , drop=FALSE]
    kappa1_long <- kappa1[par_set]
    kappa2_long <- kappa2[par_set]
    kappa3_long <- kappa3[par_set]
    mu1_long <- mu1[par_set]
    mu2_long <- mu2[par_set]
    n_x_final <- nrow(x_long)
    den <- rep(0, n_x_final)
    which_unif <- which(kappa1_long < 1e-10 & kappa2_long < 1e-10)
    n_unif <- length(which_unif)


    # browser()
    den[which_unif] <- 1/(4*pi^2)
    if (n_unif < n_x_final)
      den[-which_unif] <- c(dwnorm2_manyx_manypar(x_long[-which_unif, , drop=FALSE], kappa1_long[-which_unif],
                                                 kappa2_long[-which_unif], kappa3_long[-which_unif],
                                                 mu1_long[-which_unif], mu2_long[-which_unif],
                                                 omega.2pi))


  }

  if (log)
    den <- log(den)

  den
}

#' The bivariate Wrapped Normal mixtures
#' @inheritParams rvmsinmix
#' @inheritParams rwnorm2
#' @param mu1,mu2 vectors of mean parameters.
#' @param kappa1,kappa2,kappa3 vectors of concentration parameters; \code{kappa1, kappa2 > 0, kappa3^2 < kappa1*kappa2} for each component.
#'
#' @details All the argument vectors \code{pmix, kappa1, kappa2, kappa3, mu1} and \code{mu2} must be of the same length,
#' with \eqn{j}-th element corresponding to the \eqn{j}-th component of the mixture distribution.
#' @details The bivariate wrapped normal mixture distribution with component size \code{K = \link{length}(pmix)} has density
#' \deqn{g(x) = \sum p[j] * f(x; \kappa_1[j], \kappa_2[j], \kappa_3[j], \mu_1[j], \mu_2[j])}
#' where the sum extends over \eqn{j}; \eqn{p[j]; \kappa_1[j], \kappa_2[j], \kappa_3[j]}; and \eqn{\mu_1[j], \mu_2[j]} respectively denote the mixing proportion,
#' the three concentration parameters and the two mean parameter for the \eqn{j}-th component, \eqn{j = 1, ..., K},
#' and \eqn{f(. ; \kappa_1, \kappa_2, \kappa_3, \mu_1, \mu_2)} denotes the density function of the wrapped normal distribution
#' with concentration parameters \eqn{\kappa_1, \kappa_2, \kappa_3} and  mean parameters \eqn{\mu_1, \mu_2}.
#' @return \code{dwnorm2mix} computes the density  and \code{rwnorm2mix} generates random deviates from the mixture density.
#'
#' @examples
#' kappa1 <- c(1, 2, 3)
#' kappa2 <- c(1, 6, 5)
#' kappa3 <- c(0, 1, 2)
#' mu1 <- c(1, 2, 5)
#' mu2 <- c(0, 1, 3)
#' pmix <- c(0.3, 0.4, 0.3)
#' x <- diag(2, 2)
#' n <- 10
#'
#' # mixture densities calculated at the rows of x
#' dwnorm2mix(x, kappa1, kappa2, kappa3, mu1, mu2, pmix)
#'
#' # number of observations generated from the mixture distribution is n
#' rwnorm2mix(n, kappa1, kappa2, kappa3, mu1, mu2, pmix)
#'
#' @export

rwnorm2mix <- function(n, kappa1, kappa2, kappa3,
                       mu1, mu2, pmix, ...)
{
  allpar <- list(kappa1=kappa1, kappa2=kappa2, kappa3=kappa3,
                 mu1=mu1, mu2=mu2, pmix=pmix)

  allpar_len <- listLen(allpar)
  if(min(allpar_len) != max(allpar_len))
    stop("component size mismatch: number of components of the input parameter vectors differ")

  if(any(allpar$pmix < 0)) stop("\'pmix\' must be non-negative")
  sum_pmix <- sum(allpar$pmix)
  if(signif(sum_pmix, 5) != 1) {
    if(sum_pmix <= 0) stop("\'pmix\' must have at least one positive element")
    allpar$pmix <- allpar$pmix/sum_pmix
    warning("\'pmix\' is rescaled to add up to 1")
  }

  if(any(c(allpar$kappa1, allpar$kappa2) <= 0))
    stop("kappa1 and kappa2 must be positive in wnorm2")
  if(any(allpar$kappa1*allpar$kappa2 - allpar$kappa3^2 <= 1e-10))
    stop("abs(kappa3) must be less than sqrt(kappa1*kappa2) in wnorm2")
  if(any(allpar$mu1 < 0 | allpar$mu1 >= 2*pi)) allpar$mu1 <- prncp_reg(allpar$mu1)
  if(any(allpar$mu2 < 0 | allpar$mu2 >= 2*pi)) allpar$mu2 <- prncp_reg(allpar$mu2)

  out <- matrix(0, n, 2)
  ncomp <- allpar_len[1] # number of components
  comp_ind <- cID(tcrossprod(rep(1, n), allpar$pmix), ncomp, runif(n))
  # n samples from multinom(ncomp, pmix)
  for(j in seq_len(ncomp)) {
    obs_ind_j <- which(comp_ind == j)
    n_j <- length(obs_ind_j)
    if(n_j > 0) {
      out[obs_ind_j, ] <- rwnorm2(n_j, kappa1[j], kappa2[j],
                                  kappa3[j], mu1[j], mu2[j], ...)
    }
  }

  out
}


#' @rdname rwnorm2mix
#' @export
dwnorm2mix <- function(x, kappa1, kappa2, kappa3,
                       mu1, mu2, pmix, int.displ, log=FALSE)
{
  if(missing(int.displ)) int.displ <- 3
  else if(int.displ >= 5) int.displ <- 5
  else if(int.displ <= 1) int.displ <- 1
  displ <- floor(int.displ)
  omega.2pi.all <- expand.grid(-displ:displ,-displ:displ) * (2*pi) # 2pi * integer displacements
  omega.2pi <- as.matrix(omega.2pi.all)

  allpar <- list("kappa1"=kappa1, "kappa2"=kappa2, "kappa3"=kappa3,
                 "mu1"=mu1, "mu2"=mu2, "pmix"=pmix)

  allpar_len <- listLen(allpar)
  if(min(allpar_len) != max(allpar_len)) stop("component size mismatch: number of components of the input parameter vectors differ")

  if(any(allpar$pmix < 0)) stop("\'pmix\' must be non-negative")
  sum_pmix <- sum(allpar$pmix)
  if(signif(sum_pmix, 5) != 1) {
    if(sum_pmix <= 0) stop("\'pmix\' must have at least one positive element")
    allpar$pmix <- allpar$pmix/sum_pmix
    warning("\'pmix\' is rescaled to add up to 1")
  }

  if(any(allpar$kappa1*allpar$kappa2 - allpar$kappa3*allpar$kappa3 <= 1e-10))
    stop("abs(kappa3) must be less than sqrt(kappa1*kappa2) in wnorm2")

  if(any(c(allpar$kappa1, allpar$kappa2) <= 0))
    stop("kappa1 and kappa2 must be positive in wnorm2")
  if(any(allpar$mu1 < 0 | allpar$mu1 >= 2*pi)) allpar$mu1 <- prncp_reg(allpar$mu1)
  if(any(allpar$mu2 < 0 | allpar$mu2 >= 2*pi)) allpar$mu2 <- prncp_reg(allpar$mu2)
  if((length(dim(x)) < 2 && length(x) != 2) || (length(dim(x)) == 2 && tail(dim(x), 1) != 2)
     || (length(dim(x)) > 2)) stop("x must either be a bivariate vector or a two-column matrix")


  ncomp <- length(kappa1)

  if (length(x) == 2) x <- matrix(x, nrow=1)

  allcompden <- vapply(1:ncomp,
                       function(j) dwnorm2(x, kappa1[j], kappa2[j],
                                           kappa3[j], mu1[j], mu2[j],
                                           int.displ, FALSE),
                       rep(0, nrow(x)))

  mixden <- c(allcompden %*% pmix)

  if (log) {
    log(mixden)
  } else {
    mixden
  }
}


#' Fitting bivariate wrapped normal model mixtures using MCMC
#' @inheritParams fit_vmsinmix
#'
#' @details
#' Wrapper for \link{fit_angmix} with \code{model = "wnorm2"}.
#'
#'
#' @examples
#' # illustration only - more iterations needed for convergence
#' fit.wnorm2.10 <- fit_wnorm2mix(tim8, ncomp = 3, n.iter =  10,
#'                              n.chains = 1)
#' fit.wnorm2.10
#'
#' @export

fit_wnorm2mix <- function(...)
{
  fit_angmix(model="wnorm2", ...)
}


wnorm2_var_cor_singlepar <- function(kappa1, kappa2, kappa3) {
  den <- kappa1*kappa2 - kappa3^2
  sig1_sq <- kappa2/den
  sig2_sq <- kappa1/den
  sig12 <- -kappa3/den

  rho_fl <- sinh(2*sig12) /
    sqrt(sinh(2*sig1_sq)* sinh(2*sig2_sq))

  rho_js <- sinh(sig12) /
    sqrt(sinh(sig1_sq)* sinh(sig2_sq))

  list(var1 = 1-exp(-0.5*sig1_sq),
       var2 = 1-exp(-0.5*sig2_sq),
       rho_fl = rho_fl,
       rho_js = rho_js
  )
}


