#pragma once

#include <chrono>
#include "general.h"
#include "modelbits.hpp"
#include "randomeffects.hpp"
#include "modelmatrix.hpp"
#include "modeloptim.hpp"

namespace glmmr {

using namespace Eigen;

enum class MarginType {
  DyDx = 0,
    Diff = 1,
    Ratio = 2
};

template<class>
struct check_type : std::false_type {};

template<>
struct check_type<glmmr::ModelBits<glmmr::Covariance, glmmr::LinearPredictor> > : std::true_type {};

template<>
struct check_type<glmmr::ModelBits<glmmr::ar1Covariance, glmmr::LinearPredictor> > : std::true_type {};

template<typename modeltype>
class Model {
public:
  // model objects
  modeltype                       model;
  glmmr::RandomEffects<modeltype> re;
  glmmr::ModelMatrix<modeltype>   matrix;
  glmmr::ModelOptim<modeltype>    optim;
  
  Model(const Model&) = default;
  Model(Model&&) = default;
  Model& operator=(const Model&) = default;
  Model& operator=(Model&&) = default;
  
  // constructor
  template <typename C = modeltype, typename = std::enable_if_t<!std::is_same_v<C, bits_ar1> > >
  Model(const std::string& formula_,const ArrayXXd& data_,const strvec& colnames_,std::string family_,std::string link_) : model(formula_,data_,colnames_,family_,link_), 
  re(model), 
  matrix(model,re,check_type<modeltype>::value,check_type<modeltype>::value),  
  optim(model,matrix,re) {}
  
  template <typename C = modeltype, typename = std::enable_if_t<std::is_same_v<C, bits_ar1> > >
  Model(const std::string& formula_,
        const ArrayXXd& data_,
        const ArrayXXd& cov_data_,
        const strvec& colnames_,
        const strvec& colnames_cov_,
        std::string family_, 
        std::string link_,
        const int T) : model(formula_,data_,cov_data_,colnames_,colnames_cov_,family_,link_,T), 
        re(model), 
        matrix(model,re,check_type<modeltype>::value,check_type<modeltype>::value),  
        optim(model,matrix,re) {}
  
  //functions
  virtual void    set_offset(const VectorXd& offset_);
  virtual void    set_weights(const ArrayXd& weights_);
  virtual void    set_y(const VectorXd& y_);
  virtual void    update_beta(const dblvec &beta_);
  virtual void    update_theta(const dblvec &theta_);
  virtual void    update_u(const MatrixXd &u_, bool append = false);
  virtual void    reset_u(); // just resets the random effects samples to zero
  virtual void    set_trace(int trace_);
  void            fit(const int niter, const int max_iter, const bool start_ml_beta, const double tol, const int hist, const int k0);
  void            check_for_errors(const std::string& caller = "unknown", bool full_check = false);
  virtual dblpair marginal(const MarginType type,
                             const std::string& x,
                             const strvec& at,
                             const strvec& atmeans,
                             const strvec& average,
                             const RandomEffectMargin re_type,
                             const SE se_type,
                             const IM im_type,
                             const dblpair& xvals,
                             const dblvec& atvals,
                             const dblvec& atrevals);
                                             
};

}

template<typename modeltype>
inline void glmmr::Model<modeltype>::set_offset(const VectorXd& offset_){
  model.data.set_offset(offset_);
}

template<typename modeltype>
inline void glmmr::Model<modeltype>::set_weights(const ArrayXd& weights_){
  model.data.set_weights(weights_);
  if((weights_ != 1.0).any()){
    model.weighted = true;
  }
}

template<typename modeltype>
inline void glmmr::Model<modeltype>::set_y(const VectorXd& y_){
  model.data.update_y(y_);
}

template<typename modeltype>
inline void glmmr::Model<modeltype>::update_beta(const dblvec &beta_){
  optim.update_beta(beta_);
}

template<typename modeltype>
inline void glmmr::Model<modeltype>::update_theta(const dblvec &theta_){
  model.covariance.update_parameters(theta_);
  bool weighted = re.u_weight_.maxCoeff() - re.u_weight_.minCoeff() > 1e-10;
  re.update_zu(weighted);
}

template<typename modeltype>
inline void glmmr::Model<modeltype>::reset_u(){
  re.u_.resize(model.covariance.Q(),1);
  re.u_.setZero();
  re.scaled_u_.resize(model.covariance.Q(),1);
  re.scaled_u_.setZero();
  re.u_solve_.resize(model.covariance.Q(),1);
  re.u_solve_.setZero();
  re.zu_.resize(NoChange,1);
  re.zu_.setZero();
}

template<typename modeltype>
inline void glmmr::Model<modeltype>::update_u(const MatrixXd &u_, bool append){
  if(u_.rows()!=model.covariance.Q())throw std::runtime_error(std::to_string(u_.rows())+" rows provided, "+std::to_string(model.covariance.Q())+" expected");
  bool action_append = append;
  // if HSGP then check and update the size of u is m has changed
  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();
    }
  }
  
  int newcolsize = u_.cols();
  int currcolsize = re.u_.cols();
  // check if the existing samples are a single column of zeros - if so remove them
  if(append && re.u_.cols() == 1 && re.u_.col(0).isZero()) action_append = false;
  
  if(action_append){
    re.u_.conservativeResize(NoChange,currcolsize + newcolsize);
    re.zu_.conservativeResize(NoChange,currcolsize + newcolsize);
    re.scaled_u_.conservativeResize(NoChange,currcolsize + newcolsize);
    re.u_solve_.conservativeResize(NoChange,currcolsize + newcolsize);
    re.u_.rightCols(newcolsize) = u_;
    optim.ll_current.resize(currcolsize + newcolsize,NoChange);
  } else {
    if(u_.cols()!=re.u_.cols()){
      #if defined(ENABLE_DEBUG) && defined(R_BUILD)
        Rcpp::Rcout << "\nResize u: " << model.covariance.Q() << "x" << u_.cols();
      #endif
      re.u_.resize(NoChange,newcolsize);
      re.zu_.resize(NoChange,newcolsize);
      re.scaled_u_.resize(NoChange,newcolsize);
      re.u_solve_.resize(NoChange,newcolsize);
      re.u_weight_.resize(newcolsize);
    }
    re.u_ = u_;
    if(re.u_.cols() != optim.ll_current.rows()) optim.ll_current.resize(newcolsize,NoChange);
  }
  MatrixXd Z = model.covariance.Z();
  if(re.zu_.rows() != model.n())re.zu_.resize(model.n(),NoChange);
  if(re.scaled_u_.rows() != re.u_.rows())re.scaled_u_.resize(re.u_.rows(),NoChange);
  if(re.u_solve_.rows() != re.u_.rows())re.u_solve_.resize(re.u_.rows(),NoChange);
  re.update_zu(false);
}

template<typename modeltype>
inline void glmmr::Model<modeltype>::set_trace(int trace_){
  optim.trace = trace_;
  model.trace = trace_;
}

// marginal effects:
// type is margin type 
// x is name of variable to calculate margin
// at is fixed effects at a set value specified in atvals
// atmeans specifies the fixed effects to set at their mean value
// average specifies the fixed effects to average over
// re_type is random effects margin type
// se_type is the standard error type
// xvals values of the x variable at which to evaluate marginal effect
// atvals is the values for at argument
// atrevals is the random effects values if re_type is At
template<typename modeltype>
inline dblpair glmmr::Model<modeltype>::marginal(const MarginType type,
                                                 const std::string& x,
                                                 const strvec& at,
                                                 const strvec& atmeans,
                                                 const strvec& average,
                                                 const RandomEffectMargin re_type,
                                                 const SE se_type,
                                                 const IM im_type,
                                                 const dblpair& xvals,
                                                 const dblvec& atvals,
                                                 const dblvec& atrevals){
// #pragma omp declare reduction(vec_dbl_plus : std::vector<double> :                                    \
//   std::transform(omp_out.begin(), omp_out.end(), omp_in.begin(), omp_out.begin(), std::plus<double>())) \
//   initializer(omp_priv = decltype(omp_orig)(omp_orig.size()))
// removing parallelisation for now as it is not producing correct results: todo find problem
  
  int total_p = at.size() + atmeans.size() + average.size() + 1;
  int intercept = 1- (int)model.linear_predictor.form.RM_INT;
  
  if(total_p != (model.linear_predictor.P() - intercept))throw std::runtime_error("All variables must be named");
  if(at.size() != atvals.size())throw std::runtime_error("Not enough values specified for at");
  //if(re_type == RandomEffectMargin::Average && re.zu_.cols()<=1)Rcpp::warning("No MCMC samples of random effects. Random effects will be set at estimated values.");
    
  bool single_row = true;
  MatrixXd newXdata(1,model.linear_predictor.calc.data.cols());
  int P = model.linear_predictor.P();
  int N = 1;
  auto xidx = std::find(model.linear_predictor.calc.data_names.begin(),model.linear_predictor.calc.data_names.end(),x);
  int xcol = xidx - model.linear_predictor.calc.data_names.begin();
  
  if(average.size() > 0 || (total_p == 1 && (re_type == RandomEffectMargin::Average || re_type == RandomEffectMargin::AtEstimated))){
    single_row = false;
    N = model.n();
    newXdata.conservativeResize(model.n(),NoChange);
    newXdata.col(xcol) = model.linear_predictor.calc.data.col(xcol);
    
   if(re_type == RandomEffectMargin::At && static_cast<int>(atrevals.size()) != model.covariance.Q())throw std::runtime_error("Need to provide values for u vector");
   
   for(const auto& p: average){
     auto colidx = std::find(model.linear_predictor.calc.data_names.begin(),model.linear_predictor.calc.data_names.end(),p);
     if(colidx != model.linear_predictor.calc.data_names.end()){
       int pcol = colidx - model.linear_predictor.calc.data_names.begin();
       newXdata.col(pcol) = model.linear_predictor.calc.data.col(pcol);
     } else {
       throw std::runtime_error("Variable "+p+" not in data names");  
     }
   }
  } else {
    newXdata(0,xcol) = xvals.first;
    if(re_type == RandomEffectMargin::At && atrevals.size() != 1)throw std::runtime_error("Need to provide single value for Zu");
    if(re_type == RandomEffectMargin::AtEstimated)throw std::runtime_error("All covariates are at fixed values, cannot used estimated random effects.");
  }
  
  if(at.size() > 0){
    for(int p = 0; p < at.size(); p++){
      auto colidx = std::find(model.linear_predictor.calc.data_names.begin(),model.linear_predictor.calc.data_names.end(),at[p]);
      if(colidx != model.linear_predictor.calc.data_names.end()){
        int pcol = colidx - model.linear_predictor.calc.data_names.begin();
        for(int i = 0; i < newXdata.rows(); i++){
          newXdata(i,pcol) = atvals[p];
        }
      } else {
        throw std::runtime_error("Variable "+at[p]+" not in data names");  
      }
    }
  }
  if(atmeans.size() > 0){
    for(int p = 0; p < atmeans.size(); p++){
      auto colidx = std::find(model.linear_predictor.calc.data_names.begin(),model.linear_predictor.calc.data_names.end(),atmeans[p]);
      if(colidx != model.linear_predictor.calc.data_names.end()){
        int pcol = colidx - model.linear_predictor.calc.data_names.begin();
        double xmean = 0;
        for(int i = 0; i < model.n(); i++) xmean += model.linear_predictor.calc.data(i,pcol);
        xmean *= (1.0/model.n());
        for(int i = 0; i < newXdata.rows(); i++){
          newXdata(i,pcol) = xmean;
        }
      } else {
        throw std::runtime_error("Variable "+atmeans[p]+" not in data names");  
      }
    }
  }
  
  // now create the new calculator object
  glmmr::calculator mcalc = model.linear_predictor.calc;
  mcalc.instructions.push_back(Do::PushExtraData);
  mcalc.instructions.push_back(Do::Add);
  glmmr::linear_predictor_to_link(mcalc,model.family.link);
  mcalc.data.resize(newXdata.rows(),NoChange);
  mcalc.data = newXdata;
  
  dblpair result;
  VectorXd delta = VectorXd::Zero(P);
  MatrixXd M(P,P);
  
#if defined(R_BUILD) && defined(ENABLE_DEBUG)
  Rcpp::Rcout << "\nMARGINS\nN: " << N << " xcol: " << xcol;
  Rcpp::Rcout << "\nTest calculator\nUsing X: " << newXdata.row(0);
  dblvec test_result = mcalc.calculate<CalcDyDx::XBeta>(0,0,xcol,0);
  Rcpp::Rcout << "\nValues: ";
  for(const auto& i: test_result) Rcpp::Rcout << i << " ";
#endif
  
  switch(se_type){
    case SE::KR:
      {
        if(im_type == IM::EIM){
          CorrectionData<SE::KR> kdata = matrix.template small_sample_correction<SE::KR, IM::EIM>();
          M = kdata.vcov_beta;
        } else {
          CorrectionData<SE::KR> kdata = matrix.template small_sample_correction<SE::KR, IM::OIM>();
          M = kdata.vcov_beta;
        }
      
      break;
      }
    case SE::KR2:
    {
      if(im_type == IM::EIM){
        CorrectionData<SE::KR2> kdata = matrix.template small_sample_correction<SE::KR2, IM::EIM>();
        M = kdata.vcov_beta;
      } else {
        CorrectionData<SE::KR2> kdata = matrix.template small_sample_correction<SE::KR2, IM::OIM>();
        M = kdata.vcov_beta;
      }
      break;
    }
    case SE::Robust:
      M = matrix.sandwich_matrix();
      break;
    default:
      {
        if(im_type == IM::EIM){
          M = matrix.information_matrix();
          M = M.llt().solve(MatrixXd::Identity(M.rows(),M.cols()));
        } else {
          MatrixXd Moim = matrix.template observed_information_matrix<IM::OIM>();
          Moim = Moim.llt().solve(MatrixXd::Identity(Moim.rows(),Moim.cols()));
          M = Moim.topLeftCorner(M.rows(),M.cols());
        }
        
        break;
      }
  }
  
#if defined(R_BUILD) && defined(ENABLE_DEBUG)
  Rcpp::Rcout << "\nDone calculating SE, calculating margin...";
#endif
  
  switch(re_type){
  case RandomEffectMargin::At: case RandomEffectMargin::AtEstimated: case RandomEffectMargin::AtZero:
  {
    VectorXd zu(N);
    if(re_type == RandomEffectMargin::At){
      if(single_row){
        zu(0) = atrevals[0];
      } else {
        VectorXd u(model.covariance.Q());
        for(int i = 0; i < model.covariance.Q(); i++)u(i) = atrevals[i];
        zu = model.covariance.Z()*u;
      }
    } else if(re_type == RandomEffectMargin::AtEstimated) {
      zu = re.zu_.rowwise().mean();
    } else if(re_type == RandomEffectMargin::AtZero){
      zu.setZero();
    }
    
    switch(type){
    case MarginType::DyDx:
    {
      double d_result = 0;
      dblvec m_result(2+2*P);
      dblvec delta_vec(P,0.0);
      for(int i = 0; i < newXdata.rows(); i++)mcalc.data(i,xcol) = xvals.first;
//#pragma omp parallel for reduction(+:d_result) reduction(vec_dbl_plus:delta_vec) private(m_result)
      for(int i = 0; i < N; i++){
        newXdata(i,xcol) = xvals.first;
        m_result = mcalc.calculate<CalcDyDx::XBeta>(i,0,xcol,zu(i));
        d_result += m_result[1];
        for(int p = 0; p < P; p++)delta_vec[p] += m_result[p+2+P];
      }
      result.first = d_result/N;
      for(int p = 0; p < P; p++)delta(p) = delta_vec[p]/N;
      result.second = sqrt((delta.transpose()*M*delta)(0));
      break;
    }
    case MarginType::Diff:
      {
      double d_result = 0;
      dblvec delta_vec(P,0.0);
      dblvec m_result(1+P);
      for(int i = 0; i < newXdata.rows(); i++)mcalc.data(i,xcol) = xvals.first;
//#pragma omp parallel for reduction(+:d_result) reduction(vec_dbl_plus:delta_vec) private(m_result)
      for(int i = 0; i < N; i++)
      {
        m_result = mcalc.calculate<CalcDyDx::BetaFirst>(i,0,0,zu(i));
        d_result += m_result[0];
        for(int p = 0; p < P; p++)delta_vec[p] += m_result[p+1];
      }
      
      for(int j = 0; j < newXdata.rows(); j++)mcalc.data(j,xcol) = xvals.second;
      
//#pragma omp parallel for reduction(+:d_result) reduction(vec_dbl_plus:delta_vec) private(m_result)      
      for(int i = 0; i < N; i++)
      {
        m_result = mcalc.calculate<CalcDyDx::BetaFirst>(i,0,0,zu(i));
        d_result += -1.0*m_result[0];
        for(int p = 0; p < P; p++)delta_vec[p] += -1.0*m_result[p+1];
      }
      result.first = d_result/N;
      for(int p = 0; p < P; p++)delta(p) = delta_vec[p]/N;
      result.second = sqrt((delta.transpose()*M*delta)(0));
      break;
      }
    case MarginType::Ratio:
    {
      double d_result0 = 0;
      double d_result1 = 0;
      dblvec delta0(P,0);
      dblvec delta1(P,0);
      dblvec m_result0(P+1);
      dblvec m_result1(P+1);
      for(int i = 0; i < newXdata.rows(); i++) mcalc.data(i,xcol) = xvals.first;
//#pragma omp parallel for private(m_result0) reduction(+:d_result0) reduction(vec_dbl_plus:delta0)  
      for(int i = 0; i < N; i++)
      {
        m_result0 = mcalc.calculate<CalcDyDx::BetaFirst>(i,0,0,zu(i));
        d_result0 += m_result0[0];
        for(int p = 0; p < P; p++) delta0[p] += m_result0[p+1];
      }
      
      for(int i = 0; i < newXdata.rows(); i++) mcalc.data(i,xcol) = xvals.second;
//#pragma omp parallel for private(m_result1) reduction(+:d_result1) reduction(vec_dbl_plus:delta1) 
      for(int i = 0; i < N; i++)
      {
        m_result1 = mcalc.calculate<CalcDyDx::BetaFirst>(i,0,0,zu(i));
        d_result1 += m_result1[0];
        for(int p = 0; p < P; p++) delta1[p] += m_result1[p+1];
      }
      
      result.first = log(d_result0/N) - log(d_result1/N);
      for(int p = 0; p < P; p++)delta(p) = delta0[p]/d_result0 - delta1[p]/d_result1;
      result.second = sqrt((delta.transpose()*M*delta)(0));
      break;
    }
    
    }
    break;
  }
    case RandomEffectMargin::Average:
      {
      int iter = re.zu_.cols();
      switch(type){
      case MarginType::DyDx:
      {
        double d_result = 0;
        dblvec m_result(2+2*P);
        dblvec delta_vec(P,0.0);
        for(int i = 0; i < newXdata.rows(); i++) mcalc.data(i,xcol) = xvals.first;
//#pragma omp parallel for private(m_result) reduction(+:d_result) reduction(vec_dbl_plus:delta_vec) collapse(2)
        for(int i = 0; i < model.n(); i++){
          for(int j = 0; j < iter; j++){
            if(single_row){
              m_result = mcalc.calculate<CalcDyDx::XBeta>(0,0,xcol,re.zu_(i,j));
            } else {
              m_result = mcalc.calculate<CalcDyDx::XBeta>(i,0,xcol,re.zu_(i,j));
            }
            d_result += m_result[1];
            for(int p = 0; p < P; p++)delta_vec[p] += m_result[p+2+P];
          }
        }
        result.first = d_result/(model.n()*iter);
        for(int p = 0; p < P; p++)delta(p) = delta_vec[p];
        delta.array() *= (1.0/(N*iter));
        result.second = sqrt((delta.transpose()*M*delta)(0));
        break;
      }
      case MarginType::Diff:
      {
        double d_result = 0;
        dblvec m_result(P+1);
        dblvec delta_vec(P,0.0);
        for(int i = 0; i < newXdata.rows(); i++)mcalc.data(i,xcol) = xvals.first;
//#pragma omp parallel for private(m_result) reduction(+:d_result) reduction(vec_dbl_plus:delta_vec) collapse(2)
        for(int i = 0; i < model.n(); i++)
        {
          for(int j = 0; j < iter; j++)
          {
            if(single_row){
              m_result = mcalc.calculate<CalcDyDx::BetaFirst>(0,0,0,re.zu_(i,j));
            } else {
              m_result = mcalc.calculate<CalcDyDx::BetaFirst>(i,0,0,re.zu_(i,j));
            }
            d_result += m_result[0];
            for(int p = 0; p < P; p++)delta_vec[p] += m_result[p+1];
           }
        }
        
        for(int i = 0; i < newXdata.rows(); i++)mcalc.data(i,xcol) = xvals.second;
//#pragma omp parallel for private(m_result) reduction(+:d_result) reduction(vec_dbl_plus:delta_vec) collapse(2)
        for(int i = 0; i < model.n(); i++)
        {
          for(int j = 0; j < iter; j++)
          {
            if(single_row){
              m_result = mcalc.calculate<CalcDyDx::BetaFirst>(0,0,0,re.zu_(i,j));
            } else {
              m_result = mcalc.calculate<CalcDyDx::BetaFirst>(i,0,0,re.zu_(i,j));
            }
            d_result -= m_result[0];
            for(int p = 0; p < P; p++)delta_vec[p] -= m_result[p+1];
          }
        }
        
        result.first = d_result/(model.n()*iter);
        for(int p = 0; p < P; p++)delta(p) = delta_vec[p];
        delta.array() *= (1.0/(N*iter));
        result.second = sqrt((delta.transpose()*M*delta)(0));
        break;
      }
      case MarginType::Ratio:
      {
        double d_result0 = 0;
        double d_result1 = 0;
        dblvec delta0(P,0);
        dblvec delta1(P,0);
        dblvec m_result0(1+P);
        dblvec m_result1(1+P);
        for(int i = 0; i < newXdata.rows(); i++)mcalc.data(i,xcol) = xvals.first;
//#pragma omp parallel for private(m_result0) reduction(+:d_result0) reduction(vec_dbl_plus:delta0) collapse(2)
        for(int i = 0; i < model.n(); i++)
        {
          for(int j = 0; j < iter; j++)
          {
            if(N==1){
              m_result0 = mcalc.calculate<CalcDyDx::BetaFirst>(0,0,0,re.zu_(i,j));
            } else {
              m_result0 = mcalc.calculate<CalcDyDx::BetaFirst>(i,0,0,re.zu_(i,j));
            }
            d_result0 += m_result0[0];
            for(int p = 0; p < P; p++) delta0[p] += m_result0[p+1];
          }
        }
        
        for(int i = 0; i < newXdata.rows(); i++)mcalc.data(i,xcol) = xvals.second;
//#pragma omp parallel for private(m_result1) reduction(+:d_result1) reduction(vec_dbl_plus:delta1) collapse(2)
        for(int i = 0; i < model.n(); i++)
        {
          for(int j = 0; j < iter; j++)
          {
            if(N==1){
              m_result1 = mcalc.calculate<CalcDyDx::BetaFirst>(0,0,0,re.zu_(i,j));
            } else {
              m_result1 = mcalc.calculate<CalcDyDx::BetaFirst>(i,0,0,re.zu_(i,j));
            }
            d_result1 += m_result1[0];
            for(int p = 0; p < P; p++) delta1[p] += m_result1[p+1];
          }
        }
        
        result.first = log(d_result0) - log(d_result1);
        for(int p = 0; p < P; p++){
          delta(p) = delta0[p]/d_result0 - delta1[p]/d_result1;
        }
        result.second = sqrt((delta.transpose()*M*delta)(0));
        break;
      }
      }
      break;
      }
  }
  
  return result;
  
}

template<typename modeltype>
inline void glmmr::Model<modeltype>::fit(const int niter, const int max_iter, const bool start_ml_beta, const double tol, const int hist,const int k0)
{
  using std::chrono::high_resolution_clock;
  using std::chrono::duration_cast;
  using std::chrono::duration;
  using std::chrono::milliseconds;
  
  
  bool converged = false;    
  int iter = 1;
  double ll_old = -std::numeric_limits<double>::infinity();
  // start at ml values
  if(start_ml_beta) optim.template ml_beta<BOBYQA>();
  VectorXd beta = Map<VectorXd>(model.linear_predictor.parameters.data(), model.linear_predictor.parameters.size());
  VectorXd theta = Map<VectorXd>(model.covariance.parameters_.data(), model.covariance.parameters_.size());
  if(optim.trace > 0)Rcpp::Rcout << "\nStarting values\nBeta: " << beta.transpose() << "\ntheta: " << theta.transpose() << std::endl;
  while (!converged && iter <= max_iter) {
    if(model.family.family==Fam::gaussian){
      if(optim.trace > 0)Rcpp::Rcout << "\n-------------- ITER: " << iter << " ------------" << std::endl;
      auto t1 = high_resolution_clock::now();
      optim.nr_beta_gaussian();
      auto t2 = high_resolution_clock::now();
      duration<double, std::milli> ms_double = t2 - t1;
      if(optim.trace > 0)Rcpp::Rcout << "TIMING STEP 1 (nr beta): " << ms_double.count() << "ms" << std::endl;
      optim.nr_theta_gaussian();
      auto t3 = high_resolution_clock::now();
      ms_double = t3 - t2;
      if(optim.trace > 0)Rcpp::Rcout << "TIMING STEP 2 (nr theta): " << ms_double.count() << "ms" << std::endl;
      check_for_errors("Gaussian fit");
      double ll_new = optim.marginal_log_likelihood();
      beta = Map<VectorXd>(model.linear_predictor.parameters.data(), model.linear_predictor.parameters.size());
      theta = Map<VectorXd>(model.covariance.parameters_.data(), model.covariance.parameters_.size());
      if(optim.trace > 0) Rcpp::Rcout << "\nBeta: " << beta.transpose() << "\ntheta: " << theta.transpose() << std::endl;
      if constexpr (std::is_same_v<modeltype, bits_ar1>) {
        if(optim.trace > 0) Rcpp::Rcout << "\nrho: " << model.covariance.rho;
      }
      if(optim.trace > 0)Rcpp::Rcout << "Log likelihood (new | previous): " << ll_new << " | " << ll_old << " Diff: " << std::abs(ll_new - ll_old) << std::endl;
      converged = std::abs(ll_new - ll_old) < tol;
      ll_old = ll_new;
      if(optim.trace > 0) if (converged) Rcpp::Rcout << "\nCONVERGED!";
      
    } else {
      if(optim.trace > 0)Rcpp::Rcout << "\n-------------- ITER: " << iter << " ------------" << std::endl;
      auto t1 = high_resolution_clock::now();
      matrix.posterior_u_samples(niter,false,true,false);
      check_for_errors("u step");
      auto t2 = high_resolution_clock::now();
      duration<double, std::milli> ms_double = t2 - t1;
      if(optim.trace > 0)Rcpp::Rcout << "TIMING STEP 1 (posterior u sample): " << ms_double.count() << "ms" << std::endl;
      optim.nr_beta();
      check_for_errors("beta step");
      auto t3 = high_resolution_clock::now();
      ms_double = t3 - t2;
      if(optim.trace > 0)Rcpp::Rcout << "TIMING STEP 2 (nr beta): " << ms_double.count() << "ms" << std::endl;
      optim.nr_theta();
      check_for_errors("theta step");
      auto t4 = high_resolution_clock::now();
      ms_double = t4 - t3;
      if(optim.trace > 0)Rcpp::Rcout << "TIMING STEP 3 (nr theta): " << ms_double.count() << "ms" << std::endl;
      
      beta = Map<VectorXd>(model.linear_predictor.parameters.data(), model.linear_predictor.parameters.size());
      theta = Map<VectorXd>(model.covariance.parameters_.data(), model.covariance.parameters_.size());
      if(optim.trace > 0) Rcpp::Rcout << "\nBeta: " << beta.transpose() << "\ntheta: " << theta.transpose();
      if constexpr (std::is_same_v<modeltype, bits_ar1>) {
        if(optim.trace > 0) Rcpp::Rcout << "\nrho: " << model.covariance.rho;
      }
      if(optim.trace > 0)Rcpp::Rcout << "\nu: " << re.u_.topRows(5).rowwise().mean().transpose();
      if (iter > 2) {
        converged = optim.check_convergence(tol,hist,iter,k0);
        if(optim.trace > 0) if (converged) Rcpp::Rcout << "\nCONVERGED!";
      }
    }
    
    iter++;
  }
}

template<typename modeltype>
void glmmr::Model<modeltype>::check_for_errors(const std::string& caller, bool full_check) {
  bool has_error = false;
  std::stringstream diagnostics;
  
  // ==================== Fast checks (always run) ====================
  // These are O(p) and O(n_theta) - very small
  
  // Beta parameters
  const auto& beta = model.linear_predictor.parameters;
  for(size_t i = 0; i < beta.size(); i++) {
    if(std::isnan(beta[i]) || std::isinf(beta[i])) {
      has_error = true;
      diagnostics << "ERROR: beta[" << i << "] is NaN/Inf: " << beta[i] << std::endl;
    }
  }
  
  // Theta parameters
  const auto& theta = model.covariance.parameters_;
  for(size_t i = 0; i < theta.size(); i++) {
    if(std::isnan(theta[i]) || std::isinf(theta[i])) {
      has_error = true;
      diagnostics << "ERROR: theta[" << i << "] is NaN/Inf: " << theta[i] << std::endl;
    }
  }
  
  // AR1 rho
  if constexpr (std::is_same_v<modeltype, bits_ar1>) {
    if(std::isnan(model.covariance.rho) || std::isinf(model.covariance.rho)) {
      has_error = true;
      diagnostics << "ERROR: rho is NaN/Inf: " << model.covariance.rho << std::endl;
    }
    if(model.covariance.rho <= -1.0 || model.covariance.rho >= 1.0) {
      has_error = true;
      diagnostics << "ERROR: rho out of bounds: " << model.covariance.rho << std::endl;
    }
  }
  
  // u_mean_ - O(Q), usually manageable
  if(re.u_mean_.hasNaN()) {
    has_error = true;
    diagnostics << "ERROR: u_mean_ contains NaN" << std::endl;
  }
  
  // Only do expensive checks if requested or if error already found
  if(full_check || has_error) {
    // ==================== Expensive checks ====================
    
    if(re.u_.size() > 0 && re.u_.hasNaN()) {
      has_error = true;
      diagnostics << "ERROR: u_ contains NaN" << std::endl;
    }
    
    if(re.scaled_u_.size() > 0 && re.scaled_u_.hasNaN()) {
      has_error = true;
      diagnostics << "ERROR: scaled_u_ contains NaN" << std::endl;
    }
    
    if(re.u_solve_.size() > 0 && re.u_solve_.hasNaN()) {
      has_error = true;
      diagnostics << "ERROR: u_solve_ contains NaN" << std::endl;
    }
    
    if(re.u_weight_.size() > 0 && (re.u_weight_.hasNaN() || !re.u_weight_.isFinite().all())) {
      has_error = true;
      diagnostics << "ERROR: u_weight_ contains NaN/Inf" << std::endl;
    }
    
    if(re.u_loglik_.size() > 0 && re.u_loglik_.hasNaN()) {
      has_error = true;
      diagnostics << "ERROR: u_loglik_ contains NaN" << std::endl;
    }
    
    // Data checks
    if(model.data.y.hasNaN()) {
      has_error = true;
      diagnostics << "ERROR: y contains NaN" << std::endl;
    }
    
    if(model.data.offset.hasNaN()) {
      has_error = true;
      diagnostics << "ERROR: offset contains NaN" << std::endl;
    }
  }
  
  // ==================== Print diagnostics if error ====================
  if(has_error) {
    diagnostics << std::endl << "=== CONTEXT (from " << caller << ") ===" << std::endl;
    
    // Dimensions
    diagnostics << "Dimensions: n=" << model.n() 
                << ", p=" << model.linear_predictor.P() 
                << ", Q=" << model.covariance.Q() << std::endl;
    
    // Parameters
    diagnostics << "beta: ";
    for(size_t i = 0; i < beta.size(); i++) diagnostics << beta[i] << " ";
    diagnostics << std::endl;
    
    diagnostics << "theta: ";
    for(size_t i = 0; i < theta.size(); i++) diagnostics << theta[i] << " ";
    diagnostics << std::endl;
    
    if constexpr (std::is_same_v<modeltype, bits_ar1>) {
      diagnostics << "rho: " << model.covariance.rho << std::endl;
    }
    
    // Data ranges
    diagnostics << "y range: [" << model.data.y.minCoeff() << ", " << model.data.y.maxCoeff() << "]" << std::endl;
    diagnostics << "offset range: [" << model.data.offset.minCoeff() << ", " << model.data.offset.maxCoeff() << "]" << std::endl;
    
    // Random effects ranges (only if they exist)
    if(re.u_.size() > 0) {
      diagnostics << "u_ range: [" << re.u_.minCoeff() << ", " << re.u_.maxCoeff() << "]" << std::endl;
    }
    if(re.u_mean_.size() > 0) {
      diagnostics << "u_mean_ range: [" << re.u_mean_.minCoeff() << ", " << re.u_mean_.maxCoeff() << "]" << std::endl;
    }
    if(re.u_weight_.size() > 0) {
      double ess = 1.0 / (re.u_weight_ * re.u_weight_).sum();
      diagnostics << "u_weight_ sum: " << re.u_weight_.sum() << ", ESS: " << ess << std::endl;
    }
    
    Rcpp::Rcout << diagnostics.str();
    R_FlushConsole();
    Rcpp::stop("Numerical error detected. See diagnostics above.");
  }
}
