#' SARARgamlss: Spatial Autoregressive Generalized Additive Model for Location Scale (GAMLSS)
#'
#' This function estimates a Spatial Autoregressive Generalized Additive Model for Location Scale 
#' (SARARgamlss) using GAMLSS. The model includes both spatial dependencies and the possibility of 
#' non-parametric terms in the formulas for the mean and variance. The function supports SAR, SARAR, 
#' and SEM model types and performs the estimation through an iterative process that updates spatial 
#' dependence parameters. The variance of the spatial parameters \eqn{\hat{\rho}} and \eqn{\hat{\lambda}}
#' is estimated using the inverse of the Hessian matrix from the optimization.
#'
#' @param formula A formula specifying the mean structure of the model (response ~ explanatory variables).
#' @param sigma.formula A formula specifying the variance structure of the model (default: ~1).
#' @param W1 A spatial weights matrix for the SAR term (default: identity matrix).
#' @param W2 A spatial weights matrix for the SARAR term (default: identity matrix).
#' @param data A data.frame containing the variables used in the model.
#' @param tol Convergence tolerance (default: 1E-4).
#' @param maxiter Maximum number of iterations for optimization (default: 20).
#' @param type The type of spatial model to fit: one of "SAR", "SARAR", or "SEM".
#' @param weights Optional weights for the observations (default: NULL).
#' 
#' @return A fitted GAMLSS model object with spatial autoregressive terms. The model object also includes
#' the variance of the spatial parameters \eqn{\hat{\rho}} and \eqn{\hat{\lambda}}
#' @references Toloza-Delgado, J. D., Melo, O. O., & Cruz, N. A.
#'  Joint spatial modeling of mean and non-homogeneous variance combining semiparametric SAR 
#'  and GAMLSS models for hedonic prices. Spatial Statistics, 65, 100864 (2025)
#'  @source https://doi.org/10.1016/j.spasta.2024.100864
#' @examples
#' library(spdep)
#' library(gamlss)
#' data(oldcol)
#' # Create spatial weight matrices W1 and W2
#' W1 <- spdep::nb2mat(COL.nb, style = "W")
#' W2 <- W1  # In this case, assume the same spatial weights for both
#' # Fit a SARARgamlss model
#' result <- SARARgamlss(formula = CRIME ~ INC + cs(HOVAL), 
#' sigma.formula = ~ INC + pb(HOVAL), W1 = W1, W2 = W2,data = COL.OLD, 
#' tol = 1E-4,  maxiter = 20, type = "SARAR")
#' summary_SAR(result)
#' gamlss::term.plot(result$gamlss, what="mu")
#' 
#' @export
#' @importFrom methods is
#' @importFrom stats .getXlevels binomial dpois fitted gaussian glm.fit 
#' @importFrom stats make.link model.matrix model.offset model.response coef model.frame
#' @importFrom stats optim pchisq pnorm predict printCoefmat rpois update
#' @import gamlss
#' @import splines
#' @importFrom gamlss.dist NO

SARARgamlss <- function(formula, sigma.formula = ~1,
                        W1 = diag(0, nrow(data)), W2 = diag(0, nrow(data)),
                        data, tol = 1E-4, maxiter = 20,
                        type = c("SAR", "SARAR", "SEM"),
                        weights = NULL) {
  mf <- stats::model.frame(formula, data=data)
  data_name <- substitute(data) 
  if (type == "SAR") {
    W2 = 0 * W2
  }
  if (type == "SEM") {
    W1 = 0 * W1
  }
  if (type == "SARAR" & sum(W2) == 0) {
    W2 <- W1
  }
  
  n <- nrow(data)
  # Fit the initial GAMLSS model for mean (mu) and variance (sigma)
  m0 <- gamlss::gamlss(formula = formula, sigma.formula = sigma.formula, 
                       data = data, family = NO())
  
  profile_loglik <- function(par, W1, W2, formula, sigma.formula, data){
    
    rho <- par[1]
    lambda <- par[2]
    
    n <- nrow(data)
    
    AA <- Matrix::Diagonal(n) - rho * W1
    BB <- Matrix::Diagonal(n) - lambda * W2
    
    Y <- model.response(model.frame(formula,data))
    Ytemp <- as.vector(BB %*% AA %*% Y)
    
    mf <- model.frame(formula,data)
    X <- model.matrix(attr(mf,"terms"),mf)
    
    Xtemp <- as.matrix(BB %*% X)
    
    datatemp <- data
    datatemp$Ytemp <- Ytemp
    
    m <- gamlss::gamlss(
      Ytemp ~ Xtemp -1,
      sigma.formula = sigma.formula,
      family = NO(),
      data = datatemp,
      trace = FALSE
    )
    
    mu <- predict(m,"mu",type="response")
    sigma2 <- predict(m,"sigma",type="response")^2
    
    detAA <- Matrix::determinant(AA,logarithm=TRUE)
    detBB <- Matrix::determinant(BB,logarithm=TRUE)
    
    if(detAA$sign <=0 | detBB$sign <=0) return(1e12)
    
    ll <- detAA$modulus + detBB$modulus -
      0.5*sum(log(sigma2)) -
      0.5*sum((Ytemp-mu)^2/sigma2)
    
    return(-ll)
  }
  
  opt <- optim( c(0,0),  profile_loglik,W1=W1,W2=W2,formula=formula,
    sigma.formula=sigma.formula,data=data,method="L-BFGS-B",
    lower=c(-0.999,-0.999), upper=c(0.999,0.999), hessian = TRUE)
  
  Y <- model.response(model.frame(formula,data))
  p0 <- opt$par
  AA <- Matrix::Diagonal(n) - p0[1] * W1
  BB <- Matrix::Diagonal(n) - p0[2] * W2
  Ytemp <- as.matrix(BB %*% AA %*% Y)
  Xtemp <- as.matrix(BB %*% model.matrix(m0, what = "mu"))
  Ztemp <- model.matrix(m0, what = "sigma")
  Hessian <- opt$hessian  # Extract Hessian matrix
  
  colnames(Xtemp) <- colnames(model.matrix(m0, what = "mu"))
  colnames(Ztemp) <- colnames(model.matrix(m0, what = "sigma"))
  
  
  # Fit updated GAMLSS model with transformed data (dependent and independent)
  m1 <- gamlss::gamlss(Ytemp ~ Xtemp - 1, ~Ztemp - 1, family = NO())
  var1 <- predict(m1, what = "sigma", type = "response")^2
  Xbeta <- Matrix::solve(BB,predict(m1, what = "mu", type = "response"))
  # Store updated model parameters
  
  names_to_update <- c("G.deviance", "residuals", "mu.fv", "mu.lp", "mu.wv", 
                       "mu.wt", "mu.qr", "mu.coefficients", "sigma.fv",
                       "sigma.wv", "sigma.wt", "sigma.coefficients", "sigma.qr",
                       "P.deviance", "aic", "sbc")
  
  for(valu in names_to_update){
    names(m1[[valu]]) <- names(m0[[valu]])
  }
  
  for (names in names_to_update) {
    m0[[names]] <- m1[[names]]
  }
  
  #model.frame(m0, what="mu") <- 
  # Return the final model object with spatial parameters and variance estimates
  if (type == "SARAR") {
    spamu <- c(p0[1],p0[2])
    var_cov_matrix <- solve(Hessian)
    spacov = var_cov_matrix
  } else if (type == "SAR") {
    spamu <- c(p0[1], NA)
    var_cov_matrix <- matrix(c( solve(Hessian[1,1]),NA, NA,NA), ncol=2, nrow=2)
    spacov = var_cov_matrix
  } else {
    spamu <- c(NA, p0[2])
    var_cov_matrix <- matrix(c(NA, NA,NA, solve(Hessian[2,2])), ncol=2, nrow=2)
    spacov =var_cov_matrix 
  }
  # model.frame(m0)[,1] <- as.vector(model.frame(m1)[,1])
  m0$call$data <- data_name
  
  X <- model.matrix(m0, what = "mu")
  beta <- stats::coef(m0, "mu")
  Omega <- Matrix::diag(predict(m0, what = "sigma")^2)
  rho <- p0[1]
  lambda <- p0[2]
  AA <- Matrix::Diagonal(n) - rho * W1
  BB <- Matrix::Diagonal(n) - lambda * W2
  BBInv <- Matrix::solve(BB)
  varU <- BBInv%*%Omega %*% Matrix::t(BBInv)
  
  inv_cov <- Matrix::solve(lambda^2 * W2 %*% varU %*% Matrix::t(W2) +
                     lambda *  Omega %*% Matrix::t(BBInv) %*% Matrix::t(W2) +
                     lambda * W2 %*% BBInv %*% Omega +
                     Omega)
  y_true <- Y
  y_signal <- (lambda* W2 %*% varU + BBInv %*% Omega) %*% inv_cov %*% 
    (AA%*%y_true -  X %*% beta)
  
  y_trend <- rho * W1 %*% y_true + X %*% beta
  y_blup <- y_trend + y_signal
  residuals <- AA%*%y_true -  X %*% beta
  y_noise <- y_true - y_blup
  
  out1 <- list(gamlss=m0, model = mf,
               spatial=list(spatial=spamu, sdspatial=spacov, type=type),
               gamlssAY=m1,  y_signal = y_signal, y_trend = y_trend,
               y_blup = y_blup, residuals = residuals,y_noise = y_noise)
  #out1$call <- match.call()
  class(out1) <- "SARARgamlss"
  return(out1)
}

