#pragma once

#include <boost/math/distributions/fisher_f.hpp>
#include "modelbits.hpp"
#include "matrixw.hpp"
#include "randomeffects.hpp"
#include "matrixfield.h"

namespace glmmr {

enum class SE {
  GLS = 0,
    KR = 1,
    Robust = 2,
    BW = 3,
    KR2 = 4,
    Sat = 5,
    KRBoth = 6 // used for when two types of correction are required
};

enum class IM {
  EIM = 0,
    OIM = 1,
    GEE_IND = 2,
    EIM2 = 3
};

}

struct BoxResults {
  dblvec dof;
  dblvec scale;
  dblvec test_stat;
  dblvec p_value;
  BoxResults(const int r) : dof(r), scale(r), test_stat(r), p_value(r) {};
};

struct CorrectionDataBase {
public:
  MatrixXd vcov_beta;
  MatrixXd vcov_theta;
  VectorXd dof;
  VectorXd lambda;
  CorrectionDataBase(int n1, int m1, int n2, int m2): vcov_beta(n1,m1), vcov_theta(n2,m2), dof(n1), lambda(n1) {};
  CorrectionDataBase(const MatrixXd& vcov_beta_, const MatrixXd& vcov_theta_, const MatrixXd& dof_, const MatrixXd& lambda_) : 
    vcov_beta(vcov_beta_), vcov_theta(vcov_theta_), dof(dof_), lambda(lambda_)  {};
  CorrectionDataBase(const CorrectionDataBase& x) : vcov_beta(x.vcov_beta), vcov_theta(x.vcov_theta), dof(x.dof), lambda(x.lambda) {};
  CorrectionDataBase& operator=(const CorrectionDataBase& x) = default;
};

template<glmmr::SE corr>
struct CorrectionData : public CorrectionDataBase {
public:
  CorrectionData(int n1, int m1, int n2, int m2): CorrectionDataBase(n1,m1,n2,m2) {};
  CorrectionData(const CorrectionData& x) : CorrectionDataBase(x.vcov_beta, x.vcov_theta, x.dof, x.lambda) {};
  CorrectionData& operator=(const CorrectionData& x){
    CorrectionDataBase::operator=(x);
    return *this;
  };
};

template<>
struct CorrectionData<glmmr::SE::KRBoth> : public CorrectionDataBase {
public:
  MatrixXd vcov_beta_second;
  CorrectionData(int n1, int m1, int n2, int m2): CorrectionDataBase(n1,m1,n2,m2), vcov_beta_second(n1,m1) {};
  CorrectionData(const CorrectionData& x) : CorrectionDataBase(x.vcov_beta, x.vcov_theta, x.dof, x.lambda), vcov_beta_second(x.vcov_beta_second) {};
  CorrectionData& operator=(const CorrectionData& x){
    CorrectionDataBase::operator=(x);
    vcov_beta_second = x.vcov_beta_second;
    return *this;
  };
};

namespace glmmr {

using namespace Eigen;

template<typename modeltype>
class ModelMatrix{
public:
  modeltype&                        model;
  glmmr::MatrixW<modeltype>         W;
  glmmr::RandomEffects<modeltype>&  re;
  // constructors
  ModelMatrix(modeltype& model_, glmmr::RandomEffects<modeltype>& re_);
  ModelMatrix(const glmmr::ModelMatrix<modeltype>& matrix);
  ModelMatrix(modeltype& model_, glmmr::RandomEffects<modeltype>& re_, bool useBlock_, bool useSparse_);
  // functions
  MatrixXd                information_matrix();
  MatrixXd                ave_information_matrix();
  MatrixXd                Sigma(bool inverse = false);
  template<IM imtype>
  MatrixXd                observed_information_matrix();
  template<typename C = modeltype, 
           typename = std::enable_if_t<std::is_same_v<C, bits_hsgp> || 
             std::is_same_v<C, bits_spde>>>
  VectorXd                re_variance_obs();
  template<typename C = modeltype, 
           typename = std::enable_if_t<std::is_same_v<C, bits_spde>>>
  VectorXd                re_variance_at(const Eigen::SparseMatrix<double>& A_proj);
  MatrixXd                sandwich_matrix(); 
  std::vector<MatrixXd>   sigma_derivatives();
  template<IM imtype>
  MatrixXd                information_matrix_theta();
  template<SE corr, IM imtype>
  CorrectionData<corr>    small_sample_correction();
  MatrixXd                linpred();
  VectorMatrix            b_score();
  VectorMatrix            re_score();
  VectorXd                log_gradient(const VectorXd &v,bool beta = false);
  MatrixXd                gradient_eta(const MatrixXd &v);
  std::vector<glmmr::SigmaBlock> get_sigma_blocks();
  BoxResults              box();
  int                     P() const;
  int                     Q() const;
  MatrixXd                residuals(const int type, bool conditional = true);
  void                    posterior_u_samples(const int niter, const bool reml, const bool loglik = true, const bool append = false);                  
private:
  std::vector<glmmr::SigmaBlock>  sigma_blocks;
  void                            gen_sigma_blocks();
  MatrixXd                        sigma_block(int b, bool inverse = false);
  MatrixXd                        sigma_builder(int b, bool inverse = false);
  MatrixXd                        information_matrix_by_block(int b);
  bool                            useBlock = true;
  bool                            useSparse = true;
};

}

template<typename modeltype>
inline glmmr::ModelMatrix<modeltype>::ModelMatrix(modeltype& model_, glmmr::RandomEffects<modeltype>& re_): model(model_), W(model_), re(re_) { gen_sigma_blocks();};

template<typename modeltype>
inline glmmr::ModelMatrix<modeltype>::ModelMatrix(const glmmr::ModelMatrix<modeltype>& matrix) : model(matrix.model), W(matrix.W), re(matrix.re) { gen_sigma_blocks();};

template<typename modeltype>
inline glmmr::ModelMatrix<modeltype>::ModelMatrix(modeltype& model_, glmmr::RandomEffects<modeltype>& re_, bool useBlock_, bool useSparse_): model(model_), W(model_), re(re_) { 
  useBlock = useBlock_;
  useSparse = useSparse_;
  if(useBlock)gen_sigma_blocks();};

template<typename modeltype>
inline std::vector<glmmr::SigmaBlock> glmmr::ModelMatrix<modeltype>::get_sigma_blocks(){
  return sigma_blocks;
}

template<typename modeltype>
inline MatrixXd glmmr::ModelMatrix<modeltype>::information_matrix_by_block(int b){
  ArrayXi rows = Map<ArrayXi,Unaligned>(sigma_blocks[b].RowIndexes.data(),sigma_blocks[b].RowIndexes.size());
  MatrixXd X = glmmr::Eigen_ext::submat(model.linear_predictor.X(),rows,ArrayXi::LinSpaced(P(),0,P()-1));
  MatrixXd S = sigma_block(b,true);
  MatrixXd M = X.transpose()*S*X;
  return M;
}

template<>
inline MatrixXd glmmr::ModelMatrix<bits_hsgp>::information_matrix(){
  
  W.update(re.centred_u_mean());
  const int P_ = P();
  const int M = model.covariance.Q();
  
  MatrixXd X = model.linear_predictor.X();          // n × P
  MatrixXd A = model.covariance.ZPhi();              // n × M
  ArrayXd Lambda = model.covariance.LambdaSPD();     // M
  ArrayXd inv_lambda = 1.0 / Lambda;
  
  // W X and W A
  MatrixXd WX = X;
  MatrixXd WA = A;
  bool nonlinear_w = model.family.family != Fam::gaussian || (model.data.weights != 1).any();
  if(nonlinear_w){
    WX.applyOnTheLeft(W.W_.asDiagonal());
    WA.applyOnTheLeft(W.W_.asDiagonal());
  } else {
    WX /= model.data.var_par;
    WA /= model.data.var_par;
  }
  
  // X^T W X  (P × P)
  MatrixXd XtWX = X.transpose() * WX;
  
  // A^T W A + diag(1/Λ)  (M × M)
  MatrixXd AtWA = A.transpose() * WA;
  AtWA.diagonal() += inv_lambda.matrix();
  
  // X^T W A  (P × M)
  MatrixXd XtWA = X.transpose() * WA;
  
  // Woodbury: M_β = X^T W X - X^T W A (A^T W A + Λ^{-1})^{-1} A^T W X
  Eigen::LLT<MatrixXd> llt(AtWA);
  MatrixXd solve_XtWA = llt.solve(XtWA.transpose());  // M × P
  
  return XtWX - XtWA * solve_XtWA;
}

template<>
inline MatrixXd glmmr::ModelMatrix<bits_spde>::information_matrix(){
  auto& cov = model.covariance;
  if(!cov.spde_loaded){
    throw std::runtime_error("SPDE information_matrix: spde_data() not loaded.");
  }
  
  W.update(re.centred_u_mean());
  const int P_ = P();
  
  MatrixXd X = model.linear_predictor.X();           // n × P
  const SparseMatrix<double>& ZA = cov.ZA_;          // n × n_v, sparse
  
  // Weights vector (n)
  VectorXd W_diag;
  bool nonlinear_w = model.family.family != Fam::gaussian || (model.data.weights != 1).any();
  if(nonlinear_w){
    W_diag = W.W_;
  } else {
    W_diag = VectorXd::Constant(model.n(), 1.0 / model.data.var_par);
  }
  
  // Weighted X  (n × P), dense
  MatrixXd WX = W_diag.asDiagonal() * X;
  
  // X^T W X  (P × P)
  MatrixXd XtWX = X.transpose() * WX;
  
  // (ZA)^T W X  (n_v × P) — sparse × dense
  MatrixXd ZAtWX = ZA.transpose() * WX;
  
  // Posterior precision P_post = (ZA)^T W (ZA) + Q, factorised via refactor_P
  cov.refactor_P(W_diag);
  
  // solve for (P_post)^{-1} (ZA)^T W X    (n_v × P)
  MatrixXd solve_ZAtWX = cov.chol_P.solve(ZAtWX);
  if(cov.chol_P.info() != Success){
    throw std::runtime_error("SPDE information_matrix: P solve failed.");
  }
  
  // M_β = X^T W X - (X^T W ZA) * solve_ZAtWX
  //     = XtWX - ZAtWX^T * solve_ZAtWX
  return XtWX - ZAtWX.transpose() * solve_ZAtWX;
}

template<typename modeltype>
inline MatrixXd glmmr::ModelMatrix<modeltype>::information_matrix(){
  W.update(re.centred_u_mean());
  MatrixXd M = MatrixXd::Zero(P(),P());
  for(int i = 0; i< sigma_blocks.size(); i++){
    M += information_matrix_by_block(i);
  }
  return M;
}

template<>
inline MatrixXd glmmr::ModelMatrix<bits_spde>::ave_information_matrix(){
  auto& cov = model.covariance;
  if(!cov.spde_loaded){
    throw std::runtime_error("SPDE ave_information_matrix: spde_data() not loaded.");
  }
  
  W.update(re.centred_u_mean());
  const int P_ = P();
  const int n  = model.n();
  const int K  = re.u_.cols();
  MatrixXd u_samples = cov.ZLu(re.u_);
  
  if(K < 2) throw std::runtime_error(
      "louis_information_beta requires at least 2 MC samples in re.u_; "
      "for Laplace-fit models, draw a posterior sample first via "
      "posterior_u_samples().");
  
  ArrayXd w = re.u_weight_;
  double wsum = w.sum();
  if(wsum <= 0 || !std::isfinite(wsum)){
    w = ArrayXd::Constant(K, 1.0/K);
  } else {
    w /= wsum;
  }
  
  const MatrixXd X = model.linear_predictor.X();
  ArrayXd xb = model.linear_predictor.xb().array() + model.data.offset.array();
  MatrixXd s_mat(n, K);
  VectorXd s_bar = VectorXd::Zero(n);
  VectorXd W_bar = VectorXd::Zero(n);
  ArrayXd eta_k(n), mu_k(n), W_k(n), s_k(n);
  
  if(model.family.family == Fam::poisson){
    ArrayXd sigma2_i = cov.diag_ZA_Qinv_ZAt_hutch();
    xb -= 0.5 * sigma2_i;
  }
  
  for(int k = 0; k < K; ++k){
    eta_k = xb + u_samples.col(k).array() - u_samples.col(k).mean();
    
    switch(model.family.family){
    case Fam::binomial: case Fam::bernoulli: {
      if(model.family.link == Link::logit){
      mu_k = 1.0 / (1.0 + (-eta_k).exp());
    } else {
      mu_k = glmmr::maths::mod_inv_func(eta_k.matrix(), model.family.link).array();
    }
    s_k = model.data.variance * mu_k;
    W_k = model.data.variance * mu_k * (1.0 - mu_k);
    break;
    }
    case Fam::poisson: {
      mu_k = eta_k.exp();
      s_k = mu_k;
      W_k = mu_k;
      break;
    }
    case Fam::gaussian: {
      mu_k = eta_k;
      s_k = mu_k;
      W_k = (model.data.variance.inverse() * model.data.weights);
      break;
    }
    case Fam::gamma: case Fam::beta: case Fam::exponential: {
      mu_k = glmmr::maths::mod_inv_func(eta_k.matrix(), model.family.link).array();
      s_k = mu_k;
      W_k = W.W_.array();
      break;
    }
    default:
      throw std::runtime_error("ave_information_matrix (SPDE): unsupported family");
    }
    
    s_mat.col(k) = s_k.matrix();
    s_bar += w(k) * s_k.matrix();
    W_bar += w(k) * W_k.matrix();
  }
  
  MatrixXd bread = X.transpose() * W_bar.asDiagonal() * X;
  MatrixXd meat  = MatrixXd::Zero(P_, P_);
  VectorXd g_k(P_);
  for(int k = 0; k < K; ++k){
    g_k = X.transpose() * (s_mat.col(k) - s_bar);
    meat.noalias() += w(k) * g_k * g_k.transpose();
  }
  
  return bread - meat;
}

template<typename modeltype>
inline MatrixXd glmmr::ModelMatrix<modeltype>::ave_information_matrix(){
  
  W.update(re.centred_u_mean());
  const int P_ = P();
  const int n  = model.n();
  const int K  = re.u_.cols();
  MatrixXd u_samples = model.covariance.ZLu(re.u_);
  
  if(K < 2) throw std::runtime_error(
      "louis_information_beta requires at least 2 MC samples in re.u_; "
      "for Laplace-fit models, draw a posterior sample first via "
      "posterior_u_samples().");
  
  // Importance weights, normalised
  ArrayXd w = re.u_weight_;
  double wsum = w.sum();
  if(wsum <= 0 || !std::isfinite(wsum)){
    w = ArrayXd::Constant(K, 1.0/K);
  } else {
    w /= wsum;
  }
  
  const MatrixXd X = model.linear_predictor.X();
  ArrayXd xb = model.linear_predictor.xb().array() + model.data.offset.array();
  // Per-sample s_k (n-vector) and weighted accumulators for s_bar and W_bar
  MatrixXd s_mat(n, K);          // s(u^k) per sample
  VectorXd s_bar = VectorXd::Zero(n);
  VectorXd W_bar = VectorXd::Zero(n);
  
  // re.zu_ has columns Z*L*u^k (for HSGP, A * u^k); same as in linpred()
  ArrayXd eta_k(n), mu_k(n), W_k(n), s_k(n);
  
  if(model.family.family == Fam::poisson){
    MatrixXd D = model.covariance.D(false, false);
    MatrixXd Z(model.n(), model.covariance.Q());
    if constexpr (std::is_same_v<modeltype, bits_hsgp>){
      Z = model.covariance.ZPhi();
    } else {
      Z = model.covariance.Z();
    }
    ArrayXd sigma2_i(model.n());
    MatrixXd ZD = Z * D;  // n x Q
    for(int i = 0; i < model.n(); i++){
      sigma2_i(i) = ZD.row(i).dot(Z.row(i));
    }
    xb -= 0.5 * sigma2_i;
  }
  
  for(int k = 0; k < K; ++k){
    eta_k = xb + u_samples.col(k).array() - u_samples.col(k).mean();
    
    switch(model.family.family){
    case Fam::binomial: case Fam::bernoulli: {
      if(model.family.link == Link::logit){
      mu_k = 1.0 / (1.0 + (-eta_k).exp());
    } else {
      // probit or other: route through generic helper
      mu_k = glmmr::maths::mod_inv_func(eta_k.matrix(), model.family.link).array();
    }
    s_k = model.data.variance * mu_k;                 // m_i * mu_i
    W_k = model.data.variance * mu_k * (1.0 - mu_k);  // m_i * mu_i (1-mu_i)
    break;
    }
    case Fam::poisson: {
      mu_k = eta_k.exp();
      s_k = mu_k;
      W_k = mu_k;
      break;
    }
    case Fam::gaussian: {
      mu_k = eta_k;                                     // identity link
      s_k = mu_k;
      W_k = (model.data.variance.inverse() * model.data.weights);
      break;
    }
    case Fam::gamma: case Fam::beta: case Fam::exponential: {
      mu_k = glmmr::maths::mod_inv_func(eta_k.matrix(), model.family.link).array();
      s_k = mu_k;
      W_k = W.W_.array();
      break;
    }
    default:
      throw std::runtime_error("louis_information_beta: unsupported family");
    }
    
    s_mat.col(k) = s_k.matrix();
    s_bar += w(k) * s_k.matrix();
    W_bar += w(k) * W_k.matrix();
  }
  // Bread: X' diag(W_bar) X
  MatrixXd bread = X.transpose() * W_bar.asDiagonal() * X;
  
  // Meat: sum_k w_k g_k g_k', g_k = X' (s_k - s_bar)
  MatrixXd meat = MatrixXd::Zero(P_, P_);
  VectorXd g_k(P_);
  for(int k = 0; k < K; ++k){
    g_k = X.transpose() * (s_mat.col(k) - s_bar);
    meat.noalias() += w(k) * g_k * g_k.transpose();
  }
  
  return bread - meat;
}



template<typename modeltype>
inline void glmmr::ModelMatrix<modeltype>::gen_sigma_blocks(){
  int block_counter = 0;
  intvec2d block_ids(model.n());
  int block_size;
  SparseMatrix<double> Z = model.covariance.Z_sparse();
  int i,j,k;
  for(int b = 0; b < model.covariance.B(); b++){
    block_size = model.covariance.block_dim(b);
    for(i = 0; i < block_size; i++){
      for(j = 0; j < model.n(); j++){
        if(Z.coeff(j,i+block_counter)!=0){
          block_ids[j].push_back(b);
        }
      }
    }
    block_counter += block_size;
  }
  intvec idx_matches;
  int n_matches;
  for(i = 0; i < model.n(); i++){
    if(sigma_blocks.size() == 0){
      glmmr::SigmaBlock newblock(block_ids[i]);
      newblock.add_row(i);
      sigma_blocks.push_back(newblock);
    } else {
      for(j = 0; j < sigma_blocks.size(); j++){
        if(sigma_blocks[j] == block_ids[i]){
          idx_matches.push_back(j);
        }
      }
      n_matches = idx_matches.size();
      if(n_matches==0){
        glmmr::SigmaBlock newblock(block_ids[i]);
        newblock.add_row(i);
        sigma_blocks.push_back(newblock);
      } else if(n_matches==1){
        sigma_blocks[idx_matches[0]].add(block_ids[i]);
        sigma_blocks[idx_matches[0]].add_row(i);
      } else if(n_matches>1){
        std::reverse(idx_matches.begin(),idx_matches.end());
        for(k = 0; k < (n_matches-1); k++){
          sigma_blocks[idx_matches[n_matches-1]].merge(sigma_blocks[idx_matches[k]]);
          sigma_blocks[idx_matches[n_matches-1]].add_row(i);
          sigma_blocks.erase(sigma_blocks.begin()+idx_matches[k]);
        }
      }
    }
    idx_matches.clear();
  }
}
template<typename modeltype>
inline int glmmr::ModelMatrix<modeltype>::P() const {
  return model.linear_predictor.P();
}


template<typename modeltype>
inline int glmmr::ModelMatrix<modeltype>::Q() const {
  return model.covariance.Q();
}

template<typename modeltype>
inline MatrixXd glmmr::ModelMatrix<modeltype>::Sigma(bool inverse){
  W.update(re.centred_u_mean());
  MatrixXd S(model.n(), model.n());
  if(useBlock){
    S = sigma_builder(0,inverse);
  } else {
    MatrixXd ZL = model.covariance.ZL();
    S = ZL * ZL.transpose();
    S += W.W().array().inverse().matrix().asDiagonal();
    if(inverse){
      S = S.llt().solve(MatrixXd::Identity(S.rows(),S.cols()));
    }
  }
  return S;
}

template<typename modeltype>
inline MatrixXd glmmr::ModelMatrix<modeltype>::sigma_block(int b,
                                                           bool inverse){
  // UPDATE THIS TO NOT USE SPARSE IF DESIRED
  MatrixXd ZLs = model.covariance.ZL();
  ArrayXi rows = Map<ArrayXi,Unaligned>(sigma_blocks[b].RowIndexes.data(),sigma_blocks[b].RowIndexes.size());
  MatrixXd ZL = glmmr::Eigen_ext::submat(ZLs,rows,ArrayXi::LinSpaced(Q(),0,Q()-1));//sparse_to_dense(ZLs,false);
  MatrixXd S = ZL * ZL.transpose();
  for(int i = 0; i < S.rows(); i++){
    S(i,i)+= 1/W.W()(sigma_blocks[b].RowIndexes[i]);
  }
  if(inverse){
    S = S.llt().solve(MatrixXd::Identity(S.rows(),S.cols()));
  }
  return S;
}

template<typename modeltype>
inline MatrixXd glmmr::ModelMatrix<modeltype>::sigma_builder(int b,
                                                             bool inverse){
  int B_ = sigma_blocks.size();
  if (b == B_ - 1) {
    return sigma_block(b,inverse);
  }
  else {
    MatrixXd mat1 = sigma_block(b,inverse);
    MatrixXd mat2;
    if (b == B_ - 2) {
      mat2 = sigma_block(b+1,inverse);
    }
    else {
      mat2 = sigma_builder(b + 1,  inverse);
    }
    int n1 = mat1.rows();
    int n2 = mat2.rows();
    MatrixXd dmat = MatrixXd::Zero(n1+n2, n1+n2);
    dmat.block(0,0,n1,n1) = mat1;
    dmat.block(n1, n1, n2, n2) = mat2;
    return dmat;
  }
}

template<typename modeltype>
template<glmmr::IM imtype>
inline MatrixXd glmmr::ModelMatrix<modeltype>::observed_information_matrix(){
  MatrixXd X = model.linear_predictor.X();
  bool nonlinear_w = model.family.family != Fam::gaussian || (model.data.weights != 1).any();
  if constexpr (imtype == IM::EIM){
    W.update(re.centred_u_mean());
    MatrixXd WX = X;
    if (nonlinear_w) WX.applyOnTheLeft(W.W_.asDiagonal());
    MatrixXd XtXW = X.transpose() * WX;
    MatrixXd ZL = model.covariance.ZL();
    MatrixXd WZL = ZL;
    if (nonlinear_w) WZL.applyOnTheLeft(W.W_.asDiagonal());
    MatrixXd XtWZL = X.transpose() * WZL;
    MatrixXd ZLWLZ = ZL.transpose() * WZL;
    if (!nonlinear_w) {
        XtXW *= (1.0 / model.data.var_par);
        XtWZL *= (1.0 / model.data.var_par);
        ZLWLZ *= (1.0 / model.data.var_par);
    }
    ZLWLZ += MatrixXd::Identity(Q(),Q());
    MatrixXd infomat(P()+Q(),P()+Q());
    infomat.topLeftCorner(P(),P()) = XtXW;
    infomat.topRightCorner(P(),Q()) = XtWZL;
    infomat.bottomLeftCorner(Q(),P()) = XtWZL.transpose();
    infomat.bottomRightCorner(Q(),Q()) = ZLWLZ;
    return infomat;
  } else if constexpr (imtype == IM::OIM){
    MatrixXd M = information_matrix();
    MatrixXd Mt = information_matrix_theta<IM::OIM>();
    int npar = Mt.cols();
    MatrixXd infomat(P()+npar,P()+npar);
    infomat.topLeftCorner(P(),P()) = -1.0*M;
    infomat.topRightCorner(P(),npar) = -1.0*Mt.bottomRows(P());
    infomat.bottomLeftCorner(npar,P()) = -1.0*Mt.bottomRows(P()).transpose();
    infomat.bottomRightCorner(npar,npar) = Mt.topRows(npar);
    return -1.0*infomat;
  } else if constexpr (imtype == IM::EIM2) {
    W.update(re.centred_u_mean());
    MatrixXd WX = X;
    if (nonlinear_w) WX.applyOnTheLeft(W.W_.asDiagonal());
    
    // Z matrix: Phi for HSGP, Z for standard
    MatrixXd Z;
    MatrixXd Dinv;
    
    if constexpr (std::is_same_v<modeltype, bits_hsgp>){
      Z = model.covariance.ZPhi();
      Dinv = model.covariance.LambdaSPD().inverse().matrix().asDiagonal();  // diag, cheap
    } else if constexpr (std::is_same_v<modeltype, bits_spde>){
      Z = MatrixXd(model.covariance.ZA_);
      Dinv = MatrixXd(model.covariance.Q_mat);       // Q_mat *is* D⁻¹
    } else {
      Z = model.covariance.Z();
      MatrixXd D = model.covariance.D(false, false);
      if (model.covariance.all_group_re()){
        Dinv = D;
        for (int i = 0; i < D.rows(); ++i) Dinv(i, i) = 1.0 / D(i, i);
      } else {
        Dinv = D.llt().solve(MatrixXd::Identity(D.rows(), D.cols()));
      }
    }
    
    if (nonlinear_w) WX.applyOnTheLeft(W.W_.asDiagonal());
    MatrixXd WZ = Z;
    if (nonlinear_w) WZ.applyOnTheLeft(W.W_.asDiagonal());
    if (!nonlinear_w){
      WX *= (1.0 / model.data.var_par);
      WZ *= (1.0 / model.data.var_par);
    }
    
    const int Qdim = Z.cols();
    MatrixXd infomat(P() + Qdim, P() + Qdim);
    infomat.topLeftCorner(P(), P())     = X.transpose() * WX;
    infomat.topRightCorner(P(), Qdim)   = X.transpose() * WZ;
    infomat.bottomLeftCorner(Qdim, P()) = infomat.topRightCorner(P(), Qdim).transpose();
    infomat.bottomRightCorner(Qdim, Qdim) = Z.transpose() * WZ + Dinv;
    return infomat;
  } else {
    MatrixXd M = information_matrix();
    M = M.llt().solve(MatrixXd::Identity(M.rows(),M.cols()));
    MatrixXd XtXW = X.transpose() * W.W_.asDiagonal() * X;
    M = XtXW * M * XtXW;
    return M;
  }
}

template<typename modeltype>
template<typename C, typename>
inline VectorXd glmmr::ModelMatrix<modeltype>::re_variance_at(
    const Eigen::SparseMatrix<double>& A_proj)
{
  using Eigen::SparseMatrix; 
  
  W.update(re.centred_u_mean());
  MatrixXd X = model.linear_predictor.X();
  const bool nonlinear_w = model.family.family != Fam::gaussian
  || (model.data.weights != 1).any();
  
  const int n = model.n();
  VectorXd W_(n);
  ArrayXd xb = model.linear_predictor.xb().array() + model.data.offset.array();
  ArrayXd eta = xb + (A_proj * re.u_mean_).array();
  switch(model.family.family){
  case Fam::binomial: case Fam::bernoulli: {
    ArrayXd logitp = 1.0 / (1.0 + (-eta).exp());
    W_ = (model.data.variance * logitp * (1.0 - logitp)).matrix();
    break;
  }
  case Fam::poisson: {
    W_ = eta.exp().matrix();
    break;
  }
  case Fam::gaussian: {
    W_ = (model.data.variance.inverse() * model.data.weights).matrix();
    break;
  }
  default:
    throw std::runtime_error("SPDE posterior only for Gaussian, Poisson, Binomial");
  }
  
  MatrixXd WX = X;
  if(nonlinear_w) WX.applyOnTheLeft(W_.asDiagonal());
  else            WX *= (1.0 / model.data.var_par);
  MatrixXd XtWX = X.transpose() * WX;
  
  auto& cov = model.covariance;
  cov.refactor_P(W_);
  
  MatrixXd B     = cov.ZA_.transpose() * WX;        // n_v × P  (model-side B)
  MatrixXd PinvB = cov.chol_P.solve(B);
  MatrixXd Schur = XtWX - B.transpose() * PinvB;
  Eigen::LLT<MatrixXd> llt_Schur(Schur);
  
  // ── Substitute A_proj for ZA in the diagonal extraction ──
  MatrixXd Pinv_AT      = cov.chol_P.solve(MatrixXd(A_proj.transpose())); // n_v × n_pred
  MatrixXd Bt_Pinv_AT   = B.transpose() * Pinv_AT;                       // P × n_pred
  MatrixXd Schur_solve  = llt_Schur.solve(Bt_Pinv_AT);                   // P × n_pred
  
  // VectorXd term1 = (A_proj * Pinv_AT).diagonal();   // sparse × dense → fine
  // // or, avoiding the off-diagonals:
  VectorXd term1 = VectorXd::Zero(A_proj.rows());
  for (int k = 0; k < A_proj.outerSize(); ++k)
    for (SparseMatrix<double>::InnerIterator it(A_proj, k); it; ++it)
      term1(it.row()) += it.value() * Pinv_AT(it.col(), it.row());
  VectorXd term2 = (Bt_Pinv_AT.array() * Schur_solve.array()).colwise().sum();
  return term1 + term2;
}

template<typename modeltype>
template<typename C, typename>
inline VectorXd glmmr::ModelMatrix<modeltype>::re_variance_obs()
{
  using Eigen::SparseMatrix;
  
  W.update(re.centred_u_mean());
  const int nobs = model.n();
  const int Pdim = P();
  
  MatrixXd X = model.linear_predictor.X();
  const bool nonlinear_w = model.family.family != Fam::gaussian 
  || (model.data.weights != 1).any();
  
  const int n = model.n();
  VectorXd W_(n);
  ArrayXd eta(n);
  ArrayXd xb = model.linear_predictor.xb().array() + model.data.offset.array();
  if constexpr (std::is_same_v<C, bits_hsgp>){
    eta = xb + (model.covariance.ZPhi() * re.u_mean_).array();
  }  else if constexpr (std::is_same_v<C, bits_spde>) {
    eta = xb + (model.covariance.ZA_ * re.u_mean_).array();
  }
  
  switch(model.family.family){
  case Fam::binomial: case Fam::bernoulli: {
    ArrayXd logitp = 1.0 / (1.0 + (-eta).exp());
    W_ = (model.data.variance * logitp * (1.0 - logitp)).matrix();
    break;
  }
  case Fam::poisson: {
    W_ = eta.exp().matrix();
    break;
  }
  case Fam::gaussian: {
    W_ = (model.data.variance.inverse() * model.data.weights).matrix();
    break;
  }
  default:
    throw std::runtime_error("SPDE posterior only for Gaussian, Poisson, Binomial");
  }
  
  // --- Build WX (P × P pieces are branch-independent) ---
  MatrixXd WX = X;
  if(nonlinear_w){
    WX.applyOnTheLeft(W_.asDiagonal());
  } else {
    WX *= (1.0 / model.data.var_par);
  }
  MatrixXd XtWX = X.transpose() * WX;   // P × P
  Eigen::LLT<MatrixXd> llt_XtWX(XtWX);
  
  VectorXd var_obs(nobs);
  
  if constexpr (std::is_same_v<C, bits_hsgp>){
    MatrixXd Z = model.covariance.ZPhi();            // n × M
    const int Mdim = Z.cols();
    Eigen::ArrayXd Lambda = model.covariance.LambdaSPD();
    
    MatrixXd WZ = Z;
    if(nonlinear_w){
      WZ.applyOnTheLeft(W_.asDiagonal());
    } else {
      WZ *= (1.0 / model.data.var_par);
    }
    
    MatrixXd A_mat = Z.transpose() * WZ;
    A_mat.diagonal().array() += 1.0 / Lambda;
    Eigen::LLT<MatrixXd> llt_A(A_mat);
    
    MatrixXd B = Z.transpose() * WX;                 // M × P
    MatrixXd AinvB = llt_A.solve(B);                 // M × P
    MatrixXd Schur = XtWX - B.transpose() * AinvB;   // P × P
    Eigen::LLT<MatrixXd> llt_Schur(Schur);
    MatrixXd Ainv_Zt = llt_A.solve(Z.transpose());    // M × n
    MatrixXd Bt_Ainv_Zt = B.transpose() * Ainv_Zt;    // P × n
    MatrixXd Schur_solve = llt_Schur.solve(Bt_Ainv_Zt); // P × n
    for(int i = 0; i < nobs; ++i){
      double term1 = Z.row(i).dot(Ainv_Zt.col(i));
      double term2 = Bt_Ainv_Zt.col(i).dot(Schur_solve.col(i));
      var_obs(i) = term1 + term2;
    }
    
  } else if constexpr (std::is_same_v<C, bits_spde>){
    auto& cov = model.covariance;
    const SparseMatrix<double>& ZA = cov.ZA_;         // n × n_v, sparse
    cov.refactor_P(W_);
    MatrixXd B = ZA.transpose() * WX;                 // n_v × P
    MatrixXd PinvB = cov.chol_P.solve(B);             // n_v × P  (P sparse solves)
    MatrixXd Schur = XtWX - B.transpose() * PinvB;    // P × P
    Eigen::LLT<MatrixXd> llt_Schur(Schur);
    MatrixXd Pinv_ZAt = cov.chol_P.solve(MatrixXd(ZA.transpose()));  // n_v × n
    MatrixXd Bt_Pinv_ZAt = B.transpose() * Pinv_ZAt;                 // P × n
    MatrixXd Schur_solve = llt_Schur.solve(Bt_Pinv_ZAt);             // P × n
    VectorXd term1_vec = VectorXd::Zero(nobs);
    {
      for(int j = 0; j < ZA.outerSize(); ++j){
        for(SparseMatrix<double>::InnerIterator it(ZA, j); it; ++it){
          term1_vec(it.row()) += it.value() * Pinv_ZAt(j, it.row());
        }
      }
    }
    VectorXd term2_vec = (Bt_Pinv_ZAt.array() * Schur_solve.array()).colwise().sum();
    var_obs = term1_vec + term2_vec;
  }
  
  return var_obs;
}

template<typename modeltype>
inline MatrixXd glmmr::ModelMatrix<modeltype>::sandwich_matrix(){
  
  MatrixXd infomat = information_matrix();
  infomat = infomat.llt().solve(MatrixXd::Identity(P(),P()));
  MatrixXd X = model.linear_predictor.X();
  MatrixXd S = Sigma(true);
  MatrixXd SX = S*X;
  // MatrixXd resid_sum = MatrixXd::Zero(X.rows(),X.rows());
  // MatrixXd zd = linpred();
  VectorXd resid = model.linear_predictor.xb()+model.data.offset;
  resid = model.data.y - glmmr::maths::mod_inv_func(resid, model.family.link);
  MatrixXd resid_sum = resid * resid.transpose();
  
#if defined(R_BUILD) && defined(ENABLE_DEBUG)
  Rcpp::Rcout << "\nSANDWICH\n";
  int size = resid_sum.rows() < 10 ? resid_sum.rows() : 10;
  Rcpp::Rcout << "\nResidual cross prod: \n" << resid_sum.block(0,0,size,size).transpose();
#endif
  
  MatrixXd robust = infomat * SX.transpose() * resid_sum * SX * infomat;//(infomat.rows(),infomat.cols());
  
  return robust;
}

template<typename modeltype>
inline std::vector<MatrixXd> glmmr::ModelMatrix<modeltype>::sigma_derivatives()
{
  std::vector<MatrixXd> derivs;
  model.covariance.derivatives(derivs,2);
  return derivs;
}

template<typename modeltype>
template<glmmr::IM imtype>
inline MatrixXd glmmr::ModelMatrix<modeltype>::information_matrix_theta()
{
  MatrixXd Mc = model.covariance.information_matrix();
  if (imtype == IM::EIM && Mc(0,0)!=0){
    return Mc;
  } else {
    int n = model.n();
    std::vector<MatrixXd> derivs;
    model.covariance.derivatives(derivs,1);
    int R = model.covariance.npar();
    int Rmod = model.family.family==Fam::gaussian ? R+1 : R;
    MatrixXd SigmaInv = Sigma(true);
    MatrixXd Z = model.covariance.Z();
    MatrixXd M_theta = MatrixXd::Zero(Rmod,Rmod);
    glmmr::MatrixField<MatrixXd> S;
    VectorXd resid(1);
    
    // residuals in case of OIM
    if constexpr (imtype == IM::OIM){
      if(model.data.y.size() != model.n()) throw std::runtime_error("y data not correct size");
      resid.resize(model.n());
      resid = model.linear_predictor.xb()+model.data.offset;
      resid = model.data.y - glmmr::maths::mod_inv_func(resid, model.family.link);
    }
    
    for(int i = 0; i < Rmod; i++){
      MatrixXd partial0(model.n(),model.n());
      if(i < R){
        partial0 = Z*derivs[1+i]*Z.transpose();
      } else {
        partial0 = MatrixXd::Identity(n,n);
        if((model.data.weights != 1).any())partial0 = model.data.weights.inverse().matrix().asDiagonal();
      }
      S.add(partial0);
    }
    
    if(useBlock){
      for(int b = 0; b< sigma_blocks.size(); b++){
        ArrayXi rows = Map<ArrayXi,Unaligned>(sigma_blocks[b].RowIndexes.data(),sigma_blocks[b].RowIndexes.size());
        MatrixXd SigmaInvsub = glmmr::Eigen_ext::submat(SigmaInv,rows,rows);
        VectorXd residsub(1);
        if constexpr (imtype == IM::OIM){
          residsub.resize(rows.size());
          for(int r = 0; r < rows.size(); r++) residsub(r) = resid(rows(r));
        }
        for(int i = 0; i < Rmod; i++){
          MatrixXd Ssub1 = glmmr::Eigen_ext::submat(S(i),rows,rows);
          MatrixXd SPS = SigmaInvsub * Ssub1 * SigmaInvsub;
          for(int j = i; j < Rmod; j++){
            MatrixXd Ssub2 = glmmr::Eigen_ext::submat(S(j),rows,rows);
            double oim_adj;
            if constexpr (imtype == IM::OIM){
              oim_adj = (residsub.transpose() * SPS * Ssub2 * SigmaInvsub * residsub)(0);
            }
            for(int k = 0; k < rows.size(); k++){
              for(int l = 0; l < rows.size(); l++){
                M_theta(i,j) += 0.5*SPS(k,l)*Ssub2(l,k);
                if constexpr (imtype == IM::OIM) M_theta(i,j) -= oim_adj;
              }
            }
            if(i!=j)M_theta(j,i)=M_theta(i,j);
          }
        }
      }
    } else {
      for(int i = 0; i < Rmod; i++){
        MatrixXd SPS = SigmaInv * S(i) * SigmaInv;
        for(int j = i; j < Rmod; j++){
          MatrixXd Ssub2 = S(j);
          double oim_adj;
          if constexpr (imtype == IM::OIM){
            oim_adj = (resid.transpose() * SPS * Ssub2 * SigmaInv * resid)(0);
          }
          // M_theta(i,j) = 0.5*(SPS*Ssub2).trace();
          // if constexpr (imtype == IM::OIM) M_theta(i,j) -= oim_adj;
          for(int k = 0; k < model.n(); k++){
            for(int l = 0; l <  model.n(); l++){
              M_theta(i,j) += 0.5*SPS(k,l)*Ssub2(l,k);
              if constexpr (imtype == IM::OIM) M_theta(i,j) -= oim_adj;
            }
          }
          if(i!=j)M_theta(j,i)=M_theta(i,j);
        }
      }
    }
    
    if constexpr (imtype == IM::OIM){
      // add in the beta-theta part of the matrix
      MatrixXd X = model.linear_predictor.X();
      MatrixXd Mbt(X.cols(),M_theta.cols());
      if(useBlock){
        for(int b = 0; b< sigma_blocks.size(); b++){
          Mbt.setZero();
          ArrayXi rows = Map<ArrayXi,Unaligned>(sigma_blocks[b].RowIndexes.data(),sigma_blocks[b].RowIndexes.size());
          MatrixXd Xsub = glmmr::Eigen_ext::submat(X,rows,ArrayXi::LinSpaced(P(),0,P()-1));
          MatrixXd SigmaInvsub = glmmr::Eigen_ext::submat(SigmaInv,rows,rows);
          VectorXd residsub(rows.size());
          for(int r = 0; r < rows.size(); r++) residsub(r) = resid(rows(r));
          for(int i =0; i < M_theta.cols(); i++){
            MatrixXd Ssub1 = glmmr::Eigen_ext::submat(S(i),rows,rows);
            Mbt.col(i) += Xsub.transpose() * SigmaInvsub * Ssub1 * SigmaInvsub * residsub;
          }
        }
      } else {
        for(int i =0; i < M_theta.cols(); i++) Mbt.col(i) = X.transpose() * SigmaInv * S(i) * SigmaInv * resid;
      }
      M_theta.conservativeResize(M_theta.cols()+X.cols(),M_theta.cols());
      M_theta.bottomRows(X.cols()) = Mbt;
    }
    return M_theta;
  }
}

template<typename modeltype>
template<glmmr::SE corr, glmmr::IM imtype>
inline CorrectionData<corr> glmmr::ModelMatrix<modeltype>::small_sample_correction()
{
  using namespace glmmr;
  static_assert(corr == SE::KR || corr == SE::KR2 || corr == SE::Sat || corr == SE::KRBoth,"Only Kenward-Roger or Satterthwaite allowed for small sample correction");
  
  int n = model.n();
  std::vector<MatrixXd> derivs;
  model.covariance.derivatives(derivs,2);
  int R = model.covariance.npar();
  int Rmod = model.family.family==Fam::gaussian ? R+1 : R;
  
  MatrixXd M = information_matrix();
  M = M.llt().solve(MatrixXd::Identity(M.rows(),M.cols()));
  MatrixXd M_new(M);
  MatrixXd SigmaInv = Sigma(true);
  MatrixXd Z = model.covariance.Z();
  MatrixXd X = model.linear_predictor.X();
  MatrixXd M_theta = MatrixXd::Zero(Rmod,Rmod);
  MatrixXd meat = MatrixXd::Zero(X.cols(),X.cols());
  MatrixField<MatrixXd> Pmat;
  MatrixField<MatrixXd> Q;
  MatrixField<MatrixXd> RR;
  MatrixField<MatrixXd> S;
  // in case of OIM
  MatrixXd M_theta_oim = MatrixXd::Zero(Rmod,Rmod);
  MatrixXd M_theta_kr = MatrixXd::Zero(Rmod,Rmod);
  VectorXd resid(1);
  int counter, counter_rr;
  
  // residuals in case of OIM
  if constexpr (imtype == IM::OIM){
    if(model.data.y.size() != model.n()) throw std::runtime_error("y data not correct size");
    resid.resize(model.n());
    resid = model.linear_predictor.xb()+model.data.offset;
    resid = model.data.y - glmmr::maths::mod_inv_func(resid, model.family.link);
  }
  
  for(int i = 0; i < Rmod; i++){
    MatrixXd partial0(model.n(),model.n());
    if(i < R){
      partial0 = Z*derivs[1+i]*Z.transpose();
    } else {
      partial0 = MatrixXd::Identity(n,n);
      if((model.data.weights != 1).any())partial0 = model.data.weights.inverse().matrix().asDiagonal();
    }
    S.add(partial0);
  }
  
  for(int i = 0; i < Rmod; i++){
    for(int j = i; j < Rmod; j++){
      if(i < R && j < R){
        int scnd_idx = i + j*(R-1) - j*(j-1)/2;
        S.add(Z*derivs[R+1+scnd_idx]*Z.transpose());
      }
    }
  }
  
  if(useBlock){
    for(int b = 0; b< sigma_blocks.size(); b++){
      ArrayXi rows = Map<ArrayXi,Unaligned>(sigma_blocks[b].RowIndexes.data(),sigma_blocks[b].RowIndexes.size());
      MatrixXd SigmaInvsub = glmmr::Eigen_ext::submat(SigmaInv,rows,rows);
      MatrixXd Xsub = glmmr::Eigen_ext::submat(X,rows,ArrayXi::LinSpaced(P(),0,P()-1));
      VectorXd residsub(1);
      if constexpr (imtype == IM::OIM){
        residsub.resize(rows.size());
        for(int r = 0; r < rows.size(); r++) residsub(r) = resid(rows(r));
      }
      counter = 0;
      counter_rr = 0;
      for(int i = 0; i < Rmod; i++){
        MatrixXd Ssub1 = glmmr::Eigen_ext::submat(S(i),rows,rows);
        MatrixXd SPS = SigmaInvsub * Ssub1 * SigmaInvsub;
        // create P matrix
        if(b==0){
          Pmat.add(MatrixXd::Zero(P(),P()));
        }
        Pmat.sum(i,-1.0*Xsub.transpose()*SPS*Xsub);
        
        for(int j = i; j < Rmod; j++){
          MatrixXd Ssub2 = glmmr::Eigen_ext::submat(S(j),rows,rows);
          
          // Q and R matrices
          if constexpr (corr == SE::KR || corr == SE::KR2 || corr == SE::KRBoth){
            if(b==0){
              Q.add(MatrixXd::Zero(P(),P()));
              if(i < R && j < R) RR.add(MatrixXd::Zero(P(),P()));
            }
            Q.sum(counter,Xsub.transpose()*SPS*Ssub2*SigmaInvsub*Xsub);
            counter++;
            // RR here
            if(i < R && j < R){
              int scnd_idx = i + j*(R-1) - j*(j-1)/2;
              MatrixXd Ssub3 = glmmr::Eigen_ext::submat(S(Rmod+scnd_idx),rows,rows);
              RR.sum(counter_rr,Xsub.transpose() * SigmaInvsub * Ssub3 * SigmaInvsub * Xsub);
              counter_rr++;
            }
          }
          
          double oim_adj;
          if constexpr (imtype == IM::OIM){
            oim_adj = (residsub.transpose() * SPS * Ssub2 * SigmaInvsub * residsub)(0);
          }
          for(int k = 0; k < rows.size(); k++){
            for(int l = 0; l < rows.size(); l++){
              M_theta(i,j) += 0.5*SPS(k,l)*Ssub2(l,k);
              if constexpr (imtype == IM::OIM) M_theta_oim(i,j) -= oim_adj;
            }
          }
          if(i!=j)M_theta(j,i)=M_theta(i,j);
          if constexpr (imtype == IM::OIM) if(i!=j)M_theta_oim(j,i)=M_theta_oim(i,j);
        }
      }
    }
  } else {
    for(int i = 0; i < Rmod; i++){
      MatrixXd SPS = SigmaInv * S(i) * SigmaInv;
      Pmat.add(X.transpose() * SPS * X);
      for(int j = i; j < Rmod; j++){
        MatrixXd Ssub2 = S(j);
        
        if constexpr (corr == SE::KR || corr == SE::KR2 || corr == SE::KRBoth){
          Q.add(X.transpose()*SPS*Ssub2*SigmaInv*X);
          // RR here
          if(i < R && j < R){
            int scnd_idx = i + j*(R-1) - j*(j-1)/2;
            RR.add(X.transpose() * SigmaInv * S(Rmod+scnd_idx) * SigmaInv * X);
          }
        }
        
        double oim_adj;
        if constexpr (imtype == IM::OIM){
          oim_adj = (resid.transpose() * SPS * Ssub2 * SigmaInv * resid)(0);
        }
        M_theta(i,j) = 0;
        for(int k = 0; k < model.n(); k++){
          for(int l = 0; l <  model.n(); l++){
            M_theta(i,j) += 0.5*SPS(k,l)*Ssub2(l,k);
            if constexpr (imtype == IM::OIM) M_theta_oim(i,j) -= oim_adj;
          }
        }
        if(i!=j)M_theta(j,i)=M_theta(i,j);
        if constexpr (imtype == IM::OIM) if(i!=j)M_theta_oim(j,i)=M_theta_oim(i,j);
      }
    }
  }
  
  if constexpr (corr == SE::KR || corr == SE::KR2 || corr == SE::KRBoth){
    counter = 0;
    for(int i = 0; i < Rmod; i++){
      for(int j = i; j < Rmod; j++){
        M_theta_kr(i,j) += -1.0*(M*Q(counter)).trace() + 0.5*((M*Pmat(i)*M*Pmat(j)).trace());
        if(i!=j)M_theta_kr(j,i)=M_theta_kr(i,j);
        counter++;
      }
    }
  }
  
  if constexpr (imtype == IM::OIM){
    // add in the beta-theta part of the matrix
    MatrixXd Mbt(X.cols(),M_theta.cols());
    if(useBlock){
      for(int b = 0; b< sigma_blocks.size(); b++){
        Mbt.setZero();
        ArrayXi rows = Map<ArrayXi,Unaligned>(sigma_blocks[b].RowIndexes.data(),sigma_blocks[b].RowIndexes.size());
        MatrixXd Xsub = glmmr::Eigen_ext::submat(X,rows,ArrayXi::LinSpaced(P(),0,P()-1));
        MatrixXd SigmaInvsub = glmmr::Eigen_ext::submat(SigmaInv,rows,rows);
        VectorXd residsub(rows.size());
        for(int r = 0; r < rows.size(); r++) residsub(r) = resid(rows(r));
        for(int i =0; i < M_theta.cols(); i++){
          MatrixXd Ssub1 = glmmr::Eigen_ext::submat(S(i),rows,rows);
          Mbt.col(i) += Xsub.transpose() * SigmaInvsub * Ssub1 * SigmaInvsub * residsub;
        }
      }
    } else {
      for(int i =0; i < M_theta.cols(); i++) Mbt.col(i) = X.transpose() * SigmaInv * S(i) * SigmaInv * resid;
    }
    
    MatrixXd Moim(this->P()+Rmod,this->P()+Rmod);
    Moim.topLeftCorner(this->P(),this->P()) = M.llt().solve(MatrixXd::Identity(M.rows(),M.cols()));
    Moim.topRightCorner(this->P(),Rmod) = Mbt;
    Moim.bottomLeftCorner(Rmod,this->P()) = Mbt.transpose();
    Moim.bottomRightCorner(Rmod,Rmod) = -1.0*(M_theta + M_theta_oim);
    Moim = Moim.llt().solve(MatrixXd::Identity(Moim.rows(),Moim.cols()));
    M = Moim.topLeftCorner(this->P(),this->P());
    M_theta = -1.0*(M_theta + M_theta_oim);
  }
  
  if constexpr (corr == SE::KR || corr == SE::KR2 || corr == SE::KRBoth ) M_theta += M_theta_kr;
  M_theta = M_theta.llt().solve(MatrixXd::Identity(Rmod,Rmod));
  
  if constexpr (corr == SE::KR || corr == SE::KR2 || corr == SE::KRBoth ){
    for(int i = 0; i < (Rmod-1); i++){
      for(int j = (i+1); j < Rmod; j++){
        int scnd_idx = i + j*(Rmod-1) - j*(j-1)/2;
        meat += M_theta(i,j)*(Q(scnd_idx) + Q(scnd_idx).transpose() - Pmat(i)*M*Pmat(j) -Pmat(j)*M*Pmat(i));//(SigX.transpose()*partial1*PG*partial2*SigX);//
        if(i < R && j < R){
          scnd_idx = i + j*(R-1) - j*(j-1)/2;
          meat -= 0.5*M_theta(i,j)*(RR(scnd_idx));
        }
      }
    }
    for(int i = 0; i < Rmod; i++){
      int scnd_idx = i + i*(Rmod-1) - i*(i-1)/2;
      meat += M_theta(i,i)*(Q(scnd_idx) - Pmat(i)*M*Pmat(i));
      if(i < R){
        scnd_idx = i + i*(R-1) - i*(i-1)/2;
        meat -= 0.25*M_theta(i,i)*RR(scnd_idx);
      }
    }
    M_new = M + 2*M*meat*M;
    
#if defined(R_BUILD) && defined(ENABLE_DEBUG)
    Rcpp::Rcout << "\n(K-R) First correction matrix: \n" << 2*M*meat*M;
#endif
  } else {
    M_new = M;
  }
  
  CorrectionData<corr> out(this->P(),this->P(),Rmod,Rmod);
  out.vcov_beta = M_new;
  out.vcov_theta = M_theta;
  if constexpr (corr == SE::KRBoth) out.vcov_beta_second = M_new;
  
  // new improved correction
  // this code can also be sped up by doing it by blocks
  if constexpr (corr == SE::KR2 || corr == SE::KRBoth){
    MatrixXd SigX = SigmaInv*X;
    MatrixXd SS = MatrixXd::Zero(SigmaInv.rows(),SigmaInv.cols());
    for(int i = 0; i < Rmod; i++){
      for(int j = i; j < Rmod; j++){
        if(i < R && j < R){
          int scnd_idx = i + j*(R-1) - j*(j-1)/2;
          double rep = i == j ? 1.0 : 2.0;
          SS += rep * M_theta(i,j)*S(Rmod+scnd_idx)*SigmaInv;
        }
      }
    }
    SS.applyOnTheLeft(SigmaInv);
    MatrixXd XSXM = X.transpose()*SS*X*M;
    dblvec V(Rmod,0.0);
    for(int i = 0; i < Rmod; i++){
      V[i] += (SS*S(i)).trace();
      V[i] -= 2*((SigX.transpose()*S(i)*SS*X*M).trace());
      V[i] += ((XSXM*SigX.transpose()*S(i)*SigX*M).trace());
    }
    MatrixXd M_correction = MatrixXd::Zero(M.rows(),M.cols());
    for(int i = 0; i < Rmod; i++){
      for(int j = i; j < Rmod; j++){
        double rep = i == j ? 0.25 : 0.5;
        M_correction += rep*M_theta(i,j)*V[j]*M*Pmat(i)*M;
      }
    }
    
#if defined(R_BUILD) && defined(ENABLE_DEBUG)
    Rcpp::Rcout << "\n(K-R) Improved correction matrix: \n" << M_correction;
    Rcpp::Rcout << "\nV: ";
    for(const auto& v: V)Rcpp::Rcout << v << " ";
#endif
    
    if constexpr (corr == SE::KR2) {
      out.vcov_beta -= M_correction;
    } else {
      out.vcov_beta_second -= M_correction;
    }
  }
  
  // degrees of freedom correction
  
  double a1, a2, B, g, c1, c2, c3, v0, v1, v2, rhotop, rho;
  int mult = 1;
  VectorXd L = VectorXd::Zero(this->P());
  MatrixXd Theta(this->P(),this->P());
  for(int p = 0; p < L.size(); p++){
    L.setZero();
    L(p) = 1;
    double vlb = L.transpose() * M * L;
    Theta = (1/vlb)*(L*L.transpose());
    Theta = Theta*M;
    a1 = 0; 
    a2 = 0;
    for(int i = 0; i < Rmod; i++){
      for(int j = i; j < Rmod; j++){
        mult = i==j ? 1 : 2;
        a1 += mult*M_theta(i,j)*(Theta*Pmat(i)*M).trace()*(Theta*Pmat(j)*M).trace();
        a2 += mult*M_theta(i,j)*(Theta*Pmat(i)*M*Theta*Pmat(j)*M).trace();
      }
    }
    B = (a1 + 6*a2)*0.5;
    g = (2*a1 - 5*a2)/(3*a2);
    c1 = g/(3+2*(1-g));
    c2 = (1-g)/(3+2*(1-g));
    c3 = (3-g)/(3+2*(1-g));
    v0 = abs(1 + c1*B) < 1e-10 ? 0 : 1 + c1*B;
    v1 = 1 - c2*B;
    v2 = 1/(1 - c3*B);
    rhotop = abs(1-a2) < 1e-10 && abs(v1) < 1e-10 ? 1.0 : (1-a2)/v1;
    rho = rhotop*rhotop*v0*v2;
    out.dof(p) = 4 + 3/(rho-1);
    out.lambda(p) = (1-a2)*out.dof(p)/(out.dof(p) - 2);
  }
  
  return out;
}

template<typename modeltype>
inline MatrixXd glmmr::ModelMatrix<modeltype>::linpred()
{
  return (re.zu_.colwise()+(model.linear_predictor.xb()+model.data.offset));
}

template<typename modeltype>
inline VectorMatrix glmmr::ModelMatrix<modeltype>::b_score()
{
  MatrixXd zuOffset = re.Zu();
  zuOffset.colwise() += model.data.offset;
  MatrixMatrix hess = model.calc.jacobian_and_hessian(zuOffset);
  VectorMatrix out(hess.mat1.rows());
  out.mat = hess.mat1;
  out.mat *= -1.0;
  out.vec = hess.mat2.rowwise().sum();
  return out;
}

// type = 0 - raw
// type = 1 - pearson
// type = 2 - standardised
template<typename modeltype>
inline MatrixXd glmmr::ModelMatrix<modeltype>::residuals(const int type, bool conditional)
{
  int ncol_r = conditional ? re.u_.cols() : 1;
  MatrixXd resids(model.n(), ncol_r);
  VectorXd mu = glmmr::maths::mod_inv_func(model.linear_predictor.xb(), model.family.link);
  VectorXd yxb = model.data.y - mu;
  if(conditional){
    MatrixXd zuOffset = re.Zu();
    zuOffset.colwise() += model.data.offset;
    for(int i = 0; i < ncol_r; i++) resids.col(i) = yxb - zuOffset.col(i);
  } else {
    resids.col(0) = yxb;
  }
  if(type == 1){
    VectorXd var = glmmr::maths::marginal_var(mu, model.family.family, model.data.var_par);
    for(int i = 0; i < ncol_r; i++) resids.col(i).array() *= 1.0/sqrt(var(i));
  } else if(type == 2){
    VectorXd colmeans = resids.colwise().mean().transpose();
    double std_dev;
    for(int i = 0; i < ncol_r; i++){
      std_dev = sqrt((resids.col(i).array() - colmeans(i)).square().sum()/(resids.rows() - 1));
      resids.col(i).array() *= 1.0/std_dev; 
    }
  }
  return resids;
}


template<typename modeltype>
inline VectorMatrix glmmr::ModelMatrix<modeltype>::re_score(){
  VectorXd xbOffset = model.linear_predictor.xb() + model.data.offset;
  model.vcalc.parameters = dblvec(re.u(false).col(0).data(),re.u(false).col(0).data()+re.u(false).rows());
  model.vcalc.data = model.covariance.ZL();
  MatrixMatrix hess = model.vcalc.jacobian_and_hessian(Map<MatrixXd>(xbOffset.data(),xbOffset.size(),1));
  VectorMatrix out(Q());
  hess.mat1 *= -1.0;
  out.mat = hess.mat1 + MatrixXd::Identity(Q(),Q());
  out.vec = hess.mat2.rowwise().sum();
  out.vec -= re.u(false).col(0);
  return out;
}

template<typename modeltype>
inline BoxResults glmmr::ModelMatrix<modeltype>::box(){
  int r = P(); // to follow notation in Skene and Kenward
  BoxResults results(r);
  MatrixXd S = Sigma();
  MatrixXd X = model.linear_predictor.X();
  MatrixXd XtX = X.transpose()*X;
  XtX = XtX.llt().solve(MatrixXd::Identity(XtX.rows(),XtX.cols()));
  MatrixXd Px = X*XtX*X.transpose();
  MatrixXd A = MatrixXd::Identity(Px.rows(),Px.cols()) - Px;
  MatrixXd ASig = A*S;
  double asigtrace = ASig.trace();
  double trasig = ((ASig*ASig.transpose()).trace())/(asigtrace*asigtrace);
  MatrixXd Xr(model.n(),r-1);
  MatrixXd B(Px.rows(),Px.cols());
  MatrixXd XrtXr(r-1,r-1);
  for(int i = 0; i < r; i++){
    if(i == 0){
      Xr = X.rightCols(r-1);
    } else if(i == r-1){
      Xr = X.leftCols(r-1);
    } else {
      Xr.leftCols(i) = X.leftCols(i);
      Xr.rightCols(r-i-1) = Xr.rightCols(r-i-1);
    }
    XrtXr = Xr.transpose()*Xr;
    XrtXr = XrtXr.llt().solve(MatrixXd::Identity(r-1,r-1));
    B = Px - Xr*XrtXr*Xr.transpose();
    MatrixXd BSig = B*S;
    double bsigtrace = BSig.trace();
    double trbsig = ((BSig*BSig.transpose()).trace())/(bsigtrace*bsigtrace);
    double V = trbsig + trasig;
    results.dof[i] = ((4*V+1)-2)/(V-1);
    results.scale[i] = (X.rows() - r)*(results.dof[i] - 2)*bsigtrace/(results.dof[i] * asigtrace);
    results.test_stat[i] = (X.rows() - r)*((model.data.y.transpose()*B*model.data.y)(0))/((model.data.y.transpose()*A*model.data.y)(0));
    boost::math::fisher_f fdist(1.0,results.dof[i]);
    results.p_value[i] = 1 - boost::math::cdf(fdist,results.test_stat[i]/results.scale[i]);
  }
  return results;
}

template<typename modeltype>
inline MatrixXd glmmr::ModelMatrix<modeltype>::gradient_eta(const MatrixXd& v){
  
  ArrayXXd size_n_array(model.n(), v.cols());
  size_n_array.setZero();
  if(size_n_array.rows() != model.n())throw std::runtime_error("Size n array != n");
  size_n_array.colwise() += model.xb();
  SparseMatrix<double> ZL = model.covariance.ZL_sparse_new();
  if(ZL.cols() != v.rows())throw std::runtime_error("ZL cols != v rows");
  size_n_array += (ZL * v).array();
  
  switch(model.family.family){
  case Fam::poisson:
  {
    switch(model.family.link){
      case Link::identity:
      {
        size_n_array = size_n_array.inverse();
        size_n_array = size_n_array.colwise() * model.data.y.array();
        size_n_array -= 1.0;
        break;
      }
      default:
      {
        size_n_array = size_n_array.exp();
        ArrayXXd expmu(size_n_array);
        size_n_array.colwise() += -1.0 * model.data.y.array();
        size_n_array *= -1.0;
        size_n_array = size_n_array/expmu;
        break;
      }
  }
    break;
  }
  case Fam::bernoulli: case Fam::binomial:
  {
    switch(model.family.link){
  case Link::loglink:
  {
    ArrayXXd logitxb = 1.0 - size_n_array.exp();
    logitxb = logitxb.inverse();
    logitxb *= size_n_array.exp();
    size_n_array = logitxb.colwise() * (model.data.y.array() - model.data.variance);
    size_n_array.colwise() += model.data.y.array();
    break;
  }
  case Link::identity:
  {
    ArrayXXd n_array2 = 1.0 - size_n_array;
    n_array2 = n_array2.inverse();
    n_array2.colwise() *= (model.data.variance - model.data.y.array());
    size_n_array = size_n_array.inverse();
    size_n_array.colwise() *= model.data.y.array();
    size_n_array -= n_array2;
    break;
  }
  case Link::probit:
  {
    ArrayXXd n_array2(size_n_array.rows(), size_n_array.cols());
    ArrayXXd cdf(size_n_array);
    ArrayXXd pdf(size_n_array);
    cdf.unaryExpr(&glmmr::maths::gaussian_cdf);
    pdf.unaryExpr(&glmmr::maths::gaussian_pdf);
    size_n_array = pdf / cdf;
    n_array2 = -1.0 * (pdf / (1.0 - cdf));
    size_n_array.colwise() *= model.data.y.array();
    n_array2.colwise() *= (model.data.variance - model.data.y.array());
    size_n_array += n_array2;
    break;
  }
  default:
    //logit
  {
    ArrayXXd logitxb = (size_n_array.exp().inverse() + 1.0).inverse();
    // size_n_array = logitxb;
    // size_n_array.colwise() += model.data.variance * model.data.y.array();
    // size_n_array *= -1.0;
    // size_n_array = size_n_array / ()
    size_n_array = logitxb;
    size_n_array.colwise() *= (model.data.variance - model.data.y.array());
    logitxb *= -1.0;
    logitxb += 1.0;
    logitxb.colwise() *= model.data.y.array();
    size_n_array = logitxb - size_n_array;
    break;
  }
  }
    break;
  }
  case Fam::gaussian:
  {
    switch(model.family.link){
  case Link::loglink:
  {
    ArrayXXd narray2 = size_n_array.exp();
    narray2.colwise() *= model.data.weights;
    size_n_array *= -1.0;
    size_n_array.colwise() += model.data.y.array();
    size_n_array *= narray2;
    break;
  }
  default:
  {
    size_n_array *= -1.0;
    size_n_array.colwise() += model.data.y.array();
    size_n_array.colwise() *= model.data.weights/model.data.var_par;
    break;
  }
  }
    break;
  }
  case Fam::gamma:
  {
    switch(model.family.link){
  case Link::inverse:
  {
    size_n_array = size_n_array.inverse();
    size_n_array.colwise() -= model.data.y.array();
    break;
  }
  case Link::identity:
  {
    size_n_array = size_n_array.inverse();
    size_n_array.colwise() *= model.data.y.array();
    size_n_array -= 1.0;
    break;
  }
  default:
    //log
  {
    size_n_array *= -1.0;
    size_n_array = size_n_array.exp();
    size_n_array.colwise() *= model.data.y.array();
    break;
  }
  }
    break;
  }
  case Fam::beta:
  {
    throw std::runtime_error("Beta is currently disabled");
    /*#pragma omp parallel for 
     for(int i = 0; i < model.n(); i++){
     size_n_array(i) = exp(size_n_array(i))/(exp(size_n_array(i))+1);
     size_n_array(i) = (size_n_array(i)/(1+exp(size_n_array(i)))) * model.data.var_par * (log(model.data.y(i)) - log(1- model.data.y(i)) - boost::math::digamma(size_n_array(i)*model.data.var_par) + boost::math::digamma((1-size_n_array(i))*model.data.var_par));
     }
     break;*/
  }
  case Fam::quantile: case Fam::quantile_scaled: 
  {
    throw std::runtime_error("Quantile is currently disabled");
    break;
  }
  case Fam::exponential:
  {
    throw std::runtime_error("Gradient_eta not yet available with exponential distribution");
    break;
  }
  }
  return size_n_array.matrix();
}

template<typename modeltype>
inline VectorXd glmmr::ModelMatrix<modeltype>::log_gradient(const VectorXd &v,
                                                            bool betapars){
  MatrixXd vm(v.size(), 1);
  vm.col(0) = v;
  ArrayXd size_n_array(model.n());
  size_n_array = (gradient_eta(vm)).col(0).array();
  ArrayXd size_q_array = ArrayXd::Zero(Q());
  ArrayXd size_p_array = ArrayXd::Zero(P());
  SparseMatrix<double> ZLt = model.covariance.ZL_sparse_new();
  ZLt.transpose();
  
  switch(model.family.family){
  case Fam::poisson: case Fam::bernoulli: case Fam::binomial: case Fam::beta: case Fam::quantile: case Fam::quantile_scaled: case Fam::exponential:
  {
    if(betapars){
    size_p_array =  (model.linear_predictor.X().transpose()*size_n_array.matrix()).array();
  } else {
    size_q_array =  (ZLt * size_n_array.matrix()).array()-v.array();
  }
  break;
  }
  case Fam::gaussian:
  {
    if(betapars){
    size_p_array += ((1.0/(model.data.var_par))*(model.linear_predictor.X().transpose()*size_n_array.matrix())).array();
  } else {
    size_q_array = (ZLt * size_n_array.matrix()).array();
    size_q_array *= 1.0/(model.data.var_par);
    size_q_array -= v.array();
  }
  break;
  }
  case Fam::gamma:
  {
    switch(model.family.link){
  case Link::inverse:
  {
    if(betapars){
    size_p_array += (model.linear_predictor.X().transpose()*(size_n_array.matrix()-model.data.y)*model.data.var_par).array();
  } else {
    size_q_array = (ZLt * size_n_array.matrix()).array();
    size_q_array *= model.data.var_par;
    size_q_array -= v.array();
  }
  break;
  }
  case Link::identity:
  {
    if(betapars){
    size_p_array += (model.linear_predictor.X().transpose()*((model.data.y.array()*size_n_array*size_n_array).matrix() - size_n_array.matrix())*model.data.var_par).array();
  } else {
    size_q_array = (ZLt * size_n_array.matrix()).array();
    size_q_array *= model.data.var_par;
    size_q_array -= v.array();
  }
  break;
  }
  default:
    //log
  {
    if(betapars){
    size_p_array += (model.linear_predictor.X().transpose()*(model.data.y.array()*size_n_array-1).matrix()*model.data.var_par).array();
  } else {
    size_q_array = (ZLt * size_n_array.matrix()).array();
    size_q_array *= model.data.var_par;
    size_q_array -= v.array();
  }
  break;
  }
  }
    break;
  }
  }
  return betapars ? size_p_array.matrix() : size_q_array.matrix();
}

template<typename modeltype>
inline void glmmr::ModelMatrix<modeltype>::posterior_u_samples(const int niter,
                                                               const bool reml,
                                                               const bool loglik,
                                                               const bool append)
{
  
  if constexpr (std::is_same_v<modeltype,bits_hsgp>){
    if(model.covariance.Q() != re.u_.rows()){
      re.u_.resize(model.covariance.Q(),1);
      re.u_.setZero();
    }
  }
  ArrayXd xb = model.linear_predictor.xb().array() + model.data.offset.array();
  ArrayXd eta = xb;
  eta += model.covariance.ZLu(re.u_mean_).array();
  VectorXd W_(eta.size());

  switch(model.family.family){
  case Fam::gaussian:
    if(model.family.link == Link::identity){
      W_ = (model.data.variance.inverse() *  model.data.weights).matrix();
    } else {
      throw std::runtime_error("Analtyic posterior only available with canonical link");
    }
    break;
  case Fam::binomial: case Fam::bernoulli:
    if(model.family.link == Link::logit){
      ArrayXd logitp = (eta.exp().inverse() + 1.0).inverse();
      W_ = (model.data.variance * logitp * (1- logitp)).matrix();
    } else {
      throw std::runtime_error("Analtyic posterior only available with canonical link");
    }
    break;
  case Fam::poisson:
    if(model.family.link == Link::loglink){
      W_ = eta.exp().matrix();
    } else {
      throw std::runtime_error("Analtyic posterior only available with canonical link");
    }
    break;
  case Fam::exponential:
    if(model.family.link == Link::loglink){
      ArrayXd mu = eta.exp();
      W_ = ArrayXd::Ones(eta.size()).matrix();  // Constant weight = 1
    }else {
      throw std::runtime_error("Analtyic posterior only available with canonical link");
    }
  break;
  default:
    throw std::runtime_error("Analtyic posterior only available with Gaussian, Poisson, and Binomial");
    break;
  }

  const MatrixXd ZL = model.covariance.ZL();
  const MatrixXd ZLt = ZL.transpose();
  const int n_cols = ZL.cols();
  VectorXd Mb(n_cols);
  MatrixXd Vb(n_cols, n_cols);
  Vb.setIdentity();
  LLT<MatrixXd> llt_Pb;
  VectorXd yb(n_cols);
  MatrixXd WZL(W_.size(),n_cols);
  MatrixXd LWL = MatrixXd::Identity(n_cols,n_cols);

  eta = maths::mod_inv_func(eta.matrix(), model.family.link).array();
  if(model.family.family == Fam::binomial) eta.array().colwise() *= model.data.variance;
  VectorXd resid = model.data.y - eta.matrix();

  if(model.family.family == Fam::gaussian) {
    WZL.noalias() = (ZL.array().colwise() * W_.array()).matrix();
    LWL.noalias() = ZLt * WZL;
    LWL.diagonal().array() += 1.0;
    yb.noalias() = WZL.transpose() * (model.data.y - xb.matrix());
    llt_Pb.compute(LWL);
    Mb = llt_Pb.solve(yb);
    llt_Pb.solveInPlace(Vb);
  } else {
    VectorXd b = re.u_mean_;
    VectorXd bnew(b);
    double diff = 1.0;
    int itero = 0;
    VectorXd u(b.size());
    while(diff > 1e-6 && itero < 10) {
      u.noalias() = ZL * b;
      eta = xb + u.array();
      if(model.family.family == Fam::binomial || model.family.family == Fam::bernoulli) {
        ArrayXd exp_neg_eta = (-eta).exp();
        ArrayXd logitp = 1.0 / (1.0 + exp_neg_eta);
        ArrayXd var_p = model.data.variance * logitp;
        W_ = (var_p * (1.0 - logitp)).matrix();
        eta = maths::mod_inv_func(eta.matrix(), model.family.link).array();
        if(model.family.family == Fam::binomial) eta.array().colwise() *= model.data.variance;
        resid = model.data.y - eta.matrix();

      } else if(model.family.family == Fam::poisson) {
        ArrayXd exp_eta = eta.exp();
        W_ = exp_eta.matrix();
        eta = maths::mod_inv_func(eta.matrix(), model.family.link).array();
        resid = model.data.y - eta.matrix();
      } else if(model.family.family == Fam::exponential) {
        // Log link
        ArrayXd mu = eta.exp();
        W_ = ArrayXd::Ones(eta.size()).matrix();  // Constant weight = 1
        eta = maths::mod_inv_func(eta.matrix(), model.family.link).array();
        resid = ((model.data.y.array() - eta) / eta).matrix();  // (y - mu) / mu
      }
      WZL.noalias() = (ZL.array().colwise() * W_.array()).matrix();
      LWL.noalias() = ZLt * WZL;
      yb.noalias() = LWL * b + ZLt * resid;
      LWL.diagonal().array() += 1.0;
      llt_Pb.compute(LWL);
      bnew = llt_Pb.solve(yb);
      diff = (b - bnew).array().abs().maxCoeff();
      itero++;
      b.swap(bnew);
    }

    Mb = b;
    re.u_mean_ = Mb;
    if(reml){
      MatrixXd X = model.linear_predictor.X();
      MatrixXd WX(W_.size(),model.linear_predictor.P());
      WX.noalias() = (X.array().colwise() * W_.array()).matrix();
      MatrixXd XWX = X.transpose() * WX;
      MatrixXd C = XWX.llt().solve(MatrixXd::Identity(XWX.rows(), XWX.cols()));
      MatrixXd B = X.transpose() * WZL;
      MatrixXd Corr = B.transpose() * C * B;
      Vb -= Corr;
    }
    llt_Pb.solveInPlace(Vb);
    re.u_var_diag_ = Vb.diagonal(); 

  }

  MatrixXd unew(re.u_.rows(), niter);
  std::random_device rd;
  std::mt19937 gen(rd());
  std::normal_distribution<double> d(0.0, 1.0);

  double* data = unew.data();
  for(int i = 0; i < unew.size(); ++i) {
    data[i] = d(gen);
  }

  LLT<MatrixXd> llt(Vb);
  MatrixXd LVb = llt.matrixL();
  VectorXd v(re.u_.rows());
  if(append && loglik) throw std::runtime_error("SAEM and importance sampling not currently compatible");
  unew = LVb * unew;
  bool action_append = append;
  if(append && re.u_.cols() == 1)action_append = false;
  if(action_append){
    int currcolsize = re.u_.cols();
    unew.colwise() += re.u_mean_;
    re.u_.conservativeResize(NoChange,currcolsize + niter);
    re.scaled_u_.conservativeResize(NoChange,currcolsize + niter);
    re.u_solve_.conservativeResize(NoChange,currcolsize + niter);
    re.zu_.conservativeResize(NoChange,currcolsize + niter);
    re.u_.rightCols(niter).noalias() = unew;
  } else {
    if(re.u_.cols() != niter){
      re.u_.resize(NoChange, niter);
      re.u_solve_.resize(NoChange, niter);
      re.scaled_u_.resize(NoChange, niter);
      re.zu_.resize(NoChange, niter);
      re.u_weight_.resize(niter);
      if(loglik) re.u_loglik_.resize(niter);
    }
    re.u_.noalias() = unew;
    if(loglik){
#pragma omp parallel for
      for(int i = 0; i < re.u_.cols(); i++){
        v = llt.solve(re.u_.col(i));
        re.u_loglik_(i) = -0.5 * v.dot(re.u_.col(i));
      }
    }
    re.u_.colwise() += re.u_mean_;
  }
  re.update_zu(loglik);
}


template<>
inline void glmmr::ModelMatrix<bits_hsgp>::posterior_u_samples(
    const int niter, const bool reml, const bool loglik, const bool append)
{
  if(model.covariance.Q() != re.u_.rows()){
    re.u_.resize(model.covariance.Q(), 1);
    re.u_.setZero();
  }
  
  ArrayXd xb = model.linear_predictor.xb().array() + model.data.offset.array();
  const int M = model.covariance.Q();
  const int n = model.n();
  ArrayXd eta(n);
  VectorXd W_(n);
  
  // A = Z * Phi: maps u to observation space (no Lambda scaling)
  MatrixXd A = model.covariance.ZPhi();    // n × M
  MatrixXd At = A.transpose();             // M × n
  ArrayXd inv_lambda = 1.0 / model.covariance.LambdaSPD();
  
  // IRLS for posterior mode in u-space
  VectorXd b = re.u_mean_;
  if(b.size() != M){ b.resize(M); b.setZero(); }
  VectorXd bnew(M);
  double diff = 1.0;
  int itero = 0;
  MatrixXd AtWA(M, M);
  MatrixXd P(M, M);
  VectorXd rhs(M);
  LLT<MatrixXd> llt_P;
  
  while(diff > 1e-6 && itero < 20){
    eta = xb + (A * b).array();
    VectorXd W_(n);
    switch(model.family.family){
    case Fam::binomial: case Fam::bernoulli: {
      ArrayXd logitp = 1.0 / (1.0 + (-eta).exp());
      W_ = (model.data.variance * logitp * (1.0 - logitp)).matrix();
      ArrayXd mu = model.data.variance * logitp;
      rhs = At * (model.data.y - mu.matrix()).matrix();
      break;
    }
    case Fam::poisson: {
      W_ = eta.exp().matrix();
      rhs = At * (model.data.y - eta.exp().matrix()).matrix();
      break;
    }
    case Fam::gaussian: {
      W_ = (model.data.variance.inverse() * model.data.weights).matrix();
      rhs = At * (model.data.y - eta.matrix()).matrix();
      break;
    }
    default:
      throw std::runtime_error("HSGP posterior only for Gaussian, Poisson, Binomial");
    }
    
    // P = A^T W A + diag(1/Lambda)
    AtWA.noalias() = At * W_.asDiagonal() * A;
    P = AtWA;
    P.diagonal() += inv_lambda.matrix();
    
    // rhs = AtWA * b + A^T * residuals
    rhs += AtWA * b;
    
    llt_P.compute(P);
    bnew = llt_P.solve(rhs);
    diff = (b - bnew).array().abs().maxCoeff();
    b.swap(bnew);
    itero++;
  }
  
  re.u_mean_ = b;
  MatrixXd Vu = MatrixXd::Identity(M, M);
  llt_P.solveInPlace(Vu);
  re.u_var_diag_ = Vu.diagonal(); 
  
  if(re.u_.cols() != niter){
    re.u_.resize(M, niter);
    re.u_solve_.resize(M, niter);
    re.scaled_u_.resize(M, niter);
    re.zu_.resize(n, niter);
    re.u_weight_.resize(niter);
    if(loglik) re.u_loglik_.resize(niter);
  }
  
  if(niter == 1){
    // Laplace: posterior mode only
    re.u_.col(0) = re.u_mean_;
    if(loglik) re.u_loglik_(0) = 0.0;
    re.update_zu(loglik);
  } else {
    // Existing MC sampling code
    LLT<MatrixXd> llt_Vu(Vu);
    MatrixXd L_Vu = llt_Vu.matrixL();
    
    MatrixXd znorm(M, niter);
    std::random_device rd;
    std::mt19937 gen(rd());
    std::normal_distribution<double> d(0.0, 1.0);
    for(int i = 0; i < znorm.size(); ++i) znorm.data()[i] = d(gen);
    
    if(loglik){
      for(int i = 0; i < niter; i++){
        re.u_loglik_(i) = -0.5 * znorm.col(i).squaredNorm();
      }
    }
    
    re.u_ = (L_Vu * znorm).colwise() + re.u_mean_;
    re.update_zu(loglik);
  }
}

template<>
inline void glmmr::ModelMatrix<bits_spde>::posterior_u_samples(
    const int niter, const bool reml, const bool loglik, const bool append)
{
  auto& cov = model.covariance;
  
  if(niter == 1)throw std::runtime_error("SPDE requires MC > 1");
  
  if(!cov.spde_loaded){
    throw std::runtime_error("SPDE posterior_u_samples: spde_data() has not been loaded.");
  }
  if(cov.Q() != re.u_.rows()){
    re.u_.resize(cov.Q(), 1);
    re.u_.setZero();
  }
  
  ArrayXd xb = model.linear_predictor.xb().array() + model.data.offset.array();
  const int M = cov.Q();               // n_v
  const int n = model.n();
  ArrayXd eta(n);
  VectorXd W_(n);
  
  // Observation-to-vertex projector: n × n_v, cached in the covariance class
  const SparseMatrix<double>& A  = cov.ZA_;        // n × M
  SparseMatrix<double> At        = A.transpose(); // M × n
  
  // Ensure Q factorisation is current (the θ step upstream should have refreshed it)
  if(!cov.chol_Q_current) cov.refactor_Q();
  
  // IRLS for posterior mode in u-space. We solve, at each step,
  //   P (b_new - b) = A^T (y - μ(b)) - Q b
  // equivalently   P b_new = A^T W A b + A^T (y - μ(b))
  VectorXd b = re.u_mean_;
  if(b.size() != M){ b.resize(M); b.setZero(); }
  VectorXd bnew(M);
  double diff = 1.0;
  int itero = 0;
  
  SparseMatrix<double> AtWA(M, M);
  SparseMatrix<double> P(M, M);
  VectorXd rhs(M);
  
  while(diff > 1e-6 && itero < 20){
    eta = xb + (A * b).array();
    
    switch(model.family.family){
    case Fam::binomial: case Fam::bernoulli: {
      ArrayXd logitp = 1.0 / (1.0 + (-eta).exp());
      W_  = (model.data.variance * logitp * (1.0 - logitp)).matrix();
      ArrayXd mu = model.data.variance * logitp;
      rhs = At * (model.data.y - mu.matrix());
      break;
    }
    case Fam::poisson: {
      W_  = eta.exp().matrix();
      rhs = At * (model.data.y - eta.exp().matrix());
      break;
    }
    case Fam::gaussian: {
      W_  = (model.data.variance.inverse() * model.data.weights).matrix();
      rhs = At * (model.data.y - eta.matrix());
      break;
    }
    default:
      throw std::runtime_error("SPDE posterior only for Gaussian, Poisson, Binomial.");
    }
    
    // Compute AtWA once — used both for rhs update and inside P (via refactor_P)
    SparseMatrix<double> AtWA = At * W_.asDiagonal() * A;
    rhs.noalias() += AtWA * b;
    
    // Covariance class assembles P = AtWA + Q_mat internally and factorises it
    cov.refactor_P(W_);
    
    bnew = cov.chol_P.solve(rhs);
    if(cov.chol_P.info() != Success){
      throw std::runtime_error("SPDE posterior: P solve failed.");
    }
    diff = (b - bnew).array().abs().maxCoeff();
    b.swap(bnew);
    itero++;
  }
  
  re.u_mean_ = b;
  
  // Allocate sample storage
  if(re.u_.cols() != niter){
    re.u_.resize(M, niter);
    re.u_solve_.resize(M, niter);
    re.scaled_u_.resize(M, niter);
    re.zu_.resize(n, niter);
    re.u_weight_.resize(niter);
    if(loglik) re.u_loglik_.resize(niter);
  }
  
  // Posterior draws: u_k = b + x_k, where x_k ~ N(0, P^{-1}).
  // Since P = R_P R_P^T  (NaturalOrdering, no permutation),
  //   z_k ~ N(0, I); solve R_P^T x_k = z_k  ⇒  x_k ~ N(0, P^{-1}).
  MatrixXd znorm(M, niter);
  std::random_device rd;
  std::mt19937 gen(rd());
  std::normal_distribution<double> d(0.0, 1.0);
  for(int i = 0; i < znorm.size(); ++i) znorm.data()[i] = d(gen);
  
  if(loglik){
    // log q(u_k) up to iteration-wise constant = -0.5 * z_k^T z_k
    for(int i = 0; i < niter; ++i){
      re.u_loglik_(i) = -0.5 * znorm.col(i).squaredNorm();
    }
  }
  
  // Solve R_P^T X = Z  (upper triangular solve with L^T = R_P^T).
  // SimplicialLLT exposes matrixU() = L^T, and .solve() on it handles the triangular system.
  MatrixXd X = cov.chol_P.matrixU().solve(znorm);
  X = cov.chol_P.permutationP().transpose() * X;
  re.u_ = X.colwise() + re.u_mean_;
  re.update_zu(loglik);
  
}
