#pragma once

#include "covariance.hpp"
#include "modelbits.hpp"

namespace glmmr {

using namespace Eigen;

enum class RandomEffectMargin {
  AtEstimated = 0,
    At = 1,
    AtZero = 2,
    Average = 3
};


template<typename modeltype>
class RandomEffects{
public:
  MatrixXd    u_;
  MatrixXd    scaled_u_;
  MatrixXd    zu_;
  VectorXd    u_mean_;
  MatrixXd    u_solve_;
  ArrayXd     u_weight_;
  VectorXd    u_loglik_;
  modeltype&  model;
  int         mcmc_block_size = 1; // for saem
  
  RandomEffects(modeltype& model_) : 
    u_(MatrixXd::Zero(model_.covariance.Q(),1)),
    scaled_u_(MatrixXd::Zero(model_.covariance.Q(),1)),
    zu_(model_.n(),1), u_mean_(VectorXd::Zero(model_.covariance.Q())), 
    u_solve_(MatrixXd::Zero(model_.covariance.Q(),1)),
    u_weight_(VectorXd::Zero(1)), u_loglik_(VectorXd::Zero(1)), model(model_) {};
  
  RandomEffects(modeltype& model_, int n, int Q) : 
    u_(MatrixXd::Zero(Q,1)),
    scaled_u_(MatrixXd::Zero(Q,1)),
    zu_(n,1), u_mean_(Q), u_solve_(Q,1),
    u_weight_(1), u_loglik_(1), model(model_) {};
  
  RandomEffects(const glmmr::RandomEffects<modeltype>& re) : u_(re.u_), scaled_u_(re.scaled_u_), 
    zu_(re.zu_), u_mean_(re.u_mean_), u_solve_(re.u_solve_), u_weight_(re.u_weight_),
    u_loglik_(re.u_loglik_), model(re.model) {};                
  
  MatrixXd      Zu(){return zu_;};
  MatrixXd      u(bool scaled = true);
  
  
  VectorMatrix  predict_re(const ArrayXXd& newdata_,const ArrayXd& newoffset_);
  void          update_zu(const bool weights);
};

}

template<typename modeltype>
inline MatrixXd glmmr::RandomEffects<modeltype>::u(bool scaled){
  if(scaled){
    return model.covariance.Lu(u_);
  } else {
    return u_;
  }
}

template<typename modeltype>
inline void glmmr::RandomEffects<modeltype>::update_zu(const bool weights){
  scaled_u_ = model.covariance.Lu(u_);
  MatrixXd Z = model.covariance.Z();
  zu_ = Z * scaled_u_;
  ArrayXd xb = model.xb();
  u_solve_ = model.covariance.solve(scaled_u_);
  u_weight_.setZero();
  if(weights){
#pragma omp parallel for 
    for(int i = 0; i < scaled_u_.cols(); i++){
      double llmod = maths::log_likelihood(model.data.y.array(),xb + zu_.col(i).array(), model.data.variance,model.family);
      double llprior = -0.5 * scaled_u_.col(i).dot(u_solve_.col(i)); 
      u_weight_(i) = llmod + llprior - u_loglik_(i);
    }
    u_weight_ -= u_weight_.maxCoeff();
    u_weight_ = u_weight_.exp();
    double weightsum = u_weight_.sum();
    u_weight_ *= 1.0/weightsum;
  } else {
    u_weight_.setConstant(1.0/scaled_u_.cols());
  }
}

template<>
inline void glmmr::RandomEffects<bits_hsgp>::update_zu(const bool weights){
  // u_ is M-dimensional u-space spectral coefficients, u ~ N(0, diag(Λ))
  // zu_ = Z * Phi * u (no sqrt(Lambda) scaling)
  zu_ = model.covariance.ZLu(u_);  // ZLu now returns ZPhi * u
  scaled_u_ = u_;
  u_solve_ = model.covariance.solve(u_);  // diag(1/Λ) * u
  ArrayXd xb = model.xb();
  u_weight_.setZero();
  if(weights){
    ArrayXd Lambda = model.covariance.LambdaSPD();
    double half_logdet = 0.5 * Lambda.log().sum();
    
#pragma omp parallel for
    for(int i = 0; i < u_.cols(); i++){
      double llmod = maths::log_likelihood(model.data.y.array(), xb + zu_.col(i).array(),
                                           model.data.variance, model.family);
      // log prior: -½ Σ log(Λ_k) - ½ Σ u_k²/Λ_k
      double llprior = -half_logdet - 0.5 * (u_.col(i).array().square() / Lambda).sum();
      u_weight_(i) = llmod + llprior - u_loglik_(i);
    }
    u_weight_ -= u_weight_.maxCoeff();
    u_weight_ = u_weight_.exp();
    u_weight_ *= 1.0 / u_weight_.sum();
  } else {
    u_weight_.setConstant(1.0 / u_.cols());
  }
}

template<typename modeltype>
inline VectorMatrix glmmr::RandomEffects<modeltype>::predict_re(const ArrayXXd& newdata_,
                                                            const ArrayXd& newoffset_){
  if(model.covariance.data_.cols()!=newdata_.cols())throw std::runtime_error("Different numbers of columns in new data");
  // generate the merged data
  int nnew = newdata_.rows();
  ArrayXXd mergedata(model.n()+nnew,model.covariance.data_.cols());
  mergedata.topRows(model.n()) = model.covariance.data_;
  mergedata.bottomRows(nnew) = newdata_;
  ArrayXd mergeoffset(model.n()+nnew);
  mergeoffset.head(model.n()) = model.data.offset;
  mergeoffset.tail(nnew) = newoffset_;
  
  Covariance covariancenew(model.covariance.form_,
                           mergedata,
                           model.covariance.colnames_);
  
  covariancenew.update_parameters(model.covariance.parameters_);
  // //generate sigma
  int newQ = covariancenew.Q() - model.covariance.Q();
  VectorMatrix result(newQ);
  result.vec.setZero();
  result.mat.setZero();
  MatrixXd D = covariancenew.D(false,false);
  result.mat = D.block(model.covariance.Q(),model.covariance.Q(),newQ,newQ);
  MatrixXd D22 = D.block(0,0,model.covariance.Q(),model.covariance.Q());
  D22 = D22.llt().solve(MatrixXd::Identity(model.covariance.Q(),model.covariance.Q()));
  MatrixXd D12 = D.block(model.covariance.Q(),0,newQ,model.covariance.Q());
  MatrixXd Lu = model.covariance.Lu(u(false));
  MatrixXd SSV = D12 * D22 * Lu;
  result.vec = SSV.rowwise().mean();
  result.mat -= D12 * D22 * D12.transpose();
  return result;
}

template<>
inline VectorMatrix glmmr::RandomEffects<bits_hsgp>::predict_re(const ArrayXXd& newdata_,
                                                                const ArrayXd& newoffset_){
  if(model.covariance.data_.cols() != newdata_.cols())
    throw std::runtime_error("Different numbers of columns in new data");
  
  hsgpCovariance covnew(model.covariance.form_,
                        newdata_,
                        model.covariance.colnames_);
  
  // Copy basis function settings from fitted model
  covnew.update_approx_parameters(model.covariance.m, model.covariance.L_factor_);
  covnew.update_parameters(model.covariance.parameters_);
  
  // Phi_new is n_new × M, evaluated at new locations
  MatrixXd Phi_new = covnew.PhiSPD(false, false);
  
  // Map u-space samples to predictions at new locations: Phi_new * u
  MatrixXd umat = u(false);  // M × n_iter
  MatrixXd pred = Phi_new * umat;  // n_new × n_iter
  
  int n_new = newdata_.rows();
  int iter = pred.cols();
  VectorMatrix result(n_new);
  result.vec = pred.rowwise().mean();
  result.mat.setZero();
  for(int i = 0; i < iter; i++){
    VectorXd diff = pred.col(i) - result.vec;
    result.mat.noalias() += diff * diff.transpose();
  }
  result.mat.array() *= (1.0 / (double)iter);
  return result;
}

