#' Model Specification
#'
#' Method for creating dsp specification object prior to fitting.
#'
#' @param family A character string specifying the model family. Must be one of:
#'   \itemize{
#'     \item "gaussian": Gaussian family.
#'     \item "negbinomial": Negative binomial family.
#'   }
#' @param model A character string specifying the model type:
#'   \itemize{
#'     \item \code{family} = "gaussian":
#'       \itemize{
#'         \item "changepoint": Change point detection with
#'         Adaptive Bayesian Changepoint analysis and local Outlier (ABCO),
#'         \item "smoothing": Bayesian smoothing,
#'         \item "regression": Time-varying regression,
#'         \item "bspline": Bayesian smoothing with B-spline for irregularly spaced or functional time-series.
#'       }
#'     \item \code{family} = "negbinomial":
#'       \itemize{
#'         \item "smoothing": Bayesian smoothing.
#'       }
#'   }
#' @param ... Additional arguments based on \code{family} and \code{model}.
#' For \code{family} = "gaussian" and \code{model} = "regression", \code{X} is required.
#' Other arguments are optional and default to predefined value.
#'
#' \tabular{lll}{
#'  **Family**    \tab  **Model**       \tab    **Required Arguments**  \cr
#'  "gaussian"    \tab "changepoint"     \tab \code{D}, \code{useAnom}, \code{obsSV} \code{D_asv}, \code{evol_error_asv}, \code{nugget_asv} \cr
#'  "gaussian"    \tab "smoothing"       \tab \code{D}, \code{evol_error}, \code{obsSV}, \code{D_asv}, \code{evol_error_asv}, \code{nugget_asv}\cr
#'  "gaussian"    \tab "regression"      \tab \code{D}, \code{evol_error}, \code{obsSV}, \code{X}, \code{D_asv}, \code{evol_error_asv}, \code{nugget_asv} \cr
#'  "gaussian"    \tab "bspline"        \tab \code{D}, \code{evol_error}, \code{obsSV}, \code{times}, \code{num_knots}  \cr
#'  "negbinomial"    \tab "smoothing"       \tab \code{D}, \code{evol_error}, \code{r_user}, \code{offset}  \cr
#' }
#'
#'   \itemize{
#'     \item **Shared Arguments**:
#'       \itemize{
#'         \item \code{D}: integer scalar indicating degree of differencing.
#'         Implementation is available D = 0, 1, or 2 (default) for \code{family} = "gaussian",
#'         and D = 1 and 2 (default) for \code{family} = "negbinomial".
#'         \item \code{obsSV}:Options for modeling the error variance for \code{family} = "gaussian".
#'         It must be one of the following:
#'            \itemize{
#'              \item "const": Constant error variance for all time points (default);
#'              \item "SV": Stochastic Volatility model;
#'              \item "ASV": Adaptive Stochastic Volatility model;
#'            }
#'         \item \code{D_asv}: integer; degree of differencing (0, 1 (default), or 2) for the ASV model. Only used when \code{obsSV = "ASV"}.
#'         \item \code{evol_error_asv}: character; "HS" by default. evolution error distribution for the ASV model. Must be one of the five options used in \code{evol_error}. Only used when \code{obsSV = "ASV"}
#'         \item \code{nugget_asv}: logical; if \code{TRUE} (default), fits the nugget variant of the ASV model. Only used when \code{obsSV = "ASV"}.
#'         \item \code{evol_error}: the evolution error distribution; .
#'            \itemize{
#'              \item the dynamic horseshoe prior ("DHS") (default);
#'              \item the horseshoe prior ("HS");
#'              \item the Bayesian lasso ("BL") (\code{family} = "gaussian" only);
#'              \item the normal stochastic volatility model ("SV") (\code{family} = "gaussian" only);
#'              \item the normal-inverse-gamma prior ("NIG") (\code{family} = "gaussian" only).
#'            }
#'       }
#'     \item **Model-specific Arguments**:
#'      \itemize{
#'        \item \code{model} = "changepoint":
#'          \itemize{
#'            \item \code{useAnom}: logical; Defaults to FALSE. if TRUE, include an anomaly component
#'            in the observation equation.
#'          }
#'        \item \code{family} = "gaussian", \code{model} = "regression":
#'        \itemize{
#'          \item \code{X}: matrix or data frame; T x p matrix of time series predictors
#'        }
#'        \item \code{model} = "bspline":
#'        \itemize{
#'          \item \code{times}: length \code{T} vector; observation indices; if `NULL` (Default), assume equally spaced.
#'          \item \code{num_knots}: numeric; the number of knots to be used for bspline. defaults to 20. Has to be at least 4.
#'        }
#'        \item \code{family} = "negbinomial", \code{model} = "smoothing":
#'        \itemize{
#'          \item \code{r_user}: numeric; if `NULL` (Default), MCMC sampling is done for the overdispersion parameter \code{r}; If a positive numeric value is passed, the overdispersion is fixed to that value. Observation distribution approximately Poisson if a large numeric value is passed.
#'          \item \code{offset}: length \code{T} vector; offset values for the log conditional expectation
#'        }
#'     }
#'   }
#'
#' @return A list containing the model specification.
#'
#' @examples
#' model_spec <- dsp_spec(family = "gaussian",
#'                       model = "changepoint")
#'
#'
#' @export
dsp_spec <- function(family,
                     model,
                     ...){
  if(!family %in% c("gaussian", "negbinomial")) stop('family must be gaussian or negbinomial')
  input_args <- list(...)
  if(family == "gaussian"){
    if(!model %in% c("changepoint", "regression",
                     "smoothing","bspline")){
      stop('model must be "changepoint","regression","smoothing" or "bspline for gaussian family')
    }
  }else{
    if(model != "smoothing"){
      stop('model must be "smoothing" for negative binomial family')
    }
  }
  # Extract obsSV if provided
  obsSV <- input_args$obsSV
  if (is.null(obsSV)) obsSV <- "NONE"  # default fallback if needed

  # Transform r_user to the two arguments r_init and r_sample
  if(!is.null(input_args$r_user)){
    input_args$r_init <- input_args$r_user
    input_args$r_sample <- FALSE
  }

  # if NULL, just drop the input arg and the behavior will be default
  input_args$r_user <- NULL

  #inputted arguments
  required_args <- function(family,model){
    if(family == "gaussian"){
      if(model == "changepoint" && obsSV != "ASV"){
        return(c("D", "useAnom", "obsSV"))
      }else if(model == "changepoint" && obsSV == "ASV"){
        return(c("D", "useAnom", "obsSV","D_asv","evol_error_asv","nugget_asv"))
      }else if(model == "smoothing" && obsSV == "ASV"){
        return(c("evol_error","D","obsSV","D_asv","evol_error_asv","nugget_asv"))
      }else if(model == "smoothing" && obsSV != "ASV"){
        return(c("evol_error","D","obsSV"))
      }else if(model == "regression" && obsSV == "ASV"){
        return(c("X", "evol_error", "D","obsSV","D_asv","evol_error_asv","nugget_asv"))
      }else if(model == "regression" && obsSV != "ASV"){
        return(c("X", "evol_error", "D","obsSV"))
      }else if(model == "bspline"){
        return(c("evol_error", "D","times","num_knots"))
      }
    }else{
      # negative binomial smoothing
      return(c("evol_error","D","r_init","r_sample","offset"))
    }
  }

  # Handling unnecessary arguments supplied as inputs
  extraArgs = setdiff(names(input_args),required_args(family,model))
  if(length(extraArgs) > 0){
    warning(
      sprintf("The following arguments are not used for %s family %s model and will be ignored: %s",
              family,model, paste(extraArgs, collapse = ", ")),
      call. = FALSE)
    for(arg in extraArgs){
      input_args[[arg]] = NULL
    }
  }
  # check the input type and throw error if the inputs are invalid
  validation_rules <- list(
    X = function(family,x) { if (any(is.na(x)) || !(is.matrix(x) || is.data.frame(x))) stop("X must be a dataframe or a matrix")},
    D = function(family,x) {
      if (length(x) != 1 || is.na(x) || !is.numeric(x) || x != as.integer(x)) {
        stop("D must be a single integer.")
      }
      if (family == "gaussian" && !x %in% c(0,1,2)) stop("D must be 0, 1 or 2 for gaussian models.")
      if (family != "gaussian" && !x %in% c(1,2)) stop("D must be 1 or 2 for negative binomial models.")
    },
    D_asv = function(family,x) {
      if (length(x) != 1 || is.na(x) || !is.numeric(x) || x != as.integer(x)) {
        stop("D_asv must be a single integer.")
      }
      if (!x %in% c(0, 1, 2)) stop("D_asv must be 0, 1 or 2 for ASV.")
    },
    nugget_asv = function(family,x) {
      if (is.na(x)|| !is.logical(x)) stop("useAnom must be TRUE or FALSE.")
    },
    evol_error_asv = function(family,x) {
      if (length(x) != 1 || is.na(x)) stop("evol_error_asv must be a single character value.")
      if (!x %in% c("HS", "DHS", "NIG", "SV", "BL")) {
        stop("evol_error_asv must be one of 'HS', 'DHS', 'NIG', 'SV', 'BL' for gaussian models")
      }
    },
    evol_error = function(family, x) {
      if (length(x) != 1 || is.na(x)) stop("evol_error must be a single character value.")
      if (family == "gaussian") {
        if (!x %in% c("HS", "DHS", "NIG", "SV", "BL")) {
          stop("evol_error must be one of 'HS', 'DHS', 'NIG', 'SV', 'BL' for gaussian models")
        }
      } else {
        if (!x %in% c("HS", "DHS")) {
          stop("evol_error must be one of 'HS', 'DHS' for negative binomial models")
        }
      }
    },
    useAnom = function(family,x) { if (is.na(x)|| !is.logical(x)) stop("useAnom must be TRUE or FALSE.") },
    obsSV = function(family,x) { if (!x %in% c("const", "SV", "ASV")) stop("obsSV must be one of 'const', 'SV', 'ASV'.") },
    times = function(family, x) {
        if (is.null(x)) return(invisible(TRUE))
        if (!is.numeric(x) || anyNA(x) || any(!is.finite(x))) {
          stop("times must be a numeric vector with no missing/non-finite values, or NULL")
        }
        if (length(x) < 2) stop("times must have length >= 2")
        if (diff(range(x)) <= 0) stop("times must contain at least two distinct values")
    },
    num_knots = function(family, x) {
      if (is.null(x)) return(invisible(TRUE))  # keep if you want NULL to use default
      if (length(x) != 1 || is.na(x) || !is.numeric(x) || !is.finite(x) ||
          x != as.integer(x) || x < 4) {
        stop("num_knots must be an integer >= 4, or NULL.")
      }
    },
    r_init = function(family,x){ if (!is.numeric(x) || x < 0) stop("r_user nust be a positive number.")},
    r_sample = function(family,x) { if (is.na(x) || !is.logical(x)) stop("r_sample must be TRUE or FALSE.")},
    offset = function(family,x){ if (any(is.na(x)) || !is.numeric(x)) stop("offset must be a vector") }
  )

  if(length(input_args) >0){
    for(arg in names(input_args)){
      tryCatch(
        validation_rules[[arg]](family,input_args[[arg]]),
        error = function(e) { stop(sprintf("invalid input for %s model in %s family: %s",
                                           model, family, e$message)) }
      )
    }
  }
  # If any of the required arguments are missing give them default values
  default_args <- list(D =2, useAnom = FALSE, obsSV = "const",times = NULL,
                       evol_error = "DHS", num_knots = 20,
                       r_init = 5, r_sample = FALSE, offset = 0,
                       D_asv = 1,evol_error_asv = "HS",nugget_asv = FALSE)
  requiredArgs = setdiff(required_args(family,model), names(input_args))
  # X is required to run regression
  if("X" %in% requiredArgs) stop("To run regression, the design matrix 'X' needs to be specified.")
  if(length(requiredArgs)>0){
    for(arg in requiredArgs){
      dvalue = default_args[[arg]]
      input_args[[arg]] = dvalue
      message(sprintf("Argument '%s' is missing. Using default value: %s", arg,
                      if (is.null(dvalue)) "NULL" else dvalue)
              )
    }
  }
  ret = structure(list(
    family = family,
    model = model,
    arguments = input_args
  ),
  class = c("dsp_spec"))
  return(ret)

}

#' MCMC Sampler for Models with Dynamic Shrinkage Processes
#'
#' Wrapper function for fitting models with Dynamic Shrinkage Processes (DSP), including:
#' \itemize{
#'         \item Adaptive Bayesian Changepoint analysis and local Outlier (ABCO),
#'         \item Bayesian Trend Filter for Gaussian Data
#'         \item Time-varying Regression
#'         \item Bayesian Trend Filter with B-spline for irregularly spaced or functional time-series.
#'         \item Bayesian Smoothing for Count Data
#'}
#'
#' @param y a numeric vector of the \code{T x 1} vector of time series observations
#' @param model_spec a list containing model specification generated from [dsp_spec()].
#' @param nsave integer scalar (default = 1000); number of MCMC iterations to record
#' @param nburn integer scalar (default = 1000); number of MCMC iterations to discard (burn-in)
#' @param nskip integer scalar (default = 4); number of MCMC iterations to skip between saving iterations,
#' i.e., save every (nskip + 1)th draw
#' @param computeDIC logical; if TRUE (default), compute the deviance information criterion \code{DIC}
#' and the effective number of parameters \code{p_d}
#' @param verbose logical; should extra information on progress be printed to the console? Defaults to FALSE
#' @param ... optional additional arguments to pass to the MCMC sampler.
#'
#' @return \code{dsp_fit} returns an object of class "\code{dsp}".
#'
#' An object of class "\code{dsp}" is defined as a list containing at least the following components:
#'    \item{mcmc_output}{a list of the \code{nsave} MCMC samples for the parameters named in \code{mcmc_params}}
#'    \item{DIC}{Deviance Information Criterion}
#'    \item{mcpar}{named vector of supplied nsave, nburn, and nskip}
#'    \item{model_spec}{the object supplied for model_spec argument}
#'
#' @section MCMC Output Parameters:
#' Shared parameters across wrappers include:
#' \itemize{
#'   \item \code{mu}: Conditional mean.
#'   \item \code{ypred}: Posterior predictive \eqn{\tilde{y}_t}.
#'   \item \code{evol_sigma_t2}: Variance of the state variable.
#'   \item \code{obs_sigma_t2}: Observation variance \eqn{\sigma^2_{\epsilon}} (\code{family = "gaussian"}).
#'   \item \code{dhs_phi}, \code{dhs_mean}: DHS hyperparameter draws (\code{evol_error = "DHS"}).
#' }
#'
#' Model-specific parameters include:
#' \itemize{
#'   \item \code{changepoint} (Gaussian):
#'   \itemize{
#'     \item \code{omega}: \eqn{D}-th differenced posterior mean draws, \eqn{\mu_t}.
#'     \item \code{zeta}: anomaly component (when \code{useAnom = TRUE}).
#'     \item \code{zeta_sigma_t2}: anomaly variance (when \code{useAnom = TRUE}).
#'     \item \code{gamma}: threshold parameter in the thresholded DHS log-volatility dynamics.
#'   }
#'   \item \code{regression} (Gaussian):
#'   \itemize{
#'     \item \code{beta}: time-varying regression coefficient draws.
#'   }
#'   \item \code{bspline} (Gaussian):
#'   \itemize{
#'     \item \code{beta}: B-spline basis coefficient draws.
#'   }
#'   \item \code{smoothing} (Negative Binomial):
#'   \itemize{
#'     \item \code{r}: overdispersion parameter draws.
#'   }
#' }
#' @note The data \code{y} may contain NAs, which will be treated with a simple imputation scheme
#' via an additional Gibbs sampling step. In general, rescaling \code{y} to have unit standard
#' deviation is recommended to avoid numerical issues when family is "gaussian".
#'
#' @examples
#' set.seed(200)
#' signal = c(rep(0, 50), rep(10, 50))
#' noise = rep(1, 100)
#' noise_var = rep(1, 100)
#' for (k in 2:100){
#'   noise_var[k] = exp(0.9*log(noise_var[k-1]) + rnorm(1, 0, 0.5))
#'   noise[k] = rnorm(1, 0, sqrt(noise_var[k])) }
#'
#' y = signal + noise
#' model_spec = dsp_spec(family = "gaussian", model = "changepoint",
#'                       D = 1, useAnom = TRUE, obsSV = "SV")
#' mcmc_output = dsp_fit(y, model_spec = model_spec, nsave = 500, nburn = 500)
#'
#' @export
dsp_fit = function(y, model_spec,
                      nsave = 1000, nburn = 1000, nskip = 4,
                      computeDIC = TRUE,
                      verbose = TRUE,
                      ...){
  if (!inherits(model_spec, "dsp_spec")) {
    stop("'model_spec' must be an object created by 'dsp_spec()'.")
  }

  family = model_spec$family
  model = model_spec$model
  # if(!is.null(model_spec$arguments$zero_error)){
  #   model = "smoothing_sparse"
  # }
  fitter_list <- list(
    gaussian = list(
      changepoint = abco, #changepoint model,
      regression = btf_reg, #Bayesian Trend Filter with Regression,
      smoothing = btf, #Bayesian Trend Filter,
      bspline = btf_bspline #Bayesian Trend Filter with Bspline
    ),
    negbinomial = list(
      smoothing = btf_nb #  Bayesian Trend Filter with Negative Binomial
    )
  )
  fitter = fitter_list[[family]][[model]]
  if (is.null(fitter)) {
    stop("No valid fitter function found for the specified family and model.")
  }
  ## This is a temporary case until we figure out what to do about the DIC thing

  if(model == "changepoint"){
    input_args <- c(model_spec$arguments, list(y = y, nsave = nsave, nburn = nburn,
                                               nskip = nskip,
                                               verbose = verbose, ...))
  }else{
    input_args <- c(model_spec$arguments, list(y = y, nsave = nsave, nburn = nburn,
                                               nskip = nskip, computeDIC = computeDIC,
                                               verbose = verbose, ...))
  }

  mcmc_output = do.call(fitter, input_args)
  structure(list(mcmc_output = mcmc_output,
                 DIC = mcmc_output[c("DIC", "p_d")],
                 mcpar = c(nsave = nsave, nburn = nburn, nskip = nskip),
                 model_spec = model_spec),
            class = c(c("dsp"), mcmc_output$class))

}
