#pragma once

#define _USE_MATH_DEFINES
#include <sstream>
#include "maths.h"
#include "interpreter.h"
#include "formula.hpp"
#include "calculator.hpp"
#include "linearpredictor.hpp"
#include "matrixfield.h"

using namespace Eigen;

namespace glmmr {

inline bool validate_fn(const str& fn){
  bool not_fn = str_to_covfunc.find(fn) == str_to_covfunc.end();
  return not_fn;
}

inline bool is_compact_fn(const CovFunc& fn){
  bool compact = false;
  if(fn == CovFunc::truncpow2 || fn == CovFunc::truncpow3 || fn == CovFunc::truncpow4 || fn == CovFunc::truncpow20 || fn == CovFunc::truncpow30
       || fn == CovFunc::truncpow40 || fn == CovFunc::cauchy || fn == CovFunc::cauchy3 || fn == CovFunc::cauchy0  || fn == CovFunc::cauchy30) compact = true;
  return compact;
}

struct ZNonZero {
  int     col;
  intvec  rows;
  int     xcol;
};

template<typename T, typename ... Vals>
inline bool any_match(T t, Vals ...vals)
{
  return (... || (t == vals));
}



class CovarianceLLT {
public:
  LLT<MatrixXd>                        matL;
  
  CovarianceLLT(const MatrixXd& matD){
    compute(matD);
  }
  
  CovarianceLLT(){};
  
  VectorXd solve(const VectorXd& x) const {
    return matL.solve(x);
  }
  
  MatrixXd solve(const MatrixXd& x) const {
    return matL.solve(x);
  }
  
  void solveInPlace(MatrixXd& x) const {
    matL.solveInPlace(x);
  }
  
  MatrixXd matrixL() const {
    return matL.matrixL();
  }
  
  MatrixXd productR(const MatrixXd& Z) const {
    return Z * matL.matrixL();
  }
  
  MatrixXd productL(const MatrixXd& Z) const {
    return matL.matrixL() * Z;
  }
  
  SparseMatrix<double> productR(const SparseMatrix<double>& Z) const {
    MatrixXd L = MatrixXd(matL.matrixL());
    SparseMatrix<double> LS = L.sparseView();
    return Z * LS;
  }
  
  void compute(const MatrixXd& matD){
    matL.compute(matD);
  }
  
  MatrixXd Ltrisolve(const MatrixXd& x){
    MatrixXd L = matL.matrixL();
    return L.triangularView<Lower>().solve(x);
  }
  
  MatrixXd LtrisolveT(const MatrixXd& x){
    MatrixXd L = matL.matrixL();
    return L.transpose().triangularView<Upper>().solve(x);
  }
};


class Covariance {
public:
  // objects
  glmmr::Formula  form_;
  const ArrayXXd  data_;
  const strvec    colnames_;
  dblvec          parameters_;
  dblvec          other_pars_;
  
  // constructors
  Covariance(const str& formula,const ArrayXXd &data,const strvec& colnames);
  Covariance(const glmmr::Formula& form,const ArrayXXd &data,const strvec& colnames);
  Covariance(const str& formula,const ArrayXXd &data,const strvec& colnames,const dblvec& parameters);
  Covariance(const glmmr::Formula& form,const ArrayXXd &data,const strvec& colnames,const dblvec& parameters);
  Covariance(const str& formula,const ArrayXXd &data,const strvec& colnames,const ArrayXd& parameters);
  Covariance(const glmmr::Formula& form,const ArrayXXd &data,const strvec& colnames,const ArrayXd& parameters);
  Covariance(const glmmr::Covariance& cov);
  
  // functions
  virtual void      update_parameters(const dblvec& parameters);
  virtual void      update_parameters_extern(const dblvec& parameters);
  virtual void      update_parameters(const ArrayXd& parameters);
  virtual int       parse();
  double            get_val(int b, int i, int j) const;
  double            get_val(int i, int j) const;
  virtual MatrixXd  Z();
  virtual MatrixXd  D(bool chol = false, bool upper = false);
  virtual VectorXd  sim_re();
  virtual double    log_likelihood(const VectorXd &u);
  virtual double    log_determinant();
  virtual int       npar() const;
  virtual int       B() const;
  virtual int       Q() const;
  virtual int       max_block_dim() const;
  virtual int       block_dim(int b) const;
  virtual void      make_sparse();
  virtual MatrixXd  ZL();
  virtual MatrixXd  ZLu(const MatrixXd& u);
  virtual MatrixXd  Lu(const MatrixXd& u);
  virtual MatrixXd  Zu(const MatrixXd& u);
  virtual void      set_sparse(bool sparse);
  bool              any_group_re() const;
  bool              all_group_re() const;
  bool              all_log_re() const;
  bool              any_log_re() const;
  intvec            parameter_fn_index() const;
  virtual intvec    re_count() const;
  virtual SparseMatrix<double>    ZL_sparse_new();
  virtual SparseMatrix<double>    Z_sparse();
  strvec            parameter_names();
  virtual void      derivatives(std::vector<MatrixXd>& derivs,int order = 1);
  virtual void      nr_step(const MatrixXd &umat,const MatrixXd &vmat, ArrayXd& logl, ArrayXd& gradients, const ArrayXd& uweight);
  void              linear_predictor_ptr(glmmr::LinearPredictor* ptr);
  MatrixXd          information_matrix();
  virtual MatrixXd  solve(const MatrixXd& u);
  
  // the functions below are deprecated
  // virtual MatrixXd LZWZL(const VectorXd& w);
  // virtual sparse ZL_sparse();
  // virtual MatrixXd  log_gradient(const MatrixXd& u, VectorXd& logl);
  // virtual VectorXd  log_gradient(const MatrixXd& u, double& logl);
  
protected:
  // data
  std::vector<glmmr::calculator>      calc_;
  std::vector<std::vector<CovFunc> >  fn_;
  glmmr::LinearPredictor*             linpred_ptr = nullptr;
  intvec                              re_fn_par_link_;
  intvec                              re_count_;
  intvec                              re_order_;
  intvec                              block_size;
  intvec                              block_nvar;
  intvec3d                            re_cols_data_;
  dblvec3d                            re_temp_data_;
  intvec                              z_;
  int                                 Q_;
  SparseMatrix<double>                matZ;
  MatrixXd                            matD;
  int                                 n_;
  int                                 B_;
  int                                 npars_;
  MatrixXd                            dmat_matrix;
  VectorXd                            zquad;
  bool                                isSparse = true;
public:
  CovarianceLLT                       matL;
  MatrixXd                            infomat_theta;

protected:  
  // functions
  void                            update_parameters_in_calculators();
  MatrixXd                        get_block(int b);
  MatrixXd                        get_chol_block(int b,bool upper = false);
  MatrixXd                        D_builder(int b,bool chol = false,bool upper = false);
  void                            update_ax();
  //void                            L_constructor();
  void                            Z_constructor();
  void                            Z_updater();
  MatrixXd                        D_sparse_builder(bool chol = false, bool upper = false);
  // logical flags
  bool                            sparse_initialised = false;
  bool                            sufficient_sparse = true;
public:
  bool                            z_requires_update = false;
protected:
  std::vector<ZNonZero>           z_nonzero;
};

}

inline glmmr::Covariance::Covariance(const str& formula,
                                     const ArrayXXd &data,
                                     const strvec& colnames) :
  form_(formula), data_(data), colnames_(colnames), Q_(parse()), matZ(data_.rows(),Q_), matD(Q_,Q_),
  dmat_matrix(max_block_dim(),max_block_dim()),
  zquad(max_block_dim()), infomat_theta(MatrixXd::Zero(npar(),npar())) {
  Z_constructor();
};

inline glmmr::Covariance::Covariance(const glmmr::Formula& form,
                                     const ArrayXXd &data,
                                     const strvec& colnames) :
  form_(form), data_(data), colnames_(colnames),
  Q_(parse()),matZ(data_.rows(),Q_),matD(Q_,Q_), dmat_matrix(max_block_dim(),max_block_dim()),
  zquad(max_block_dim()), infomat_theta(MatrixXd::Zero(npar(),npar())) {
  Z_constructor();
};

inline glmmr::Covariance::Covariance(const str& formula,
                                     const ArrayXXd &data,
                                     const strvec& colnames,
                                     const dblvec& parameters) :
  form_(formula), data_(data), colnames_(colnames), parameters_(parameters), 
  Q_(parse()), matZ(data_.rows(),Q_),matD(Q_,Q_), dmat_matrix(max_block_dim(),max_block_dim()),
  zquad(max_block_dim()), infomat_theta(MatrixXd::Zero(npar(),npar())) {
  make_sparse();
  Z_constructor();
};

inline glmmr::Covariance::Covariance(const glmmr::Formula& form,
                                     const ArrayXXd &data,
                                     const strvec& colnames,
                                     const dblvec& parameters) :
  form_(form), data_(data), colnames_(colnames), parameters_(parameters),
  Q_(parse()), matZ(data_.rows(),Q_), matD(Q_,Q_),dmat_matrix(max_block_dim(),max_block_dim()),
  zquad(max_block_dim()), infomat_theta(MatrixXd::Zero(npar(),npar())) {
  make_sparse();
  Z_constructor();
};

inline glmmr::Covariance::Covariance(const str& formula,
                                     const ArrayXXd &data,
                                     const strvec& colnames,
                                     const ArrayXd& parameters) :
  form_(formula), data_(data), colnames_(colnames),
  parameters_(parameters.data(),parameters.data()+parameters.size()),
  Q_(parse()), matZ(data_.rows(),Q_),matD(Q_,Q_),
  dmat_matrix(max_block_dim(),max_block_dim()),
  zquad(max_block_dim()), infomat_theta(MatrixXd::Zero(npar(),npar())) {
  make_sparse();
  Z_constructor();
};

inline glmmr::Covariance::Covariance(const glmmr::Formula& form,
                                     const ArrayXXd &data,
                                     const strvec& colnames,
                                     const ArrayXd& parameters) :
  form_(form), data_(data), colnames_(colnames),
  parameters_(parameters.data(),parameters.data()+parameters.size()),Q_(parse()), matZ(data_.rows(),Q_),
  matD(Q_,Q_),dmat_matrix(max_block_dim(),max_block_dim()),
  zquad(max_block_dim()), infomat_theta(MatrixXd::Zero(npar(),npar())) {
  make_sparse();
  Z_constructor();
};

inline glmmr::Covariance::Covariance(const glmmr::Covariance& cov) : form_(cov.form_), data_(cov.data_),
colnames_(cov.colnames_),
parameters_(cov.parameters_), Q_(parse()), matZ(data_.rows(),Q_),matD(Q_,Q_),
dmat_matrix(max_block_dim(),max_block_dim()),
zquad(max_block_dim()), infomat_theta(MatrixXd::Zero(npar(),npar())) {
  make_sparse();
  Z_constructor();
};

inline int glmmr::Covariance::parse(){
  intvec3d re_cols_;
  strvec2d re_par_names_;
  intvec3d re_pars_;
  // now process each step of the random effect terms
  if(static_cast<int>(colnames_.size())!= data_.cols())throw std::runtime_error("colnames length != data columns");
  
  int nre = form_.re_.size();
  
  for(int i = 0; i < nre; i++){
    strvec fn;
    intvec2d fnvars;
    std::stringstream check1(form_.re_[i]);
    std::string intermediate;
    int iter = 0;
    
    while(getline(check1, intermediate, '*')){
      intvec fnvars1;
      fnvars.push_back(fnvars1);
      std::stringstream check2(intermediate);
      std::string intermediate2;
      getline(check2, intermediate2, '(');
      fn.push_back(intermediate2);
      getline(check2, intermediate2, ')');
      if(intermediate2.find(",") != std::string::npos){
        std::stringstream check3(intermediate2);
        std::string intermediate3;
        while(getline(check3, intermediate3, ',')){
          auto colidx = std::find(colnames_.begin(),colnames_.end(),intermediate3);
          if(colidx == colnames_.end()){
            throw std::runtime_error("variable "+intermediate3+" not in data");
          } else {
            int newidx = colidx - colnames_.begin();
            fnvars[iter].push_back(newidx);
          }
        }
      } else {
        auto colidx = std::find(colnames_.begin(),colnames_.end(),intermediate2);
        if(colidx == colnames_.end()){
          throw std::runtime_error("variable "+intermediate2+" not in data");
        } else {
          int newidx = colidx - colnames_.begin();
          fnvars[iter].push_back(newidx);
        }
      }
      iter++;
    }
    
    // if any of the functions are group, then use block functions
    strvec grfn {"gr", "grlog"};
    //auto idxgr = std::find(fn.begin(),fn.end(),"gr");
    auto idxgr = std::find_first_of(fn.begin(),fn.end(),grfn.begin(),grfn.end());
    dblvec2d groups;
    dblvec vals;
    bool isgr;
    int j,k,idx,zcol;
    if(form_.z_[i].compare("1")==0){
      zcol = -1;
    } else {
      auto idxz = std::find(colnames_.begin(),colnames_.end(),form_.z_[i]);
      if(idxz == colnames_.end()){
        auto idxzpar = std::find(form_.fe_parameter_names_.begin(),form_.fe_parameter_names_.end(),form_.z_[i]);
        if(idxzpar== form_.fe_parameter_names_.end()){
          throw std::runtime_error("z variable "+form_.z_[i]+" not in column or parameter names");
        } else {
          z_requires_update = true;
          zcol = idxzpar - form_.fe_parameter_names_.begin();
          zcol = -1*(zcol+2);
        }
      } else {
        zcol = idxz - colnames_.begin();
      }
    }
    
    if(idxgr!=fn.end()){
      idx = idxgr - fn.begin();
      isgr = true;
      vals.resize(fnvars[idx].size());
      
      for(j = 0; j < data_.rows(); j++){
        for(k = 0; k < fnvars[idx].size(); k++){
          vals[k] = data_(j,fnvars[idx][k]);
        }
        if(std::find(groups.begin(),groups.end(),vals) == groups.end()){
          groups.push_back(vals);
        }
      }
    } else {
      isgr = false;
      vals.push_back(0.0);
      groups.push_back(vals);
    }
    
    intvec allcols;
    for(j = 0; j< fnvars.size();j++){
      for(k = 0; k < fnvars[j].size();k++){
        allcols.push_back(fnvars[j][k]);
      }
    }
    
    int total_vars = allcols.size();
    int gridx;
    dblvec allvals;
    intvec2d newrecols;
    allvals.resize(total_vars);
    newrecols.resize(fnvars.size());
    int currresize = calc_.size();
    calc_.resize(currresize+groups.size());
    re_temp_data_.resize(currresize+groups.size());
    re_cols_.resize(currresize+groups.size());
    re_cols_data_.resize(currresize+groups.size());
    
    int fn_var_counter = 0;
    for(int m = 0; m < fnvars.size(); m++){
      intvec iter_fn_var_index;
      for(int p = 0; p < fnvars[m].size(); p++){
        iter_fn_var_index.push_back(p + fn_var_counter);
      }
      fn_var_counter += fnvars[m].size();
      newrecols[m] = iter_fn_var_index;
    }
    
    for(j = 0; j < groups.size(); j++){
      fn_.emplace_back();
      for(const auto& fnvalue: fn){
        if(glmmr::validate_fn(fnvalue))throw std::runtime_error("Function " + fnvalue + " not valid");
        fn_.back().push_back(str_to_covfunc.at(fnvalue));
      }
      z_.push_back(zcol);
      re_order_.push_back(form_.re_order_[i]);
      re_cols_[currresize + j] = newrecols;
    }
    
    for(k = 0; k < data_.rows(); k++){
      if(isgr){
        for(int m = 0; m < vals.size(); m++) vals[m] = data_(k,fnvars[idx][m]);
        auto gridx2 = std::find(groups.begin(),groups.end(),vals);
        gridx = gridx2 - groups.begin();
      } else {
        gridx = 0;
      }
      for(int m = 0; m<total_vars; m++){
        allvals[m] = data_(k,allcols[m]);
      }
      // this can probably be sped up using an unordered map
      if(std::find(re_temp_data_[gridx + currresize].begin(),
                   re_temp_data_[gridx + currresize].end(),allvals) == re_temp_data_[gridx + currresize].end()){
        re_temp_data_[gridx + currresize].push_back(allvals);
        re_cols_data_[gridx + currresize].push_back(allcols);
      }
    }
  }
  
  // get parameter indexes
  re_pars_.resize(fn_.size());
  re_par_names_.resize(fn_.size());
  bool firsti;
  npars_ = 0;
  int remidx;
  for(int i = 0; i < nre; i++){
    firsti = true;
    for(int j = 0; j < fn_.size(); j++){
      if(re_order_[j]==i){
        if(firsti){
          intvec2d parcount1;
          strvec parnames2;
          str fn_name = "";
          for(int k = 0; k < fn_[j].size(); k++) fn_name += covfunc_to_str.at(fn_[j][k]);
          for(int k = 0; k < fn_[j].size(); k++){
            auto parget = covfunc_to_nvar.find(fn_[j][k]);
            int npars = parget->second;
            intvec parcount2;
            for(int l = 0; l < npars; l++){
              parcount2.push_back(l+npars_);
              parnames2.push_back(fn_name+"."+std::to_string(i)+".("+covfunc_to_str.at(fn_[j][k])+")."+std::to_string(l));
              re_fn_par_link_.push_back(i);
            }
            parcount1.push_back(parcount2);
            npars_ += npars;
          }
          re_pars_[j] = parcount1;
          re_par_names_[j] = parnames2;
          firsti = false;
          remidx = j;
        } else {
          re_pars_[j] = re_pars_[remidx];
          re_par_names_[j] = re_par_names_[remidx];
        }
      }
    }
  }
  
  for(int i =0; i<fn_.size();i++){
    for(int j = 0; j < re_pars_[i].size(); j++){
      for(int k = 0; k < re_pars_[i][j].size(); k++){
        calc_[i].parameter_indexes.push_back(re_pars_[i][j][k]);
      }
    }
  }
  
  //now build the reverse polish notation and add distances
  int nvarfn;
  for(int i =0; i<fn_.size();i++){
    std::vector<Do> fn_instruct;
    intvec fn_par;
    int minvalue = 100;
    int ndata = re_temp_data_[i].size();
    calc_[i].data.conservativeResize(ndata*(ndata-1)/2,fn_[i].size());
    calc_[i].data_size = ndata;
    for(int j = 0; j<fn_[i].size();j++){
      auto min_value_iterator = std::min_element(re_pars_[i][j].begin(),re_pars_[i][j].end());
      if(*min_value_iterator < minvalue) minvalue = *min_value_iterator;
    }
    for(int j = 0; j<fn_[i].size();j++){
      if(fn_[i][j]!=CovFunc::gr && fn_[i][j]!=CovFunc::grlog ){
        nvarfn = re_cols_[i][j].size();
        double dist_val;
        double dist_ab;
        for(int k = 0; k < (ndata-1); k++){
          for(int l = k+1; l < ndata; l++){
            dist_val = 0;
            for(int m = 0; m < nvarfn; m++){
              dist_ab = re_temp_data_[i][k][re_cols_[i][j][m]]-re_temp_data_[i][l][re_cols_[i][j][m]];
              dist_val += dist_ab * dist_ab;
            }
            int idxval = (ndata-1)*k - ((k-1)*k/2) + (l-k-1);
#if defined(R_BUILD) && defined(ENABLE_DEBUG)
            if(idxval > calc_[i].data.rows())Rcpp::stop("idxval out of range, i: "+std::to_string(i)+" j: "+std::to_string(j)+" k: "+std::to_string(k));   
            if(j > calc_[i].data.cols())Rcpp::stop("j out of range of cols");
#endif
            calc_[i].data(idxval,j) = sqrt(dist_val);
          }
        }
      }
      std::vector<Do> B = glmmr::interpret_re(fn_[i][j], calc_[i]);
      intvec re_par_less_min_ = re_pars_[i][j];
      for(int k = 0; k < re_pars_[i][j].size(); k++)re_par_less_min_[k] -= minvalue;
      intvec Bpar = glmmr::interpret_re_par(fn_[i][j],j,re_par_less_min_);
      fn_instruct.insert(fn_instruct.end(),B.begin(),B.end());
      fn_par.insert(fn_par.end(),Bpar.begin(),Bpar.end());
    }
    if(fn_[i].size() > 1){
      for(int j = 0; j < (fn_[i].size()-1); j++){
        fn_instruct.push_back(Do::Multiply);
        calc_[i].push_back_function<Do::Multiply>();
      }
    }
    calc_[i].instructions = fn_instruct;
    calc_[i].indexes = fn_par;
    calc_[i].parameter_names = re_par_names_[i];
    calc_[i].parameter_count = re_par_names_[i].size();
    calc_[i].parameters.resize(re_par_names_[i].size());
  }
  //get the number of random effects
  int Qn = 0;
  block_size.resize(calc_.size());
  block_nvar.resize(calc_.size());
  for(int i = 0; i < calc_.size(); i++){
    Qn += re_temp_data_[i].size();
    block_size[i] = re_temp_data_[i].size();
    block_nvar[i] = re_temp_data_[i][0].size();
  }
  re_count_.resize(form_.re_terms().size());
  std::fill(re_count_.begin(), re_count_.end(), 0);
  for(int i = 0; i < calc_.size(); i++) re_count_[re_order_[i]] += re_temp_data_[i].size();
  B_ = calc_.size();
  n_ = data_.rows();
  return Qn;
}

inline MatrixXd glmmr::Covariance::information_matrix(){
  return infomat_theta;
}

inline void glmmr::Covariance::Z_constructor()
{
  // matZ.n = data_.rows();
  // matZ.m = Q_;
  // matZ.Ap = intvec(data_.rows()+1,0);
  typedef Eigen::Triplet<double> T;
  std::vector<T> tripletList;
  int zcount = 0;
  double insertval;
  for(int i = 0; i < B_; i++){
    dblvec vals(block_nvar[i]);
    dblvec valscomp(block_nvar[i]);
    for(int j = 0; j < block_size[i]; j++){
      ZNonZero nonzero;
      if (z_[i]< -1){
        nonzero.col = zcount;
        int xidx = -1*(z_[i]+2);
        nonzero.xcol = xidx;
      } 
      for(int m = 0; m < block_nvar[i]; m++){
        valscomp[m] = re_temp_data_[i][j][m];
      }
      for(int k = 0; k < data_.rows(); k++){
        for(int m = 0; m < block_nvar[i]; m++){
          vals[m] = data_(k,re_cols_data_[i][j][m]);
        }
        if(valscomp==vals){
          if(z_[i]==-1){
            insertval = 1.0;
          } else if (z_[i]< -1){
            insertval = 999.0;
            nonzero.rows.push_back(k);
          } else {
            insertval = data_(k,z_[i]); 
          }
          //matZ.insert(k,zcount,insertval);
          tripletList.push_back(T(k,zcount,insertval));
        }
      }
      zcount++;
      if (z_[i]< -1) z_nonzero.push_back(nonzero);
    }
  }
  matZ.setFromTriplets(tripletList.begin(),tripletList.end());
  re_temp_data_.clear();
}

inline void glmmr::Covariance::linear_predictor_ptr(glmmr::LinearPredictor* ptr){
  linpred_ptr = ptr;
  if(z_requires_update)Z_updater();
}

inline void glmmr::Covariance::Z_updater(){
  if(z_nonzero.size() > 0)z_requires_update = true;
  if(z_requires_update){
    if(linpred_ptr == nullptr)throw std::runtime_error("Linpred ptr not initialised");
    MatrixXd X = linpred_ptr->X();
    if(z_nonzero.size() == 0)throw std::runtime_error("Non non-zero data");
    for(int i = 0; i < z_nonzero.size(); i++){
      for(int j = 0; j < z_nonzero[i].rows.size(); j++){
        matZ.coeffRef(z_nonzero[i].rows[j],z_nonzero[i].col) = X(z_nonzero[i].rows[j],z_nonzero[i].xcol);
      }
    }
  }
}

inline void glmmr::Covariance::update_parameters_in_calculators(){
  for(int i = 0; i < B_; i++){
    calc_[i].update_parameters(parameters_);
  }
}

inline MatrixXd glmmr::Covariance::D(bool chol, bool upper)
{
  MatrixXd D(Q_,Q_);
  if(isSparse){
    D = D_sparse_builder(chol,upper);
  } else {
    D = D_builder(0,chol,upper);
  }
  return D;
};

inline int glmmr::Covariance::npar() const
{
  return npars_;
};


inline int glmmr::Covariance::B() const
{
  return B_;
}

inline int glmmr::Covariance::Q() const
{
  if(Q_==0)throw std::runtime_error("Random effects not initialised");
  return Q_;
}

inline int glmmr::Covariance::max_block_dim() const
{
  int max = 0;
  for(const auto& b: block_size)if(b > max)max = b;
  return max;
}

inline int glmmr::Covariance::block_dim(int b) const
{
  return block_size[b];//re_data_[b].rows();
};

inline intvec glmmr::Covariance::parameter_fn_index() const
{
  return re_fn_par_link_;
}

inline intvec glmmr::Covariance::re_count() const
{
  return re_count_;
}

inline void glmmr::Covariance::update_parameters(const dblvec& parameters)
{
  if(parameters_.size()==0){
    parameters_.resize(npar());
  }
  
  parameters_ = parameters;
  update_parameters_in_calculators();
  
  if(!sparse_initialised){
    make_sparse();
  } else {
    update_ax();
  }
};

inline void glmmr::Covariance::update_parameters_extern(const dblvec& parameters)
{
  if(static_cast<int>(parameters.size())!=npar())throw std::runtime_error(std::to_string(parameters.size())+" covariance parameters provided, "+std::to_string(npar())+" required");
  if(parameters_.size()==0){
    parameters_.resize(npar());
  }
  
  parameters_ = parameters;
  update_parameters_in_calculators();
  
  if(!sparse_initialised){
    make_sparse();
  } else {
    update_ax();
  }
};

inline void glmmr::Covariance::update_parameters(const ArrayXd& parameters)
{
  if(parameters_.size()==0){
    for(int i = 0; i < parameters.size(); i++){
      parameters_.push_back(parameters(i));
    }
    update_parameters_in_calculators();
  } else if(parameters_.size() == parameters.size()){
    for(int i = 0; i < parameters.size(); i++){
      parameters_[i] = parameters(i);
    }
    update_parameters_in_calculators();
    update_ax();
  } else {
    throw std::runtime_error(std::to_string(parameters.size())+" covariance parameters provided, "+std::to_string(parameters_.size())+" required");
  }
};

inline double glmmr::Covariance::get_val(int b, int i, int j) const
{
  return calc_[b].calculate<CalcDyDx::None>(i,j,0,0)[0];
}

inline double glmmr::Covariance::get_val(int i, int j) const
{
  int b_cumul = 0;
  int b_size;
  int b;
  for(b=0; b< B(); b++){
    b_size = block_dim(b);
    if(i < b_cumul+b_size) break;
    b_cumul += b_size;
  }
  int row = i - b_cumul;
  int col = j - b_cumul;
  return calc_[b].calculate<CalcDyDx::None>(row,col,0,0)[0];
}

inline MatrixXd glmmr::Covariance::get_block(int b)
{
  
#if defined(R_BUILD) && defined(ENABLE_DEBUG)
  if(b > calc_.size()-1)Rcpp::stop("b larger than number of blocks");
  if(parameters_.size()==0)Rcpp::stop("no parameters");
  if(b > B_-1)Rcpp::stop("b is too large");
#endif
  
  int dim = block_dim(b);
  MatrixXd D(dim,dim);
  D.setZero();
  //diagonal
  for(int k = 0; k < dim; k++){
    D(k,k) = get_val(b,k,k);
  }
  if(dim>1){
    for(int i = 0; i < (dim-1); i++){
      for(int j = (i+1); j < dim; j++){
        D(j,i) = get_val(b,j,i);
        D(i,j) = D(j,i);
      }
    }
  }
  return D;
}

inline MatrixXd glmmr::Covariance::Z()
{
  Z_updater();
  return MatrixXd(matZ);
}

inline MatrixXd glmmr::Covariance::get_chol_block(int b,bool upper)
{
  int n = block_dim(b);//re_data_[b].rows();
  std::vector<double> L(n * n, 0.0);
  
  for (int j = 0; j < n; j++) {
    double s = glmmr::algo::inner_sum(&L[j * n], &L[j * n], j);
    L[j * n + j] = sqrt(get_val(b, j, j) - s);
    for (int i = j + 1; i < n; i++) {
      double s = glmmr::algo::inner_sum(&L[j * n], &L[i * n], j);
      L[i * n + j] = (1.0 / L[j * n + j] * (get_val(b, j, i) - s));
    }
  }
  MatrixXd M = Map<MatrixXd>(L.data(), n, n);
  if (upper) {
    return M;
  } else {
    return M.transpose();
  }
}

inline VectorXd glmmr::Covariance::sim_re()
{
  if(parameters_.size()==0)throw std::runtime_error("no parameters, cannot simulate random effects");
#if defined(R_BUILD) && defined(ENABLE_DEBUG)
  Rcpp::Rcout << "\nSim";
#endif
  VectorXd samps(Q_);
  int idx = 0;
  int ndim;
  if(!isSparse){
    for(int i=0; i< B(); i++){
      MatrixXd L = get_chol_block(i);
      ndim = block_dim(i);//re_data_[i].rows();
      VectorXd zz(ndim);      
      std::random_device rd{};
      std::mt19937 gen{ rd() };
      std::normal_distribution d{ 0.0, 1.0 };
      auto random_norm = [&d, &gen] { return d(gen); };
      for (int j = 0; j < zz.size(); j++) zz(j) = random_norm();
      samps.segment(idx,ndim) = L*zz;
      idx += ndim;
    }
  } else {
    VectorXd zz(Q_);
    std::random_device rd{};
    std::mt19937 gen{ rd() };
    std::normal_distribution d{ 0.0, 1.0 };
    auto random_norm = [&d, &gen] { return d(gen); };
    for (int j = 0; j < zz.size(); j++) zz(j) = random_norm();
    samps = matL.productL(zz);
  }  
  return samps;
}

inline MatrixXd glmmr::Covariance::D_builder(int b, bool chol, bool upper)
{
  if (b == B_ - 1) {
    return chol ? get_chol_block(b,upper) : get_block(b);
  }
  else {
    MatrixXd mat1 = chol ? get_chol_block(b,upper) : get_block(b);
    MatrixXd mat2;
    if (b == B_ - 2) {
      mat2 = chol ? get_chol_block(b+1,upper) : get_block(b+1);
    }
    else {
      mat2 = D_builder(b + 1, chol, upper);
    }
    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;
  }
}

inline SparseMatrix<double> glmmr::Covariance::ZL_sparse_new() 
{
  Z_updater();
  return matL.productR(matZ);
}

inline SparseMatrix<double> glmmr::Covariance::Z_sparse() {
  Z_updater();
  return matZ;
}

inline MatrixXd glmmr::Covariance::ZL() {
  SparseMatrix<double> ZL = ZL_sparse_new();
  return MatrixXd(ZL);
}

inline MatrixXd glmmr::Covariance::ZLu(const MatrixXd& u){
  SparseMatrix<double> ZL = ZL_sparse_new();
  return ZL * u;
}

inline MatrixXd glmmr::Covariance::Lu(const MatrixXd& u){
  return matL.productL(u);
}

inline MatrixXd glmmr::Covariance::Zu(const MatrixXd& u){
  return matL.productL(u);
}

inline double glmmr::Covariance::log_likelihood(const VectorXd &u){
  //if(parameters_.size()==0)throw std::runtime_error("no covariance parameters, cannot calculate log likelihood");
  double logdet_val=0.0;
  double loglik_val=0.0;
  int obs_counter=0;
  ArrayXd size_B_array(B_);
  
  static const double LOG2PI = log(2.0 * M_PI);
  
  if(!isSparse)
  {
    int blocksize;
    size_B_array.setZero();
    for(int b=0;b<B_;b++){
      blocksize = block_dim(b);
      if(blocksize==1){
        double var = get_val(b,0,0);
        size_B_array[b] = -0.5*(log(var) + LOG2PI + u(obs_counter)*u(obs_counter)/(var));
      } else {
        zquad.setZero();
        dmat_matrix.block(0,0,blocksize,blocksize).noalias() = get_chol_block(b);
        logdet_val = 0.0;
        for(int i = 0; i < blocksize; i++){
          logdet_val += 2*log(dmat_matrix(i,i));
        }
        zquad.segment(0,blocksize).noalias() = glmmr::algo::forward_sub(dmat_matrix,u.segment(obs_counter,blocksize),blocksize);
        size_B_array[b] = -0.5*(blocksize * LOG2PI + logdet_val + zquad.transpose()*zquad);
      }
      obs_counter += blocksize;
    }
    loglik_val = size_B_array.sum();
    
  } else {
    logdet_val = log_determinant();
    VectorXd v = matL.solve(u);
    double quad = v.dot(u);
    loglik_val = -0.5*(Q_ * LOG2PI + logdet_val + quad);
  }
  return loglik_val;
}

inline double glmmr::Covariance::log_determinant(){
  if(parameters_.size()==0)throw std::runtime_error("no covariance parameters, cannot calculate log determinant");
  int blocksize;
  double logdet_val = 0.0;
  if(!isSparse){
    for(int b=0;b<B_;b++){
      blocksize = block_dim(b);
      dmat_matrix.block(0,0,blocksize,blocksize) = get_chol_block(b);
      for(int i = 0; i < blocksize; i++){
        logdet_val += 2*log(dmat_matrix(i,i));
      }
    }
  } else {
    //for (auto k : spchol.D) logdet_val += log(k);
    MatrixXd L(matL.matrixL());
    for(int i = 0; i < Q_; i++) logdet_val += 2*log(L(i,i));
  }
  
  return logdet_val;
}


inline void glmmr::Covariance::make_sparse(){
  if(parameters_.size()==0)throw std::runtime_error("no covariance parameters, cannot make sparse");
  int dim;
  double val;
  int col_counter=0;
  matD.setZero();
  bool compact_fn;
  int compact_col;
  
  for(int b = 0; b < B(); b++){
    compact_fn = false;
    compact_col = 0;
    for(const auto& fn: fn_[b]){
      if(glmmr::is_compact_fn(fn)){
        compact_fn = true;
        break;
      }
      compact_col++;
    }
    dim = block_dim(b);
    
#pragma omp parallel for 
    for (int idx = 1; idx <= dim * (dim + 1) / 2; idx++) {
      double p = (sqrt(8.0 * idx + 1) - 1) / 2;
      int i = (int)p;
      int j;
      if (i == p) {
        i--;
        j = i;
      }
      else {
        j = idx - i * (i + 1) / 2 - 1;
      }
      val = get_val(b, i, j);
      if (compact_fn && i != j) {
        double dist = calc_[b].get_covariance_data(i, j, compact_col);
        if (dist >= 1)val = 0;
      }
      matD(col_counter + i, col_counter + j) = val;
      if (i != j) matD(col_counter + j, col_counter + i) = val;
    }
    
    col_counter += dim;
  }
  matL.compute(matD);
  sparse_initialised = true;
};


inline void glmmr::Covariance::set_sparse(bool sparse){
  isSparse = sparse;
  if(sparse)make_sparse();
}

inline void glmmr::Covariance::update_ax(){
  make_sparse();
};

inline MatrixXd glmmr::Covariance::D_sparse_builder(bool chol,
                                                    bool upper){
  MatrixXd D = MatrixXd::Zero(Q_,Q_);
  if(!chol){
    D = matD;//sparse_to_dense(spchol.A_,true);
  } else {
    D = matL.matrixL();//sparse_to_dense(matL,false);
  }
  return D;
}

inline bool glmmr::Covariance::any_group_re() const{
  bool gr = false;
  const static std::vector<CovFunc> grfns{CovFunc::gr, CovFunc::grlog};
  for(int i = 0; i < fn_.size(); i++){
    auto idxgr = std::find_first_of(fn_[i].begin(), fn_[i].end(),grfns.begin(),grfns.end());
    if(idxgr != fn_[i].end()){
      gr = true;
      break;
    }
    if(gr)break;
  }
  return gr;
}

inline bool glmmr::Covariance::all_group_re() const {
  bool gr = true;
  for (int i = 0; i < fn_.size(); i++) {
    for (int j = 0; j < fn_[i].size(); j++) {
      bool isgr = any_match(fn_[i][j], CovFunc::gr, CovFunc::grlog);
      gr = gr && isgr;
    }
  }
  return gr;
}

inline bool glmmr::Covariance::all_log_re() const {
  bool logre = true;
  for (int i = 0; i < fn_.size(); i++) {
    for (int j = 0; j < fn_[i].size(); j++) {
      bool anylog = any_match(fn_[i][j], CovFunc::grlog, CovFunc::arlog, CovFunc::ar0log, CovFunc::fexplog);// (fn_[i][j] == CovFunc::grlog || fn_[i][j] == CovFunc::arlog || fn_[i][j] == CovFunc::fexplog || fn_[i][j] == CovFunc::ar0log);
      logre = logre && anylog;
    }
  }
  return logre;
}

inline bool glmmr::Covariance::any_log_re() const{
  bool gr = false;
  const static std::vector<CovFunc> logfns {CovFunc::grlog, CovFunc::arlog, CovFunc::fexplog, CovFunc::ar0log};
  
  for(int i = 0; i < fn_.size(); i++){
    auto idxgr = std::find_first_of(fn_[i].begin(), fn_[i].end(),logfns.begin(),logfns.end());
    if(idxgr != fn_[i].end()){
      gr = true;
      break;
    }
    if(gr)break;
  }
  return gr;
}


inline void glmmr::Covariance::nr_step(const MatrixXd &umat, const MatrixXd &vmat, ArrayXd& logl, ArrayXd& gradients, 
                                       const ArrayXd& uweight){
  static const double LOG_2PI = log(2*M_PI);
  static const double NEG_HALF_LOG_2PI = -0.5 * LOG_2PI;
  std::vector<MatrixXd> derivs;
  derivatives(derivs, 1);
  const int npars = derivs.size() - 1;
  const int niter = umat.cols();
  VectorXd grad(npars);
  grad.setZero();
  double logdet_val = 0.0;
  dblvec dqf(npars, 0.0);
  logl.setZero();
  
  if(!isSparse)make_sparse();
  logdet_val = log_determinant();
  logl.array() += NEG_HALF_LOG_2PI * Q_ - 0.5 * logdet_val;
  std::vector<MatrixXd> S;
  std::vector<MatrixXd> BD;
  for(int i = 0; i < npars; i++)
  {
    S.emplace_back(derivs[i + 1].rows(), derivs[i + 1].cols());
    S[i] = matL.solve(derivs[i+1]);
    grad(i) = -0.5 * S[i].trace();
  }
  
  
  MatrixXd M = MatrixXd::Zero(npars, npars);
  

#pragma omp parallel
{
  VectorXd dqf_thread = VectorXd::Zero(npars);
  MatrixXd Sprod(S[0].rows(), S[0].cols());
#pragma omp for 
  for (int i = 0; i < niter; i++)
  {
    double qf = vmat.col(i).dot(umat.col(i));
    logl(i) += -0.5 * qf;
  }
}

for (int j = 0; j < npars; j++){
  double grad_j = 0.0;
#pragma omp parallel for reduction(+:grad_j)
  for (int i = 0; i < niter; i++)
  {
    grad_j += 0.5 * uweight(i) * (umat.col(i).dot(S[j] * vmat.col(i)));
  }
  grad(j) += grad_j;
  
  for(int k = j; k < npars; k++){
    MatrixXd Sprod = S[j] * S[k];
    M(j,k) += -0.5 * Sprod.trace();
    if(j != k) M(k,j) = M(j,k);
    double m_jk = 0.0;
#pragma omp parallel for reduction(+:m_jk)
    for (int i = 0; i < niter; i++)
    {
      m_jk += uweight(i) * (umat.col(i).dot(Sprod * vmat.col(i)));
    }
    M(j,k) += m_jk;
    if(j != k) M(k,j) += m_jk;
  }
}

  gradients.tail(grad.size()) = grad;
  infomat_theta = M;
  VectorXd theta_curr = Map<VectorXd>(parameters_.data(), parameters_.size());
  theta_curr +=  M.llt().solve(grad);
  update_parameters(theta_curr.array());
}

inline void glmmr::Covariance::derivatives(std::vector<MatrixXd>& derivs,
                                           int order){
  // get unique parameters
  strvec pars = parameter_names();
  int R = pars.size();
  int matrix_n = order==2 ? R + R*(R+1)/2 + 1 : R+1;
  // initialise all the matrices to zero
  for(int i = 0; i < matrix_n; i++)derivs.push_back(MatrixXd::Zero(Q_,Q_));
  int block_count = 0;
  // iterate over the blocks and insert if the parameter is in the list.
  for(int b = 0; b < B_; b++){
    int block_dimension = block_dim(b);
    int R_block = calc_[b].parameter_names.size();
    intvec par_index;
    for(int k = 0; k < R_block; k++){
      auto par_pos = std::find(pars.begin(),pars.end(),calc_[b].parameter_names[k]);
      int par_pos_int = par_pos - pars.begin();
      par_index.push_back(par_pos_int);
    }
#if defined(R_BUILD) && defined(ENABLE_DEBUG)
    Rcpp::Rcout << "\nSIGMA DERIVATIVES\nBlock " << b << " with " << R_block << " parameters, " << block_dimension << " size";
    Rcpp::Rcout << "\nPar indexes: ";
    for(const auto& i: par_index)Rcpp::Rcout << i << " ";
#endif
#pragma omp parallel for 
    for (int idx = 1; idx <= block_dimension * (block_dimension + 1) / 2; idx++) {
      dblvec out(matrix_n);
      double p = (sqrt(8.0 * idx + 1) - 1) / 2;
      int i = (int)p;
      int j;
      if (i == p) {
        i--;
        j = i;
      }
      else {
        j = idx - i * (i + 1) / 2 - 1;            
      }
      if(order == 1){
        out = calc_[b].calculate<CalcDyDx::BetaFirst>(i,j,0,0);
      } else {
        out = calc_[b].calculate<CalcDyDx::BetaSecond>(i,j,0,0);
      }
      derivs[0](block_count+i,block_count+j) = out[0];
      if(i!=j)derivs[0](block_count+j,block_count+i) = out[0];
      int index_count = R_block + 1;
      for(int k = 0; k < R_block; k++){
        derivs[par_index[k]+1](block_count+i,block_count+j) = out[k+1];
        if(i!=j)derivs[par_index[k]+1](block_count+j,block_count+i) = out[k+1];
        //second order derivatives
        if(order >= 2){
          for(int l=k; l < R_block; l++){
            int second_pos = par_index[l]*(R-1) - par_index[l]*(par_index[l]-1)/2 + par_index[k];
            derivs[R+1+second_pos](block_count+i,block_count+j) = out[index_count];
            if(i!=j)derivs[R+1+second_pos](block_count+j,block_count+i) = out[index_count];
            index_count++;
          }
        }
      }
      
    }
    block_count += block_dimension;
  }
}

inline strvec glmmr::Covariance::parameter_names(){
  strvec parnames;
  for(int i = 0; i < form_.re_.size(); i++){
    for(int j = 0; j < B_; j++){
      if(re_order_[j]==i){
        parnames.insert(parnames.end(),calc_[j].parameter_names.begin(),calc_[j].parameter_names.end());
        break;
      }
    }
  }
  return parnames;
};

inline MatrixXd glmmr::Covariance::solve(const MatrixXd& u){
  return matL.solve(u);
}

// deprecated functions

// MatrixXd glmmr::Covariance::LZWZL(const VectorXd& w) {
//   return MatrixXd();  // or MatrixXd::Zero(0,0)
// }
// 
// MatrixXd  glmmr::Covariance::log_gradient(const MatrixXd& u, VectorXd& logl) {
//   return MatrixXd();
// }
// 
// VectorXd  glmmr::Covariance::log_gradient(const MatrixXd& u, double& logl) {
//   return VectorXd();
// }
// 
// sparse glmmr::Covariance::ZL_sparse(){
//   sparse ZL;
//   return ZL;
// }
