#pragma once

#include "modelbits.hpp"

namespace glmmr {

using namespace Eigen;

template<typename modeltype>
class MatrixW{
public:
  bool        attenuated = false;
  VectorXd    W_ = VectorXd::Constant(1,1.0);
  modeltype&  model;
  MatrixW(modeltype& model_): model(model_) { update(); };
  VectorXd    W() const;
  void        update();
};
}

template<typename modeltype>
inline VectorXd glmmr::MatrixW<modeltype>::W() const{
  return W_;
}

template<typename modeltype>
inline void glmmr::MatrixW<modeltype>::update(){
  if(W_.size() != model.n())W_.conservativeResize(model.n());
  ArrayXd nvar_par(model.n());
  VectorXd xb(model.n());
  switch(model.family.family){
  case Fam::gaussian: 
    nvar_par = model.data.variance;
    break;
  case Fam::gamma: 
    nvar_par = model.data.variance.inverse();
    break;
  case Fam::beta:
    nvar_par = (1+model.data.variance);
    break;
  case Fam::binomial:
    nvar_par = model.data.variance.inverse();
    break;
  case Fam::quantile: case Fam::quantile_scaled:
  {
    double qvar = (1 - 2*model.family.quantile + 2*model.family.quantile*model.family.quantile)/(model.family.quantile*model.family.quantile*(1-model.family.quantile)*(1-model.family.quantile));
    if(model.family.family == Fam::quantile_scaled) qvar *= model.data.var_par*model.data.var_par;
    nvar_par.setConstant(qvar);
    break;
  }
  default:
    nvar_par.setConstant(1.0);
    break;
  }
  
  if(attenuated){
    xb = (glmmr::maths::attenuted_xb(model.xb(), model.covariance.Z(), model.covariance.D(), model.family.link)).matrix();
  } else {
    xb = model.xb().matrix();
  }
  W_ = (glmmr::maths::dhdmu(xb, model.family)).array();
  W_ = (W_.array()*nvar_par).matrix();
  W_ = ((W_.array().inverse()) * model.data.weights).matrix();
}
