# pmm2_main.R - Main module for PMM2 linear models

#' PMM2: Main function for PMM2 (S=2)
#'
#' Fits a linear model using the Polynomial Maximization Method (order 2),
#' which is robust to non-Gaussian errors.
#'
#' @param formula R formula for the model
#' @param data data.frame containing variables in the formula
#' @param max_iter integer: maximum number of iterations for the algorithm
#' @param tol numeric: tolerance for convergence
#' @param regularize logical: add small value to diagonal for numerical stability
#' @param reg_lambda numeric: regularization parameter (if regularize=TRUE)
#' @param na.action function for handling missing values, default is na.fail
#' @param weights optional weight vector (not yet implemented)
#' @param verbose logical: whether to print progress information
#'
#' @details
#' The PMM2 algorithm works as follows:
#'
#' 1. Fits ordinary least squares (OLS) regression to obtain initial estimates
#' 2. Computes central moments (m2, m3, m4) from OLS residuals
#' 3. Iteratively improves parameter estimates using a gradient-based approach
#'
#' PMM2 is especially useful when error terms are not Gaussian.
#'
#' @return S4 object of class \code{PMM2fit}
#' @export
#'
#' @examples
#' set.seed(123)
#' n <- 80
#' x <- rnorm(n)
#' y <- 2 + 3 * x + rt(n, df = 3)
#' dat <- data.frame(y = y, x = x)
#'
#' fit <- lm_pmm2(y ~ x, data = dat)
#' summary(fit, formula = y ~ x, data = dat)
lm_pmm2 <- function(formula, data,
                    max_iter=50, tol=1e-6,
                    regularize=TRUE, reg_lambda=1e-8,
                    na.action=na.fail, weights=NULL,
                    verbose=FALSE)
{
  # Capture call
  call <- match.call()

  # Validate input
  if(missing(formula) || missing(data)) {
    stop("Both 'formula' and 'data' must be provided")
  }

  if(!is.data.frame(data)) {
    stop("'data' must be a data.frame")
  }

  if(max_iter <= 0) {
    stop("'max_iter' must be positive")
  }

  if(tol <= 0) {
    stop("'tol' must be positive")
  }

  # Handle missing values
  if(!is.null(na.action)) {
    data <- na.action(data)
  }

  # Check for weights
  if(!is.null(weights)) {
    warning("Weights are not yet implemented in PMM2. Ignoring weights.")
  }

  # 1) OLS
  if(verbose) cat("Fitting initial OLS model...\n")

  mf <- model.frame(formula, data)
  X  <- model.matrix(formula, mf)
  y  <- model.response(mf)

  # Check for rank deficiency
  qr_X <- qr(X)
  if(qr_X$rank < ncol(X)) {
    warning("Design matrix has rank deficiency, some coefficients may not be estimable")
  }

  fit_ols <- lm.fit(x=X, y=y)
  b_ols   <- fit_ols$coefficients

  # Handle NA in OLS coefficients
  if(any(is.na(b_ols))) {
    stop("OLS fit resulted in NA coefficients. Check for multicollinearity.")
  }

  # Convert b_ols to numeric vector for consistent handling
  b_ols <- as.numeric(b_ols)

  # 2) OLS residuals => m2, m3, m4
  res_ols <- y - (X %*% b_ols)
  moments <- compute_moments(res_ols)
  m2 <- moments$m2
  m3 <- moments$m3
  m4 <- moments$m4

  if(verbose) {
    cat("Initial moments from OLS residuals:\n")
    cat("  m2 =", m2, "\n")
    cat("  m3 =", m3, "\n")
    cat("  m4 =", m4, "\n")
  }

  # Check for potential issues with moments
  if(m2 <= 0) {
    warning("Second central moment (m2) is not positive. Results may be unreliable.")
  }

  if(m4 <= m2^2) {
    warning("Fourth central moment (m4) is less than m2^2. This violates the basic inequality for probability distributions.")
  }

  # 3) Launch unified PMM2 algorithm
  if(verbose) cat("Starting PMM2 iterations...\n")

  out <- pmm2_algorithm(b_ols, X, y, m2, m3, m4,
                        max_iter = max_iter, tol = tol,
                        regularize = regularize, reg_lambda = reg_lambda,
                        verbose = verbose)

  # Extract results
  b_est   <- out$b
  conv    <- out$convergence
  iter    <- out$iterations
  final_res <- out$residuals

  if(verbose) {
    cat("PMM2 algorithm completed.\n")
    cat("  Converged:", conv, "\n")
    cat("  Iterations:", iter, "\n")
  }

  # Restore coefficient names from design matrix columns
  names(b_est) <- colnames(X)

  # Return S4 object with all results
  ans <- new("PMM2fit",
             coefficients = b_est,
             residuals = final_res,
             m2 = m2,
             m3 = m3,
             m4 = m4,
             convergence = conv,
             iterations = iter,
             call = call)
  attr(ans, "model_matrix") <- X
  attr(ans, "model_frame") <- mf
  attr(ans, "response") <- as.numeric(y)
  attr(ans, "data") <- data

  return(ans)
}

#' Extract coefficients from PMM2fit object
#'
#' @param object PMM2fit object
#' @param ... Additional arguments (not used)
#'
#' @return Vector of coefficients
#' @export
setMethod("coef", "PMM2fit",
          function(object, ...) {
            object@coefficients
          })

#' Extract residuals from PMM2fit object
#'
#' @param object PMM2fit object
#' @param ... Additional arguments (not used)
#'
#' @return Vector of residuals
#' @export
setMethod("residuals", "PMM2fit",
          function(object, ...) {
            object@residuals
          })

#' Extract fitted values from PMM2fit object
#'
#' @param object PMM2fit object
#' @param data Optional data source for model reconstruction, if object does not contain saved data
#' @param ... Additional arguments (not used)
#'
#' @return Vector of fitted values
#' @export
setMethod("fitted", "PMM2fit",
          function(object, data = NULL, ...) {
            fitted_values(object, data)
          })

#' Calculate AIC for PMM2fit object
#'
#' @param object PMM2fit object
#' @param ... Additional arguments (not used)
#' @param k Penalty per parameter to be used; default is 2
#'
#' @return AIC value
#' @export
setMethod("AIC", "PMM2fit",
          function(object, ..., k = 2) {
            res <- object@residuals
            n <- length(res)
            p <- length(object@coefficients)

            # Approximate log-likelihood
            ll <- -n/2 * log(sum(res^2)/n) - n/2 * (1 + log(2*pi))

            # AIC
            -2 * ll + k * p
          })

#' Calculate BIC for PMM2fit object
#'
#' @param object PMM2fit object
#' @param ... Additional arguments (not used)
#'
#' @return BIC value
#' @export
setMethod("BIC", "PMM2fit",
          function(object, ...) {
            res <- object@residuals
            n <- length(res)
            p <- length(object@coefficients)
            ll <- -n/2 * log(sum(res^2)/n) - n/2 * (1 + log(2*pi))
            -2 * ll + log(n) * p
          })

#' Extract log-likelihood from PMM2fit object
#'
#' Returns a Gaussian approximate log-likelihood, consistent with the AIC method.
#'
#' @param object PMM2fit object
#' @param ... Additional arguments (not used)
#'
#' @return Object of class \code{logLik}
#' @export
setMethod("logLik", "PMM2fit",
          function(object, ...) {
            res <- object@residuals
            n   <- length(res)
            p   <- length(object@coefficients)
            ll  <- -n/2 * log(sum(res^2)/n) - n/2 * (1 + log(2*pi))
            attr(ll, "df")   <- p   # consistent with the AIC method (sigma not counted)
            attr(ll, "nobs") <- n
            class(ll) <- "logLik"
            ll
          })

#' Number of observations in PMM2fit object
#'
#' @param object PMM2fit object
#' @param ... Additional arguments (not used)
#'
#' @return Integer number of observations
#' @export
setMethod("nobs", "PMM2fit",
          function(object, ...) {
            length(object@residuals)
          })

#' Variance-covariance matrix for PMM2fit object
#'
#' Returns the asymptotic covariance matrix \eqn{\sigma^2 g_2 (X^\top X)^{-1}},
#' where \eqn{g_2 = 1 - c_3^2 / (2 + c_4)} is the PMM2 variance reduction factor.
#' Under Gaussian errors \eqn{g_2 = 1} and the result equals the OLS covariance matrix.
#'
#' @param object PMM2fit object
#' @param ... Additional arguments (not used)
#'
#' @return Numeric matrix of the same dimension as the number of coefficients
#' @export
setMethod("vcov", "PMM2fit",
          function(object, ...) {
            X <- attr(object, "model_matrix")
            if (is.null(X))
              stop("model_matrix not found; refit the model with lm_pmm2()")
            vm <- pmm2_variance_matrices(X, object@m2, object@m3, object@m4)
            V  <- vm$pmm2
            nms <- names(object@coefficients)
            if (!is.null(nms) && length(nms) == nrow(V)) {
              rownames(V) <- nms
              colnames(V) <- nms
            }
            V
          })

#' Confidence intervals for PMM2fit coefficients
#'
#' Computes normal-approximation confidence intervals using the asymptotic
#' covariance matrix from \code{\link{vcov,PMM2fit-method}}.
#'
#' @param object PMM2fit object
#' @param parm character or integer vector of parameter names/indices to include
#' @param level confidence level (default 0.95)
#' @param ... Additional arguments (not used)
#'
#' @return Matrix with lower and upper confidence limits
#' @export
setMethod("confint", "PMM2fit",
          function(object, parm, level = 0.95, ...) {
            cf <- coef(object)
            se <- sqrt(diag(vcov(object)))
            a  <- (1 - level) / 2
            fac <- stats::qnorm(c(a, 1 - a))
            ci  <- cf + outer(se, fac)
            colnames(ci) <- paste0(format(100 * c(a, 1 - a), trim = TRUE), " %")
            if (!missing(parm)) ci <- ci[parm, , drop = FALSE]
            ci
          })

#' Plot diagnostic plots for PMM2fit object
#'
#' @param x PMM2fit object
#' @param y Not used (compatibility with generic)
#' @param which Set of plots to display (values 1-4)
#' @param ... Additional arguments passed to plotting functions
#'
#' @return Invisibly returns the input object
#' @export
setMethod("plot", signature(x = "PMM2fit", y = "missing"),
          function(x, y, which = 1:4, ...) {
            res <- as.numeric(x@residuals)
            fitted_vals <- tryCatch({
              fitted(x)
            }, error = function(e) {
              stored_X <- attr(x, "model_matrix")
              if(!is.null(stored_X)) {
                as.vector(stored_X %*% x@coefficients)
              } else {
                seq_along(res)
              }
            })

            which <- intersect(unique(which), 1:4)
            if(length(which) == 0) {
              which <- 1:4
            }

            old_par <- graphics::par(no.readonly = TRUE)
            on.exit(graphics::par(old_par))
            n_plots <- length(which)
            graphics::par(mfrow = c(2, 2))

            for(idx in which) {
              switch(idx,
                     {
                       graphics::plot(fitted_vals, res,
                                      main = "Residuals vs Fitted",
                                      xlab = "Fitted values",
                                      ylab = "Residuals", ...)
                       graphics::abline(h = 0, col = "red", lty = 2)
                     },
                     {
                       stats::qqnorm(res, main = "Normal Q-Q", ...)
                       stats::qqline(res, col = "red", lty = 2)
                     },
                     {
                       graphics::plot(seq_along(res), res, type = "l",
                                      main = "Residuals over Index",
                                      xlab = "Observation",
                                      ylab = "Residual", ...)
                       graphics::abline(h = 0, col = "red", lty = 2)
                     },
                     {
                       graphics::hist(res,
                                      main = "Residual Histogram",
                                      xlab = "Residuals",
                                      breaks = "FD", ...)
                     })
            }

            invisible(x)
          })

#' Helper function for extracting fitted values
#'
#' @param object PMM2fit object
#' @return Vector of fitted values
#'
#' @keywords internal
fitted_values <- function(object, data = NULL) {
  if(is.null(object@call)) {
    stop("PMM2fit object does not contain call information")
  }

  # Fallback to stored attributes
  stored_X <- attr(object, "model_matrix")
  stored_response <- attr(object, "response")
  stored_mf <- attr(object, "model_frame")
  stored_data <- attr(object, "data")

  if (!is.null(stored_X)) {
    fitted_attr <- tryCatch({
      as.vector(stored_X %*% object@coefficients)
    }, error = function(e) NULL)
    if (!is.null(fitted_attr)) {
      return(fitted_attr)
    }
  }

  if (!is.null(stored_response) &&
      length(stored_response) == length(object@residuals)) {
    return(as.vector(stored_response - object@residuals))
  }

  # Try to reconstruct original data
  data_to_use <- data
  if(is.null(data_to_use)) {
    if (!is.null(stored_mf)) {
      data_to_use <- stored_mf
    } else if (!is.null(stored_data)) {
      data_to_use <- stored_data
    } else {
      # Try to get data from call, but safely handle possible errors
      tryCatch({
        data_to_use <- eval(object@call$data, envir = parent.frame())
      }, error = function(e) {
        if(is.null(data)) {
          stop("Failed to retrieve data from object. Please provide 'data' parameter.")
        }
      })
    }
  }

  if(is.null(data_to_use)) {
    stop("Data frame is required to compute fitted values")
  }

  # Reconstruct formula
  formula <- eval(object@call$formula)

  # Safe construction of design matrix
  tryCatch({
    mf <- model.frame(formula, data_to_use)
    X <- model.matrix(formula, mf)

    # Calculate fitted values
    fitted <- as.vector(X %*% object@coefficients)
    return(fitted)
  }, error = function(e) {
    stop("Error computing fitted values: ", e$message)
  })
}

#' Compare PMM2 with OLS
#'
#' @param formula Model formula
#' @param data Data frame
#' @param pmm2_args List of arguments to pass to lm_pmm2()
#'
#' @return List with OLS and PMM2 fit objects
#'
#' @examples
#' \donttest{
#' result <- compare_with_ols(mpg ~ wt, data = mtcars)
#' }
#'
#' @export
compare_with_ols <- function(formula, data, pmm2_args = list()) {
  # Fit OLS model
  fit_ols <- lm(formula, data)

  # Fit PMM2 model with default or specified arguments
  args <- c(list(formula = formula, data = data), pmm2_args)
  fit_pmm2 <- do.call(lm_pmm2, args)

  # Extract and compare coefficients
  coef_ols <- coef(fit_ols)
  coef_pmm2 <- coef(fit_pmm2)

  # Check if PMM2 coefficient names are set correctly
  if(is.null(names(coef_pmm2)) || all(names(coef_pmm2) == "")) {
    # If names are not set, use names from OLS
    if(length(coef_ols) == length(coef_pmm2)) {
      names(coef_pmm2) <- names(coef_ols)
    } else {
      # If lengths differ, use generated names
      names(coef_pmm2) <- paste0("coef", seq_along(coef_pmm2))
    }
  }

  # Calculate residual statistics
  res_ols <- residuals(fit_ols)
  res_pmm2 <- residuals(fit_pmm2)

  res_stats <- data.frame(
    Method = c("OLS", "PMM2"),
    RSS = c(sum(res_ols^2), sum(res_pmm2^2)),
    MAE = c(mean(abs(res_ols)), mean(abs(res_pmm2))),
    Skewness = c(pmm_skewness(res_ols), pmm_skewness(res_pmm2)),
    Kurtosis = c(pmm_kurtosis(res_ols), pmm_kurtosis(res_pmm2))
  )

  # Create coefficient comparison table
  # Use all unique coefficient names
  all_coef_names <- unique(c(names(coef_ols), names(coef_pmm2)))

  coef_table <- data.frame(
    Coefficient = all_coef_names,
    OLS = numeric(length(all_coef_names)),
    PMM2 = numeric(length(all_coef_names)),
    Diff_Percent = numeric(length(all_coef_names))
  )

  # Fill values for OLS
  for(i in seq_along(all_coef_names)) {
    name <- all_coef_names[i]
    if(name %in% names(coef_ols)) {
      coef_table$OLS[i] <- coef_ols[name]
    } else {
      coef_table$OLS[i] <- NA
    }

    if(name %in% names(coef_pmm2)) {
      coef_table$PMM2[i] <- coef_pmm2[name]

      # Calculate percent difference only if both values exist
      if(!is.na(coef_table$OLS[i]) && coef_table$OLS[i] != 0) {
        coef_table$Diff_Percent[i] <- 100 * (coef_table$PMM2[i] - coef_table$OLS[i]) / abs(coef_table$OLS[i])
      }
    } else {
      coef_table$PMM2[i] <- NA
      coef_table$Diff_Percent[i] <- NA
    }
  }

  return(list(
    ols = fit_ols,
    pmm2 = fit_pmm2,
    coefficients = coef_table,
    residual_stats = res_stats
  ))
}


#' Prediction method for PMM2fit objects
#'
#' Computes predictions for new data using a fitted PMM2 model. The method
#' extracts the formula from the fitted model, constructs a design matrix
#' from the new data, and computes predictions via matrix multiplication.
#' This approach works with arbitrary variable names and model specifications.
#'
#' @param object PMM2fit object returned by \code{lm_pmm2()}
#' @param newdata Data frame containing the predictor variables with the same
#'   names as those used in the original model fit. Required parameter.
#' @param debug Logical value indicating whether to output debug information
#'   about the prediction process. Default is FALSE.
#' @param ... Additional arguments (currently not used)
#'
#' @return Numeric vector of predicted values for the observations in \code{newdata}
#'
#' @details
#' The prediction is computed as \code{X \%*\% beta} where X is the design matrix
#' constructed from \code{newdata} using the model formula, and beta are the
#' estimated coefficients. The method automatically handles:
#' \itemize{
#'   \item Models with intercepts and without
#'   \item Arbitrary variable names (not limited to "x1", "x2", etc.)
#'   \item Interaction terms and transformations specified in the formula
#'   \item Automatic coefficient reordering to match design matrix columns
#' }
#'
#' @examples
#' \donttest{
#' # Fit model
#' fit <- lm_pmm2(mpg ~ wt + hp, data = mtcars)
#'
#' # Predict on new data
#' newdata <- data.frame(wt = c(2.5, 3.0), hp = c(100, 150))
#' predictions <- predict(fit, newdata = newdata)
#' }
#'
#' @export
setMethod("predict", "PMM2fit",
          function(object, newdata = NULL, debug = FALSE, ...) {
            if(is.null(newdata)) {
              stop("Parameter newdata must be provided")
            }

            if(is.null(object@call)) {
              stop("PMM2fit object does not contain call information")
            }

            # Extract formula from call
            formula <- eval(object@call$formula)

            if(debug) {
              cat("Formula:", deparse(formula), "\n")
              cat("Coefficients:", paste(names(object@coefficients), "=", object@coefficients, collapse=", "), "\n")
              cat("Size of newdata:", nrow(newdata), "x", ncol(newdata), "\n")
              cat("Variables in newdata:", paste(names(newdata), collapse=", "), "\n")
            }

            # Simple solution - directly use right side of formula to build design matrix
            rhs <- formula[[3]]
            design_formula <- as.formula(paste("~", deparse(rhs)))

            if(debug) {
              cat("Design formula:", deparse(design_formula), "\n")
            }

            # Create design matrix
            X <- model.matrix(design_formula, newdata)

            if(debug) {
              cat("Size of design matrix:", nrow(X), "x", ncol(X), "\n")
              cat("Columns of design matrix:", paste(colnames(X), collapse=", "), "\n")
            }

            # Ensure coefficient names match design matrix columns
            if(is.null(names(object@coefficients)) || all(names(object@coefficients) == "")) {
              if(debug) {
                cat("Coefficient names are missing or empty. Using design matrix column names.\n")
              }

              expected_names <- colnames(X)
              if(length(expected_names) == length(object@coefficients)) {
                names(object@coefficients) <- expected_names
              } else {
                stop("Number of coefficients does not match number of columns in design matrix.")
              }
            }

            if(debug) {
              cat("Coefficient names:", paste(names(object@coefficients), collapse=", "), "\n")
            }

            # Verify that coefficient names match design matrix columns
            if(!identical(names(object@coefficients), colnames(X))) {
              if(debug) {
                cat("Warning: Coefficient names do not match design matrix columns.\n")
                cat("Coefficients:", paste(names(object@coefficients), collapse=", "), "\n")
                cat("Design matrix:", paste(colnames(X), collapse=", "), "\n")
              }

              # Try to reorder coefficients to match design matrix
              if(all(colnames(X) %in% names(object@coefficients))) {
                object@coefficients <- object@coefficients[colnames(X)]
                if(debug) {
                  cat("Coefficients reordered to match design matrix.\n")
                }
              } else {
                stop("Design matrix columns do not match coefficient names. Cannot compute predictions.")
              }
            }

            # Calculate predictions using matrix multiplication
            # This is the general approach that works for all cases
            predictions <- as.vector(X %*% object@coefficients)

            if(debug) {
              cat("Size of prediction vector:", length(predictions), "\n")
              cat("First few predictions:", paste(head(predictions), collapse=", "), "\n")
            }

            return(predictions)
          })
