#' Lagrange-multiplier optimal phase 2 design for a two-phase genetic association study.
#'
#' Computes stratum-specific sampling fractions that maximise the anticipated
#' Fisher information for the regression parameters of interest subject to a
#' fixed total phase 2 sample size.  The optimisation is performed by the
#' COBYLA algorithm (\code{nloptr}) with a derivative-free fallback
#' (\code{dfoptim::hjkb}) if COBYLA fails to converge.
#'
#' @param formula A two-sided \code{\link{formula}} specifying the regression
#'   model of interest, e.g.\code{Y ~ G + Z}.  All variables except
#'   \code{miscov} must be present in \code{data}.
#' @param miscov Right-hand-side formula identifying the missing-by-design
#'   covariate(s) (the sequencing variant), e.g.\code{~ G}.
#' @param auxvar Right-hand-side formula identifying the auxiliary variable
#'   (the GWAS tag SNP), e.g.\code{~ Z}.  Must be present in \code{data}.
#' @param strata Right-hand-side formula defining the stratification variables
#'   used to cross-tabulate \code{data}, e.g.\code{~ Z + S}.  All variables
#'   must be present in \code{data}.
#' @param family GLM family object (see \code{\link{family}}).  Passed
#'   directly to the information-matrix calculations.  Default:
#'   \code{gaussian()}.
#' @param n Integer.  Target phase 2 sample size (number of individuals to
#'   be sequenced).
#' @param data A \code{data.frame} containing the phase 1 data.  Must include
#'   all variables in \code{formula} (except those in \code{miscov}),
#'   \code{auxvar}, and \code{strata}.
#' @param beta Named or unnamed numeric vector of assumed regression
#'   coefficients, with one element per term in \code{formula} (including
#'   the intercept).  These define the operating point at which the
#'   Fisher information is evaluated.  Setting \code{beta[G] = 0} triggers
#'   the score-test (null-model) criterion; any non-zero value triggers the
#'   Wald (alternative-model) criterion.
#' @param p_gz A \code{data.frame} with the assumed joint distribution of
#'   (\code{miscov}, \code{auxvar}).  Must contain one column per variable
#'   in \code{miscov} and \code{auxvar} plus a column \code{q} holding
#'   non-negative probability masses that sum to 1.  Rows with
#'   \code{q <= 0} are silently removed.  Typically generated by
#'   \code{\link{twoPhaseDesign}} from \code{maf_G} and \code{LD}, or
#'   supplied directly by the user.
#' @param disp Numeric dispersion parameter.  For Gaussian and Gamma
#'   families this is the variance (\eqn{\sigma^2}) or dispersion
#'   (\eqn{\phi}), respectively; for binomial and Poisson it is fixed at
#'   1.  If \code{NULL} (default) it is estimated as \code{var(Y)}.
#' @param optimMeasure Character string specifying the optimality criterion.
#'   One of \code{"Par-spec"} (parameter-specific: minimise the asymptotic
#'   variance of \code{beta[K.idx]}), \code{"A-opt"} (A-optimality:
#'   minimise \code{trace(FIM^{-1})}), or \code{"D-opt"} (D-optimality:
#'   maximise \code{log det(FIM)}).
#' @param K.idx Integer index (or vector of indices) into \code{beta}
#'   identifying the parameter(s) of primary interest for
#'   \code{optimMeasure = "Par-spec"}.  Ignored for A- and D-optimality.
#'   Required when \code{optimMeasure = "Par-spec"}.
#' @param min.nk Optional integer scalar or vector of length equal to the
#'   number of strata.  Specifies the minimum number of individuals to
#'   sample from each stratum.  If a scalar, the same minimum is applied to
#'   all strata.  Default: \code{NULL} (no minimum, i.e. zero).
#' @param logical.sub Optional logical vector of length \code{nrow(data)}.
#'   If supplied, only the rows where \code{logical.sub = TRUE} are used in
#'   the information calculations (the full \code{data} is still returned).
#'   Useful for sensitivity analyses on subsets.  Default: \code{NULL}.
#'
#' @details
#' The function computes, for each stratum defined by \code{strata}, the
#' stratum-specific sampling fraction \eqn{\pi_k = n_k / N_k} that maximises
#' the chosen information criterion.  The optimisation is constrained so that
#' (i) the total selected sample equals \code{n}, (ii) no stratum is
#' over-sampled (\eqn{n_k \le N_k}), and (iii) the \code{min.nk} lower
#' bounds are respected.
#'
#' The baseline Fisher information uses the empirical score outer-product
#' estimator \eqn{I_3/N} (see \code{\link{optimTP.GA}} for details).
#'
#' @return A \code{data.frame} with one row per stratum (as defined by
#'   \code{strata}) and the following columns from
#'   \code{xtabs(strata, data)}:
#'   \describe{
#'     \item{(strata variables)}{The stratification variable values.}
#'     \item{Freq}{Stratum size \eqn{N_k}.}
#'     \item{prR_cond_optim}{Optimal conditional sampling probability
#'       \eqn{n_k / n} (proportional allocation weight).}
#'     \item{convergence}{Convergence code from the solver (\code{0} = success;
#'       codes from \code{nloptr::cobyla} or \code{dfoptim::hjkb}).}
#'   }
#'
#' @seealso \code{\link{twoPhaseDesign}}, \code{\link{optimTP.GA}}
#'
#' @examples
#' \donttest{
#' data <- DataGeneration_TPD(N = 1000)
#' data <- data[, c("Y", "Z")]  # phase-1 data: no G column
#' p_gz <- data.frame(Z = c(0, 0, 0, 1, 1, 1, 2, 2, 2), # theoretical joint distribution
#'         G1 = c(0, 1, 2, 0, 1, 2, 0, 1, 2), # between G and Z (LD = 0.75, MAF_G = 0.2, MAF_Z = 0.3)
#'         q = c(0.486, 0.004, 0, 0.143, 0.276, 0.001, 0.011, 0.04, 0.039)) 
#'         
#' data$S <- as.numeric(cut(data$Y,
#'             breaks = quantile(data$Y, probs = c(0, 0.4, 0.8, 1)),
#'             include.lowest = TRUE))
#'
#' res_lm <- optimTP.LM(
#'   formula     = Y ~ Z + G1,
#'   miscov      = ~ G1,
#'   auxvar      = ~ Z,
#'   strata      = ~ Z + S,
#'   family      = gaussian(),
#'   n           = 200,
#'   data        = data,
#'   beta        = c(2, 0.5, 0.1),
#'   p_gz        = p_gz,
#'   optimMeasure = "Par-spec",
#'   K.idx       = 2L
#' )
#' head(res_lm)
#' }
#'
#' @export
#' @importFrom nloptr cobyla
#' @importFrom dfoptim hjkb
#' @importFrom utils combn stack tail
#' @importFrom enrichwith enrich
optimTP.LM  <-  function(formula,miscov,auxvar,strata,family,n,data,beta,p_gz,disp=NULL,optimMeasure,K.idx=NULL,min.nk=NULL,logical.sub=NULL){

  if (is.character(family))
    family <- get(family, mode = "function", envir = parent.frame())
  if (is.function(family))
    family <- family()
  if (is.null(family$family)) {
    message(deparse(family))
    stop("'family' not recognized")
  }

  ### TODO: check if all unique values in data[,Z_] are in p_gz

  terms_dat <- unique( c(all.vars(formula), attr(terms(miscov), "term.labels") , attr(terms(auxvar), "term.labels"), attr(terms(strata), "term.labels") , attr(terms(auxvar), "term.labels")) )

  Y_ <- as.character(formula[[2]])
  G_ <- all.vars(miscov)
  Z_ <- all.vars(auxvar)
  S_ <- all.vars(strata)

  if( is.null(disp) ){
    if( family$family %in% c("gaussian","Gamma") ){
      disp = var(data[,Y_])
    }else disp=1
  }

  ## remove enties with zero or negative probabilities
  p_gz <- p_gz[p_gz$q>0,]

  N=NROW(data)
  rho=n/N
  if( rho<0.1 )warning("The phase 2 sample size (n2) is smaller than recommended. Please proceed with caution.")

  stratadf <- data.frame(xtabs(strata,data=data))
  #Xs <- model.matrix(strata, stratadf)
  stratadf[S_] <- lapply(stratadf[S_], function(x)as.numeric(as.character(x))) ## may need it for data.table
  g.vals <- unique(p_gz[,G_])
  z.vals <- unique(p_gz[,Z_])

  if( !is.null(logical.sub) ){
    data <- data[logical.sub,]
    N <- NROW(data)
  }

  dat_ <- cbind(id_=rep(seq(1,N,by=1), each=NROW(g.vals)),data[rep(seq_len(N), each=NROW(g.vals)), ],matrix(rep(t(g.vals),N), ncol=NCOL(g.vals), byrow=TRUE))
  names(dat_)[(ncol(dat_)-NCOL(g.vals)+1):ncol(dat_)] <- G_
  dat_ <- droplevels(dat_)
  dat_$ord_ <- seq_len(nrow(dat_))

  wgs <- omega1_g(formula,family,dat_,beta,p_gz,disp,G_,Z_) # P(G|Y,Z) and P(Y,Z)

  dat_$wg_ <- wgs$wg
  dat_ <- dat_[!is.na(dat_$wg_),]

  ## Baseline FIM: I3/N (empirical score outer product; valid large-N approximation)
  obsIM_all_lm <- obsIMf1(formula,family,dat_,beta,p_gz,disp,G_,Z_,by.id = FALSE)
  M1 <- obsIM_all_lm$I3 / N

  ## Per-stratum I3 contributions (phase-2 improvement, by stratum)
  IM3_subs <- do.call(rbind,by(dat_, INDICES = dat_[,S_],function(df.by){
    Nk <- N
    IM <- obsIMf1(formula,family,df.by,beta,p_gz,disp,G_,Z_)
    data.frame(df.by[1,S_,drop=FALSE],t(as.vector(IM$I3)/Nk))
  }))

  ## Per-stratum I2 contributions
  IM2_subs <- do.call(rbind,by(dat_, INDICES = dat_[,S_],function(df.by){
    Nk <- N
    IM <- obsIMf1(formula,family,df.by,beta,p_gz,disp,G_,Z_)
    data.frame(df.by[1,S_,drop=FALSE],t(as.vector(IM$I2)/Nk))
  }))

  lenS_ <- length(S_)
  npars <- NCOL(M1)

  ### Indicator for the beta parameters of interest
  Ind <- rep(FALSE,npars-1)
  Betas_ids <- which(grepl(paste0(paste0("^",G_),collapse="|"),colnames(model.matrix(formula,dat_[1,]))))
  Ind[Betas_ids] <- TRUE

  if( optimMeasure=="Par-spec" & is.null(K.idx) ) stop("For a parameter-specific criterion K.idx must be provided.")
  ## asign these functions to the execution environment
  optMsuref=.optMsure(optimMeasure);  environment(optMsuref)=environment()

  ### parameters for the constained optimization
  # if( !all.equal(pS_[S_],stratadf[,S_],attributes=FALSE,check.attributes = FALSE) ) stop("Strata groups in pS and stratadf don't match!")
  # ais <- (stratadf$Freq)#/sum(stratadf$Freq)
  # ais <- N*pS_$pS_#/sum(stratadf$Freq)
  npis <- nrow(stratadf)-1
  rho <- n
  upperLM <- as.numeric(stratadf$Freq) #rep(1,npis+1)
  # M1inv <- MASS::ginv(M1)
  # In <- diag(npars)
  D <- rbind(diag(npars-1),0)
  # Dinv <- MASS::ginv(D)
  # Dtinv  <- MASS::ginv((t(D)))
  ### result on how to calcuate the inverse of A+B
  # https://math.stackexchange.com/questions/17776/inverse-of-the-sum-of-matrices
  # (A+B)^{-1} = A^{-1}(I-(I+BA^{-1})^{-1}BA^{-1})


  ### The LM approach is going to get optimal values for nk/P(R=1,Z,K) as opposed to pr(R=1|Y,Z). In additio the  objective function is determined according to whether the design regression parameters are under the null or alternative

  if( all(beta[Ind]==0) ){
    ObjFun <- function(pis){ # pis = (upperLM*n/N)[-1]
      pis1 <- c(0,pis)
      pis1[1] <- (rho-sum(pis1))
      stratadf$prReq1YZ <- (pis1/N)/stratadf$Freq

      IM2b_subs <- as.data.frame(merge(as.data.table(IM2_subs), as.data.table(stratadf[,c(S_,"prReq1YZ")]), by=S_, sort = FALSE))
      IM3b_subs <- as.data.frame(merge(as.data.table(IM3_subs), as.data.table(stratadf[,c(S_,"prReq1YZ")]), by=S_, sort = FALSE))

      FIM <- M1 + matrix(colSums(IM2b_subs[,(lenS_+1):(npars^2+lenS_)]*IM2b_subs$prReq1YZ),ncol=npars) - matrix(colSums(IM3b_subs[,(lenS_+1):(npars^2+lenS_)]*IM3b_subs$prReq1YZ),ncol=npars)

      ## to deal with the constraint on the p_g's
      F1 <-  crossprod(D, FIM %*% D)

      ###  Score test variance
      invF1_noInd <- ginv(F1[!Ind,!Ind])
      V1 <- F1[Ind,Ind] - ( F1[Ind,!Ind] %*% invF1_noInd %*% F1[!Ind,Ind] )
      invV1 <- solve( V1 )

      return( optMsuref(invV1,K.idx) )
    }
  }else {
    ObjFun <- function(pis){ # pis = (upperLM*n/N)[-1]
      pis1 <- c(0,pis)
      pis1[1] <- (rho-sum(pis1))
      stratadf$prReq1YZ <- (pis1/N)/stratadf$Freq

      IM2b_subs <- as.data.frame(merge(as.data.table(IM2_subs), as.data.table(stratadf[,c(S_,"prReq1YZ")]), by=S_, sort = FALSE))
      IM3b_subs <- as.data.frame(merge(as.data.table(IM3_subs), as.data.table(stratadf[,c(S_,"prReq1YZ")]), by=S_, sort = FALSE))

      FIM <- M1 + matrix(colSums(IM2b_subs[,(lenS_+1):(npars^2+lenS_)]*IM2b_subs$prReq1YZ),ncol=npars) - matrix(colSums(IM3b_subs[,(lenS_+1):(npars^2+lenS_)]*IM3b_subs$prReq1YZ),ncol=npars)

      ## to deal with the constraint on the p_g's
      F1 <-  crossprod(D, FIM %*% D)

      # B <- matrix(colSums(IM2b_subs[,(lenS_+1):(npars^2+lenS_)]*IM2b_subs$prReq1YZ),ncol=npars) - matrix(colSums(IM3b_subs[,(lenS_+1):(npars^2+lenS_)]*IM3b_subs$prReq1YZ),ncol=npars)
      #
      #
      # FIMinv <- M1inv %*% (In - solve(In+B%*%M1inv)) %*% (B%*%M1inv)
      #
      # vcov = Dinv %*% FIMinv %*% Dtinv

      vcov = solve(F1)

      return( optMsuref(vcov,K.idx) )

    }
  }


  if( is.null(min.nk) ){
    minLB.val <- 0 #ifelse(rho<0.05,rho*0.01,0.05)
    minLB <- rep(minLB.val,npis+1)
  }else{
    if( !length(min.nk) %in% c(1,npis+1) ) stop("The number of elements of min.nk is either 1 or ",npis+1)
    if( any(min.nk<0) )stop("The minimum stratum size must greater or equal than zero")
    if( any(min.nk>stratadf$Freq) )stop("The minimum stratum size must not be greater than the strata sample sizes determined by strata")
    minLB <- min.nk
  }
  initX0 <- (upperLM*n/N)[-1] #rep(rho,npis)
  initX1 <- (upperLM*n/N)  #rep(rho,npis+1)

  ObjFun_in <- function(pis){
    ui <- rbind(rep(-1,npis),rep(1,npis))
    ci <- c(-rho,0)
    return(as.numeric(ui%*%pis-ci))
  }

  ObjFun_eq <- function(pis1){
    return(sum(pis1)-rho)
  }

  ### Mind the warning: "For consistency with the rest of the package the inequality sign may be switched from >= to <= in a future nloptr version."

  ## To disable this warning and avoid printing do the following. However, new versions of the package may return an error when  this is enforced.
  options(nloptr.show.inequality.warning = FALSE)

  ### Note that the function can be changed when the change is implemented using
  # hin <- function(x) (-1)*f2(x, ...)  # NLOPT expects hin <= 0

  ### use this one for base, if an error then use the rest
  sol <- cobyla(initX0, fn = ObjFun, lower = minLB[-1], upper = upperLM[-1], hin = ObjFun_in, nl.info = FALSE, control = list(xtol_rel = 1e-6, maxeval = 10000))

  # return the option to the original state
  options(nloptr.show.inequality.warning = TRUE)

  if( !(sol$convergence>0 & sol$convergence<5) ){
    if( sol$convergence==5 ) message("The maximum number of evaluations in nloptr:::cobyla() has been reached. Trying dfoptim:::hjkb() now. You could also increase the number of evaluations via maxeval (although it is fairly large already.)")
    if( sol$convergence<0 ) message(paste0("Convergence in nloptr:::cobyla() was not achieved; error code ",sol$convergence,". Trying dfoptim:::hjkb() now."))

    ObjFunb <- function(pisA, k){ # pisA <- (upperLM*n/N)
      nstar <- sum(pisA)
      return( ObjFun(pisA[-1]) + k*(nstar-rho)^2 )
    }

    sol1 <- hjkb(initX1, fn=ObjFunb, lower=minLB, upper=upperLM, k=0.01, control=list(maxfeval=10000))
    sol1 <- tryCatch(hjkb(sol1$par, fn=ObjFunb, lower=minLB, upper=upperLM, k=100, control=list(maxfeval=10000)), error=function(e){ NULL })
    if( is.null(sol1) ){
      solB <- c(0,sol$par)
      solB[1] <- (rho-sum(solB))

      sol1 <- hjkb(solB, fn=ObjFunb, lower=minLB, upper=upperLM, k=100, control=list(maxfeval=10000))
    }

    if( sol1$convergence!=0 ){
      warning(paste0("Convergence in dfoptim:::hjkb() was not achieved. Error code ",sol1$convergence,". Returning last par value."))
    }

    stratadf$convergence <- c(sol1$convergence,rep(NA,nrow(stratadf)-1))
    pisopt <- sol1$par

  } else {
    stratadf$convergence <- c(sol$convergence,rep(NA,nrow(stratadf)-1))
    pisopt <- c(0,sol$par)
    pisopt[1] <- (rho-sum(pisopt))
  }

  stratadf$prR_cond_optim <- pisopt/sum(pisopt)

  return(stratadf)
}


#' Genetic-algorithm optimal phase 2 design for a two-phase genetic association study.
#'
#' Uses a genetic algorithm (GA) to search over individual-level phase 2
#' selection indicators \eqn{R_i \in \{0,1\}} for the subset of size
#' \code{n} that maximises the anticipated Fisher information for the
#' regression parameters of interest.  Wraps \code{\link[kofnGA]{kofnGA}}.
#' When a parallel cluster is supplied, multiple independent GA runs are
#' executed in parallel and the best solution across all runs is returned.
#'
#' @param formula A two-sided \code{\link{formula}} specifying the regression
#'   model of interest, e.g. \code{Y ~ G + Z}.  All variables except those in
#'   \code{miscov} must be present in \code{data}.
#' @param miscov Right-hand-side formula identifying the missing-by-design
#'   covariate(s) (the sequencing variant), e.g. \code{~ G}.
#' @param auxvar Right-hand-side formula identifying the auxiliary variable
#'   (the GWAS tag SNP), e.g. \code{~ Z}.  Must be present in \code{data}.
#' @param family GLM family object (see \code{\link{family}}).  Default:
#'   \code{gaussian()}.
#' @param n Integer.  Target phase 2 sample size.
#' @param data A \code{data.frame} containing the phase 1 data.  Must include
#'   all variables in \code{formula} (except those in \code{miscov}) and
#'   \code{auxvar}.
#' @param beta Numeric vector of assumed regression coefficients, one per
#'   term in \code{formula} (including the intercept).  Setting
#'   \code{beta[G] = 0} triggers the score-test (null-model) criterion; any
#'   non-zero value triggers the Wald (alternative-model) criterion.
#' @param p_gz A \code{data.frame} with the assumed joint distribution of
#'   (\code{miscov}, \code{auxvar}).  Must contain columns for each variable
#'   in \code{miscov} and \code{auxvar} and a column \code{q} of
#'   non-negative probability masses summing to 1.  Rows with \code{q <= 0}
#'   are silently removed.
#' @param disp Numeric dispersion parameter (variance for Gaussian/Gamma;
#'   fixed at 1 for binomial/Poisson).  If \code{NULL} (default),
#'   estimated as \code{var(Y)}.
#' @param ga.popsize Integer.  GA population size (number of candidate
#'   solutions maintained per generation).
#' @param ga.propelit Numeric in \eqn{(0, 1)}.  Proportion of the population
#'   that is elite (carried forward unchanged to the next generation).
#'   The elite count is \code{ceiling(ga.popsize * ga.propelit)}.
#' @param ga.proptourney Numeric in \eqn{(0, 1)}.  Proportion of the
#'   population used in each tournament selection.  The tournament size is
#'   \code{max(ceiling(ga.popsize * ga.proptourney), 2)}.
#' @param ga.ngen Integer.  Number of GA generations to run.
#' @param ga.mutrate Numeric in \eqn{[0, 1]}.  Per-index mutation probability.
#'   At each mutation step, each of the \code{n} selected indices is
#'   independently replaced by a randomly chosen non-selected index with
#'   probability \code{ga.mutrate}.
#' @param ga.initpop Optional \code{ga.popsize}-by-\code{n} integer matrix of
#'   starting solutions (each row is a candidate subset of size \code{n}).
#'   If \code{NULL} (default), the initial population is generated randomly.
#'   When \code{cluster} is provided, the rows are split evenly across
#'   workers.
#' @param optimMeasure Character string specifying the optimality criterion.
#'   One of \code{"Par-spec"} (parameter-specific), \code{"A-opt"}
#'   (A-optimality), or \code{"D-opt"} (D-optimality).
#' @param K.idx Integer index (or vector of indices) identifying the
#'   parameter(s) of primary interest.  Required when
#'   \code{optimMeasure = "Par-spec"}.
#' @param seed Integer random seed, set before GA initialisation.  Default: 1.
#' @param verbose Integer.  If \code{> 0}, the best objective value is printed
#'   every \code{verbose} generations.  Default: \code{0} (silent).
#' @param cluster An optional parallel cluster object created by
#'   \code{\link[parallel]{makeCluster}}.  When provided,
#'   \code{length(cluster)} independent GA searches are run in parallel
#'   (each worker receives an equal share of \code{ga.initpop} rows, or a
#'   random initial population if \code{ga.initpop = NULL}), and the best
#'   solution across all workers is returned.  The caller is responsible for
#'   starting and stopping the cluster.  Default: \code{NULL} (sequential
#'   single-run GA).
#'
#' @details
#' The fitness function evaluates the anticipated Fisher information
#' matrix (FIM) for a candidate selection indicator \eqn{R}:
#' \deqn{\mathrm{FIM}(R) = I_3/N + \sum_{i: R_i=1}(I_{2,i} - I_{3,i})/N,}
#' where \eqn{I_3} is the empirical score outer-product estimator (the
#' numerically stable approximation to the full Louis (1982) observed FIM
#' valid for large \eqn{N}, since \eqn{I_1 \approx I_2} in the genotype
#' expansion context), and \eqn{I_{2,i}}, \eqn{I_{3,i}} are the
#' per-individual complete-data score outer product and observed-data
#' score outer product, respectively.  Under \eqn{H_0} (\code{beta[G] = 0})
#' the score-test Schur complement is optimised; under \eqn{H_a} the full
#' inverse FIM is optimised.
#'
#' @return An object of class \code{"kofnGA"} (see
#'   \code{\link[kofnGA]{kofnGA}}) with the following elements:
#'   \describe{
#'     \item{bestsol}{Integer vector of length \code{n} giving the row indices
#'       (in \code{data}) of the selected phase 2 individuals.}
#'     \item{bestobj}{Numeric value of the optimality criterion at
#'       \code{bestsol}.}
#'     \item{pop}{Final population matrix (\code{ga.popsize} rows,
#'       \code{n} columns), sorted by fitness.}
#'     \item{obj}{Objective values for each row of \code{pop}.}
#'     \item{old}{List with per-generation \code{best}, \code{obj}, and
#'       \code{avg} records.}
#'     \item{time.elapsed}{Named numeric vector with the pre-processing
#'       (\code{preproc}) and GA (\code{GA}) times in seconds.}
#'   }
#'
#' @seealso \code{\link[kofnGA]{kofnGA}}, \code{\link{twoPhaseDesign}},
#'   \code{\link{optimTP.LM}}
#'
#' @examples
#' \donttest{
#' data <- DataGeneration_TPD(N = 500)
#' data <- data[, c("Y", "Z")]  # phase-1 data: no G column
#' p_gz <- data.frame(Z = c(0, 0, 0, 1, 1, 1, 2, 2, 2), # theoretical joint distribution
#'         G1 = c(0, 1, 2, 0, 1, 2, 0, 1, 2), # between G and Z (LD = 0.75, MAF_G = 0.2, MAF_Z = 0.3)
#'         q = c(0.486, 0.004, 0, 0.143, 0.276, 0.001, 0.011, 0.04, 0.039)) 
#'
#' res_ga <- optimTP.GA(
#'   formula      = Y ~ Z + G1,
#'   miscov       = ~ G1,
#'   auxvar       = ~ Z,
#'   family       = gaussian(),
#'   n            = 100,
#'   data         = data,
#'   beta         = c(2, 0.5, 0.1),
#'   p_gz         = p_gz,
#'   ga.popsize   = 40,
#'   ga.propelit  = 0.1,
#'   ga.proptourney = 0.1,
#'   ga.ngen      = 50,
#'   ga.mutrate   = 0.01,
#'   optimMeasure = "Par-spec",
#'   K.idx        = 2L
#' )
#' cat("Best objective:", res_ga$bestobj, "\n")
#' cat("Selected indices (first 10):", head(res_ga$bestsol, 10), "\n")
#'
#' ## Parallel processing (2 workers)
#' cl <- parallel::makeCluster(2L)
#' res_par <- optimTP.GA(
#'   formula      = Y ~ Z + G1,
#'   miscov       = ~ G1,
#'   auxvar       = ~ Z,
#'   family       = gaussian(),
#'   n            = 100,
#'   data         = data,
#'   beta         = c(2, 0.5, 0.1),
#'   p_gz         = p_gz,
#'   ga.popsize   = 40,
#'   ga.propelit  = 0.1,
#'   ga.proptourney = 0.1,
#'   ga.ngen      = 50,
#'   ga.mutrate   = 0.01,
#'   optimMeasure = "Par-spec",
#'   K.idx        = 2L,
#'   cluster      = cl
#' )
#' parallel::stopCluster(cl)
#' cat("Best objective (parallel):", res_par$bestobj, "\n")
#' }
#'
#' @importFrom kofnGA kofnGA
#' @importFrom parallel clusterEvalQ
#' @export
optimTP.GA  <-  function(formula,miscov,auxvar,family,n,data,beta,p_gz,disp=NULL,ga.popsize,ga.propelit,ga.proptourney,ga.ngen,ga.mutrate,ga.initpop=NULL,optimMeasure,K.idx=NULL,seed=1,verbose=0,cluster=NULL){
  ptm0 <- Sys.time()
  if (is.character(family))
    family <- get(family, mode = "function", envir = parent.frame())
  if (is.function(family))
    family <- family()
  if (is.null(family$family)) {
    message(deparse(family))
    stop("'family' not recognized")
  }
  terms_dat<-unique( c(all.vars(formula), attr(terms(miscov), "term.labels") , attr(terms(auxvar), "term.labels")) )
  Y_<-as.character(formula[[2]]); G_<-all.vars(miscov); Z_<-all.vars(auxvar)

  ## remove entries with zero or negative probabilities
  p_gz <- p_gz[p_gz$q>0,]

  ### check if names in G_ and Z_ are in the joint ditribution data.frame p_gz
  g.vals<-unique(p_gz[,G_]);
  N <- NROW(data)

  if( is.null(disp) ){
    if( family$family %in% c("gaussian","Gamma") ){
      disp = var(data[,Y_])
    }else disp=1
  }
  ### Function that determines the optimality measure (D-, A- or parameter-specfic for now)
  ### if optimMeasure=="Par-spec" then K.idx MUST be not null
  if( optimMeasure=="Par-spec" & is.null(K.idx) ) stop("For a parameter-specific criterion K.idx must be provided.")

  ## assign these functions to the execution environment
  omega1_g=omega1_g; optMsuref=.optMsure(optimMeasure)#; class.ind2=class.ind2; class.ind2.Mat=class.ind2.Mat

  # envup=environment(omega1_g)=environment(class.ind2)=environment(class.ind2.Mat)=environment(optMsuref)=environment()
  environment(optMsuref) <- environment()

  uniqZ<-unique(p_gz[,Z_]); uniqZ<-uniqZ[order(uniqZ)]

  ### calculate FIMs for all samples (hopefully, a faster way than before)
  dat_ <- cbind(id_=rep(seq(1,N,by=1), each=NROW(g.vals)),data[rep(seq_len(N), each=NROW(g.vals)), ],matrix(rep(t(g.vals),N), ncol=NCOL(g.vals), byrow=TRUE))
  names(dat_)[(ncol(dat_)-NCOL(g.vals)+1):ncol(dat_)] <- G_
  dat_ <- droplevels(dat_)

  wgs <- omega1_g(formula,family,dat_,beta,p_gz,disp,G_,Z_)

  dat_$wg_ <- wgs$wg
  dat_ <- dat_[!is.na(dat_$wg_),]

  ## Baseline FIM: I3/N (empirical score outer product; valid large-N approximation
  ## to the full Louis FIM since I1 ≈ I2 for large N in the expansion context)
  obsIM_all  <- obsIMf1(formula,family,dat_,beta,p_gz,disp,G_,Z_,by.id = FALSE)
  IM1        <- obsIM_all$I3 / N

  ## Per-individual I2 and I3 for phase-2 improvement
  obsIM23_id <- obsIMf1(formula,family,dat_,beta,p_gz,disp,G_,Z_,by.id = TRUE)

  npars <- NCOL(IM1)

  IM2_id <- obsIM23_id$I2 / N
  IM3_id <- obsIM23_id$I3 / N

  Ind <- rep(FALSE,npars-1)
  Betas_ids <- which(grepl(paste0(paste0("^",G_),collapse="|"),colnames(model.matrix(formula,dat_[1,]))))
  Ind[Betas_ids] <- TRUE
  D <- rbind(diag(npars-1),0)

  if( all(beta[Ind]==0) ){
    Assess_fitness <- function(R){

      FIM <- IM1 + matrix(colSums(IM2_id*R),ncol=npars) - matrix(colSums(IM3_id*R),ncol=npars)
      F1 <-  t(D) %*% FIM %*% D

      invF1_noInd <- ginv(F1[!Ind,!Ind])
      V1 <- F1[Ind,Ind] - F1[Ind,!Ind] %*% invF1_noInd %*% F1[!Ind,Ind]
      invV1 <- solve( V1 )

      return( optMsuref(invV1,K.idx) )
    }


  } else {
    Assess_fitness <- function(R){

      FIM <- IM1 + matrix(colSums(IM2_id*R),ncol=npars) - matrix(colSums(IM3_id*R),ncol=npars)
      F1 <-  t(D) %*% FIM %*% D

      vcov = solve(F1)
      return(optMsuref(vcov,K.idx))

    }
  }

  #### R<-rep(0,N); R[sample(N,n)] <- 1
  R0s <- rep(0,N)
  ObjFun <- function(v){
    R <- R0s; R[v]<-1
    Assess_fitness(R)
  }

  time.elapsed0 <- as.numeric(difftime(Sys.time(), ptm0, units="secs"))

  ptm <- Sys.time()

  ## When cluster is non-NULL, kofnGA distributes fitness evaluations across
  ## workers each generation via parRapply.  Pre-load MASS on each worker so
  ## that ginv() is available inside the Assess_fitness closure.
  if (!is.null(cluster)) {
    parallel::clusterEvalQ(cluster, requireNamespace("MASS", quietly = TRUE))
  }

  out <- kofnGA::kofnGA(
    n           = N,
    k           = n,
    OF          = ObjFun,
    popsize     = ga.popsize,
    keepbest    = ceiling(ga.popsize * ga.propelit),
    tourneysize = max(ceiling(ga.popsize * ga.proptourney), 2),
    ngen        = ga.ngen,
    mutprob     = ga.mutrate,
    initpop     = ga.initpop,
    verbose     = verbose,
    cluster     = cluster
  )

  time.elapsed <- as.numeric(difftime(Sys.time(), ptm, units="secs"))
  out$time.elapsed <- c(preproc=time.elapsed0, GA=time.elapsed)

  return(out)
}





