// [[Rcpp::depends(RcppArmadillo)]]
// for matrices

// for vector
#include <vector>

#include <RcppArmadillo.h>
#include <Rcpp.h>
#include <RcppArmadilloExtensions/sample.h>
#include <stdio.h>
#include <string>
#include <iostream>
#include <random>

using namespace Rcpp;

// [[Rcpp::plugins(cpp11)]]

//======================================================//

//////////////////////////////////////////////////////////
//              Stchastic Approximationes               //
//////////////////////////////////////////////////////////


//randomly pick out a integer from an uniform distribution
int rand_uniform_int(int min, int max) {
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<int> dis(min, max);

    return dis(gen);
}

// log model for simple SA
inline arma::mat LogmodelSimplecpp(arma::mat m_mat, const arma::mat& a_mat, const arma::mat& b_mat){
  m_mat *= b_mat;
  const arma::mat& res_mat = 1/(exp(-a_mat - m_mat.each_row())+1);
  
  return res_mat;
}

// log model for full SA
inline arma::mat Logmodelcpp(arma::mat m,  const arma::mat& alpha, const arma::mat& beta){
  
  m *= beta;
  const arma::mat& res = 1/(exp(-alpha - m)+1);
  
  return res;
}


// One step "Heat Bath Sampling"
arma::mat OnestepHBScpp(arma::mat y, const arma::mat& lm) {
  
  // Preparation // 
  const int& itr1=y.n_cols;
  const int& itr2=y.n_rows;
  
  for(int r=0; r < itr2; ++r){
    
    // Initial state //
    int uni=R::runif(0, itr1);
    //int uni=rand_uniform_int(0, itr1-1);
    double ne=lm(r, uni);
    
    double randomNum=arma::randu<double>();
    
    // adjacency state//
    if(ne > randomNum){
      y(r, uni) = 1;
    }else{
      y(r, uni) = 0;
    }
  
  }
  return y;
}

// Log Prior function
inline arma::mat Logprior(const arma::mat& alpha, const double& x) {
  
  arma::mat y = arma::zeros(alpha.n_rows, alpha.n_cols);
  y.fill(x);
  arma::mat lp = -arma::tanh((alpha/y)/2)/y;
  return lp;
}

//////////////////////////////////////////////////////////

//========  only Inplicit environmetanl data ===========//




// [[Rcpp::export]]
arma::mat SA(const arma::mat& ocData, const double& we=0.001,const int& totalit=1000, 
             bool runadamW=true, bool sparse=true, const double& lambda=0.01,
             const int& intv=100, const double& maxlearningrate = 0.005, const double& warmup_tt = 1){
  // ========================================= //
  // -- Initial parameter setting

  const double& nlocation = ocData.n_rows;
  const double& nspecies = ocData.n_cols;
  const double& nutri = 1 / ( nspecies * (nspecies + 1) * 0.5);
  //double learningrate0 = 0.0001;
  double learningrate0 = maxlearningrate;

  arma::mat ystats=arma::trans(ocData);
  ystats *= ocData;
  arma::mat ysim = arma::zeros(nlocation, nspecies);
  arma::mat alphas = arma::zeros(1, nspecies);
  arma::mat beta = arma::zeros(nspecies, nspecies);
  arma::mat delalphas=arma::zeros(1, nspecies);
  arma::mat delbeta = arma::zeros(nspecies, nspecies);
  arma::mat betagrad;
  arma::mat alphasgrad;
  arma::mat logmat; 
  arma::mat ydif;
  arma::mat grad;
  arma::mat paramsmon;
  
  if (lambda < 0.00000001){
    sparse = false;
  }

  const double& beta1 = 0.8;
  const double& beta2 = 0.995;
  int total_tt = 0;
  double beta1t = 1;
  double beta2t = 1;
  double minu = std::numeric_limits<double>::infinity();
  double uu = 0;
  int totalrows = nspecies*std::floor(totalit/intv);
  int updidx = 0;

  arma::mat vt = arma::zeros(nspecies,nspecies);
  arma::mat st = arma::zeros(nspecies,nspecies);
  arma::mat wt = arma::zeros(nspecies,nspecies);
  arma::mat qi = arma::zeros(nspecies,nspecies);
  arma::mat vth, sth, pos_wt, neg_wt, zz;
  
  
  //arma::mat betaconst=(nlocation * arma::abs(arma::eye(nspecies, nspecies)-1));
  arma::mat betaconst= (arma::mat(nspecies, nspecies).fill(nlocation));
  arma::rowvec asconst=(arma::mat(1, nspecies).fill(nlocation));
  std::vector<double> upds(totalrows);
  std::vector<double> points(totalrows);

  double upd;
  const double& momtm=0.3; //
  //double upcrit;
  // ========================================= //
  // Main part
  for(int tt=0; tt < totalit; ++tt){
    // -- Preference for co-occurence of species i and j
    logmat=LogmodelSimplecpp(ysim, alphas, beta);
    ysim=OnestepHBScpp(ysim, logmat);
    ydif = ystats - (arma::trans(ysim) * ysim);
    if(runadamW){
      // -- Update learningrate
      const double& learningrate = learningrate0 + (maxlearningrate - learningrate0)*std::min(tt/warmup_tt,1.0);
      
      // cosine annealing
      //if (tt > warmup_tt){
      //  const double& learningrate = learningrate0 + 0.5*(maxlearningrate - learningrate0)*(1 + std::cos(arma::datum::pi*(tt - warmup_tt)/totalit)); 
      //}
      grad = ydif / betaconst;
      vt = beta1*vt + (1 - beta1)*grad;
      st = beta2*st + (1 - beta2)*(arma::square(grad));
      beta1t = beta1t*beta1; 
      beta2t = beta2t*beta2;
      vth = vt*(1/(1 - beta1t));
      sth = st*(1/(1 - beta2t));
      wt = wt + learningrate*((vth/(arma::sqrt(sth) + 0.00000001))-we*wt);
      if(sparse){
        //uu += learningrate*10*lambda/nspecies;
        // Takano changed 2025/5/29
        uu += learningrate*lambda*nspecies;
        zz = wt;
        pos_wt = wt - (uu + qi);
        pos_wt.elem(arma::find(pos_wt < 0.0)).zeros();
        neg_wt = wt + (uu - qi);
        neg_wt.elem(arma::find(neg_wt > 0.0)).zeros();
        wt.elem(arma::find(zz < 0.0)) = neg_wt.elem(arma::find(zz < 0.0));
        wt.elem(arma::find(zz > 0.0)) = pos_wt.elem(arma::find(zz > 0.0));
        wt.diag() = zz.diag();
        qi += wt - zz;
      }
      alphas = arma::diagvec(wt).t();
      beta = wt;
      beta.diag().fill(0);
      //std::cout << "alpha " << arma::mean(arma::mean(alphas)) << std::endl;

    }else{
      // -- Beta gradient
       const double& learningrate = learningrate0 + (maxlearningrate - learningrate0)*std::min(tt/warmup_tt,1.0);
      
      // cosine annealing
      //if (tt > warmup_tt){
      //  const double& learningrate = learningrate0 + 0.5*(maxlearningrate - learningrate0)*(1 + std::cos(arma::datum::pi*(tt - warmup_tt)/totalit)); 
      //}
      //const double& learningrate=learningrate0;
      betagrad = ydif / betaconst;
      betagrad.diag().fill(0);
      // -- Alpha gradient
      alphasgrad = ydif.diag().t()/asconst;
      // -- delta
      betagrad %= arma::mat(nspecies, nspecies).fill((1-momtm)*learningrate);
      delbeta %= arma::mat(nspecies, nspecies).fill(momtm);
      delbeta += betagrad;
      alphasgrad %= arma::mat(1, nspecies).fill((1-momtm)*learningrate);
      delalphas %= arma::mat(1, nspecies).fill(momtm);
      delalphas += alphasgrad ;
      beta+=delbeta;
      alphas+=delalphas;
      if(sparse){
        //uu += learningrate*lambda/nspecies;
        // Takano changed 2025/5/29
        uu += learningrate*lambda*nspecies/10;
        zz = beta;
        pos_wt = beta - (uu + qi);
        pos_wt.elem(arma::find(pos_wt < 0.0)).zeros();
        neg_wt = beta + (uu - qi);
        neg_wt.elem(arma::find(neg_wt > 0.0)).zeros();
        beta.elem(arma::find(zz < 0.0)) = neg_wt.elem(arma::find(zz < 0.0));
        beta.elem(arma::find(zz > 0.0)) = pos_wt.elem(arma::find(zz > 0.0));
        qi += beta - zz;
      }
    }
    // -- save intermediate status
    if (tt == 0 || rand_uniform_int(0,totalit) < totalrows*1.05){
      if (updidx < totalrows){
        upds[updidx] = arma::sum(arma::vectorise(arma::abs(arma::trimatu(ydif)))) / nlocation * nutri;
        points[updidx] = tt;
        updidx += 1;
      }
    }

    // Fixed 2025/10/6
     if ((tt+1)%intv == 0){
      //std::cout << "tt: " << tt << std::endl;
      paramsmon = arma::join_cols(paramsmon,arma::join_rows(alphas.t(),beta.t()));
    }
    }
    
    // Add (2025/5/6)
    int samplepoints = paramsmon.n_rows;
    upds.resize(samplepoints);
    points.resize(samplepoints);
    
    upds[samplepoints-1] = arma::sum(arma::vectorise(arma::abs(arma::trimatu(ydif)))) / nlocation * nutri;
    points[samplepoints-1] = totalit-1;

    // ========================================= //
    arma::colvec upd_arma = arma::conv_to< arma::colvec >::from(upds);
    arma::colvec point_arma = arma::conv_to< arma::colvec >::from(points);
    arma::mat ystatsList=arma::join_rows(paramsmon,upd_arma,point_arma);
    
    return ystatsList;
}


//////////////////////////////////////////////////////////

//========  including Explicit environmetanl data ======//

// [[Rcpp::export]]
arma::mat fullSA(const arma::mat& ocData, const arma::mat& envData,
                 bool runadamW=true, bool sparse=true, 
                 const double& we=0.001, const int& totalit=1000,const double& lambda=0.01, 
                 const int& intv=100, const double& maxlearningrate = 0.005, const double& warmup_tt = 1){
                                 
  // ========================================= //
  // -- Initial parameter setting
  
  //double learningrate0 = maxlearningrate;
  double learningrate0 = maxlearningrate;
  //const double& momentum=0.3;
  const double& nlocation = ocData.n_rows;
  const double& nspecies = ocData.n_cols;
  const double& nenvironment = envData.n_cols;
  const double& nutri = 1 / ( nspecies * (nspecies + 1) * 0.5);

  //bool runadamW=true;
  //bool sparse=true;
  if (lambda < 0.00000001){
    sparse = false;
  }

  // common parameters and variables
  arma::mat ystats = arma::trans(ocData) * ocData;
  arma::mat yenvstats = arma::trans(envData) * ocData;
  arma::mat ysim = arma::zeros(nlocation, nspecies);

  // Parameters and variables for normal SA
  //const arma::mat& lp = (arma::sum(ocData, 0)+1)/(nlocation+2);
  //arma::mat alphas = arma::log(lp/(1-lp));
  arma::mat alphas = arma::zeros(1,nspecies);
  arma::mat beta = arma::zeros(nspecies, nspecies);
  arma::mat alphae = arma::zeros(nenvironment, nspecies);
  
  arma::mat delalphas= arma::zeros(1, nspecies); 
  arma::mat delalphae = arma::zeros(nenvironment, nspecies); 
  arma::mat delbeta = arma::zeros(nspecies, nspecies);
  
  arma::mat betagrad = arma::zeros(nspecies, nspecies);
  arma::mat alphasgrad = arma::zeros(1, nspecies);
  arma::mat alphaegrad = arma::zeros(nenvironment, nspecies);
  arma::mat alpha; arma::mat logmat; arma::mat ydif; arma::mat yenvdiff;
  
  arma::mat betaconst=(arma::mat(nspecies, nspecies).fill(nlocation));
  arma::rowvec asconst=(arma::mat(1, nspecies).fill(nlocation));
  arma::mat aeconst=(arma::mat(nenvironment, nspecies).fill(nlocation));
  
  // AdamW parameters and variables (occurence matrix)
  const double& obeta1 = 0.8; const double& obeta2 = 0.995;
  double obeta1t = 1; double obeta2t = 1;
  arma::mat ovt = arma::zeros(nspecies,nspecies);
  arma::mat ost = arma::zeros(nspecies,nspecies);
  arma::mat owt = arma::zeros(nspecies,nspecies);
  arma::mat ograd, ovth, osth;
  
  // AdamW parameters and variables (environmental matrix)
  const double& ebeta1 = 0.8; const double& ebeta2 = 0.995;
  double ebeta1t = 1; double ebeta2t = 1;
  arma::mat evt = arma::zeros(nenvironment,nspecies);
  arma::mat est = arma::zeros(nenvironment,nspecies);
  arma::mat ewt = arma::zeros(nenvironment,nspecies);
  arma::mat egrad,evth,esth;

  // Parameters for sparse 
  double minu = std::numeric_limits<double>::infinity();
  double uu = 0; 
  arma::mat qi = arma::zeros(nspecies,nspecies);
  arma::mat pos_owt, neg_owt, zz;
  
  // for storing intermediates
  int totalrows = nspecies*std::floor(totalit/intv);
  std::vector<double> upds(totalrows);
  std::vector<double> points(totalrows);
  int updidx = 0; 
  arma::mat paramsmon;

  const double& momtm=0.3; //
  // ========================================= //
  // Main part
  
  for(int tt=0; tt < totalit; ++tt){
    // -- Preference for co-occurence of species i and j
    alpha=arma::mat(envData * alphae).each_row() + alphas;
    
    // -- 
    logmat=Logmodelcpp(ysim, alpha, beta);
    ysim=OnestepHBScpp(ysim, logmat);

    ydif = ystats - (arma::trans(ysim) * ysim);
    yenvdiff= yenvstats-(arma::trans(envData) * ysim);

    if(runadamW){
      // h and J (occurence matrix)
      const double& learningrate = learningrate0 + (maxlearningrate - learningrate0)*std::min(tt/warmup_tt,1.0);
      ograd = ydif / betaconst;
      ovt = obeta1*ovt + (1 - obeta1)*ograd;
      ost = obeta2*ost + (1 - obeta2)*(arma::square(ograd));
      obeta1t = obeta1t*obeta1; 
      obeta2t = obeta2t*obeta2;
      ovth = ovt*(1/(1 - obeta1t));
      osth = ost*(1/(1 - obeta2t));
      owt = owt + learningrate*((ovth/(arma::sqrt(osth) + 0.00000001))-we*owt);
      // e (environmental matrix)
      egrad = yenvdiff / aeconst;
      evt = ebeta1*evt + (1 - ebeta1)*egrad;
      est = ebeta2*est + (1 - ebeta2)*(arma::square(egrad));
      ebeta1t = ebeta1t*ebeta1; 
      ebeta2t = ebeta2t*ebeta2;
      evth = evt*(1/(1 - ebeta1t));
      esth = est*(1/(1 - ebeta2t));
      ewt = ewt + learningrate*((evth/(arma::sqrt(esth) + 0.00000001)));

      if(sparse){
        //uu += learningrate*10*lambda/nspecies;
        // Takano changed 2025/5/29
        uu += learningrate*lambda*nspecies;
        zz = owt;
        pos_owt = owt - (uu + qi);
        pos_owt.elem(arma::find(pos_owt < 0.0)).zeros();
        neg_owt = owt + (uu - qi);
        neg_owt.elem(arma::find(neg_owt > 0.0)).zeros();
        owt.elem(arma::find(zz < 0.0)) = neg_owt.elem(arma::find(zz < 0.0));
        owt.elem(arma::find(zz > 0.0)) = pos_owt.elem(arma::find(zz > 0.0));
        owt.diag() = zz.diag();
        qi += owt - zz;
      }
      alphas = arma::diagvec(owt).t();
      beta = owt;
      beta.diag().fill(0);
      alphae = ewt;

    }else{
      const double& learningrate = learningrate0 + (maxlearningrate - learningrate0)*std::min(tt/warmup_tt,1.0);
      // -- Beta gradient 
      betagrad = ydif / betaconst;
      betagrad.diag().fill(0);
      
      // -- Alpha gradient
      alphasgrad = (ydif.diag().t())/asconst;
      alphaegrad = yenvdiff/aeconst;

      // -- delta beta
      betagrad %= arma::mat(nspecies, nspecies).fill((1-momtm)*learningrate);
      delbeta %= arma::mat(nspecies, nspecies).fill(momtm);
      delbeta += betagrad;
      
      // -- delta alphas
      alphasgrad %= arma::mat(1, nspecies).fill((1-momtm)*learningrate) ;
      delalphas %= arma::mat(1, nspecies).fill(momtm);
      delalphas+=alphasgrad ;   
      
      // -- delta alphae
      alphaegrad %= arma::mat(nenvironment, nspecies).fill((1-momtm)*learningrate);
      delalphae %= arma::mat(nenvironment, nspecies).fill(momtm);
      delalphae += alphaegrad;
      
      beta+=delbeta; 
      alphas+=delalphas; 
      alphae+=delalphae;
      if(sparse){
        //uu += learningrate*lambda/nspecies;
        // Takano changed 2025/5/29
        uu += learningrate*lambda*nspecies/10;
        zz = beta;
        pos_owt = beta - (uu + qi);
        pos_owt.elem(arma::find(pos_owt < 0.0)).zeros();
        neg_owt = beta + (uu - qi);
        neg_owt.elem(arma::find(neg_owt > 0.0)).zeros();
        beta.elem(arma::find(zz < 0.0)) = neg_owt.elem(arma::find(zz < 0.0));
        beta.elem(arma::find(zz > 0.0)) = pos_owt.elem(arma::find(zz > 0.0));
        qi += beta - zz;
      }
    }

    // -- save intermediate status
    if (tt == 0 || rand_uniform_int(0,totalit) < totalrows*1.05){
      if (updidx < totalrows){
        upds[updidx] = arma::sum(arma::vectorise(arma::abs(arma::trimatu(ydif)))) / nlocation * nutri;
        points[updidx] = tt;
        updidx += 1;
      }
    }

    // Fixed: 2025/10/6
    if ((tt+1)%intv == 0){
      //std::cout << "tt: " << tt << std::endl;
      paramsmon = arma::join_cols(paramsmon,arma::join_rows(alphas.t(),alphae.t(),beta.t()));
      
    }
    
  }
  // -- summarize results 
  // Add (2025/5/6)
  int samplepoints = paramsmon.n_rows;
  upds.resize(samplepoints);
  points.resize(samplepoints);


  upds[samplepoints-1] = arma::sum(arma::vectorise(arma::abs(arma::trimatu(ydif)))) / nlocation * nutri;
  points[samplepoints-1] = totalit-1;
// ========================================= //
  arma::colvec upd_arma = arma::conv_to< arma::colvec >::from(upds);
  arma::colvec point_arma = arma::conv_to< arma::colvec >::from(points);
  arma::mat ystatsList=arma::join_rows(paramsmon,upd_arma,point_arma);
  return ystatsList;
}

//////////////////////////////////////////////////////////

//////////////////////////////////////////////////////////
//              Energy landscape analysis               //
//////////////////////////////////////////////////////////

// -- Functions for Energy landscape analysis 

// -- Community Energy
// [[Rcpp::export]]
inline double cEnergy(const arma::rowvec& state, const arma::rowvec& alpha, const arma::mat& beta){
  arma::mat res = -state * alpha.t() - (state* (state * beta).t() ) / 2;
  return as_scalar(res);
}

//////////////////////////////////////////////////////////

// [[Rcpp::export]]
arma::mat SteepestDescent_cpp(arma::rowvec state, arma::rowvec alpha, arma::mat beta){
  
  // ================================ //
  int term=0; 
  arma::mat energi; 
  arma::mat y1 = arma::conv_to<arma::mat>::from(state);
  double y2 = cEnergy(state, alpha, beta); 
  
  arma::mat ystat;
  double yene;
  arma::mat sign;
  double minene;
  arma::uword mp;
  // ================================ //
  
  do{
    // ============================= //
    ystat= y1;
    yene = y2;
    // ============================= //
    
    // -- Energy
    energi = -alpha-ystat*beta;
    sign = (-2 * ystat +1);
    energi %=sign; energi +=yene;
    
    minene =energi.min() ;
    if( minene < yene ){
      mp = energi.index_min();
      ystat(0, mp) = std::abs(ystat(0, mp)-1);
      
      y1.swap(ystat);
      y2 = minene;
      
    }else{
      term+=1;
    }
  } while (term==0);
  
  return arma::join_rows(y1, arma::mat(1,1).fill(y2));
  
}
// [[Rcpp::export]]
arma::mat SSestimate(arma::rowvec alpha, arma::mat beta, int itr=20000){
  
  arma::mat res = arma::zeros(itr, beta.n_cols+1);
  arma::imat intmat;
  arma::rowvec state;
  arma::mat ss;
  
  for(int i=0; i<itr; ++i ){
    
    intmat=arma::randi(1, beta.n_cols, arma::distr_param(0, 1));
    state=arma::conv_to<arma::rowvec>::from(intmat);
    
    ss = SteepestDescent_cpp(state, alpha, beta);
    
    res.row(i) = ss;
  }
  return res;
}

//////////////////////////////////////////////////////////


// [[Rcpp::export]]
arma::rowvec FindingTippingpoint_cpp(arma::rowvec s1,arma:: rowvec s2, 
                                     arma::rowvec alpha, arma::mat jj,
                                     int tmax=10000){
  
  // ======================================= //
  // Distance between stable states
  arma::uvec pd=find( abs(s1-s2)==1);
  arma::irowvec sequ=arma::conv_to<arma::irowvec>::from(shuffle(pd));
  arma::rowvec tipState(alpha.n_elem+1);

  const int& sn=sequ.n_elem;

  // ======================================= //
  // ||||||||||||||||||||||||||||||||||| //
  // -- Definition
  double bde=arma::datum::inf;
  int tt=0 ;
  double minbde=arma::datum::inf;
  arma::uvec rpl;
  arma::irowvec seqnew;
  double bdenew;
  const int& SMrow=sn+1;
  const int& SMcol=s1.n_elem;
  arma::vec::fixed<2>c; 
  arma::mat samplePath;
  int r;
  int column;
  
  arma::mat samplePathtmp=arma::zeros(SMrow, SMcol);
  samplePathtmp.row(0)=s1;
  const arma::uvec& v = arma::linspace<arma::uvec>(1,SMrow-1, SMrow-1);
  arma::vec CE=arma::vec(SMrow);
  
  // ||||||||||||||||||||||||||||||||||| //
  do{
    tt+=1;
    // ||||||||||||||||||||||||||||||||||| //
    seqnew=sequ;
    rpl = arma::randperm(sequ.n_elem, 2);
    
    seqnew.swap_cols(rpl(0), rpl(1));
    
    samplePath=samplePathtmp;
    
    // ||||||||||||||||||||||||||||||||||| //
    r=1;
    
    for(int l=0; l<sn; ++l){
      column=seqnew(l);
      
      for(int i=r; i<SMrow; ++i){
  
        samplePath(i, column)=1;
        
      }
      r+=1;
    }
      
    samplePath.each_row(v) -= s1;
    samplePath=abs(samplePath);
    
    // ||||||||||||||||||||||||||||||||||| //
    
    for(int i=0; i<SMrow; ++i){
      
      CE(i)=cEnergy(samplePath.row(i), alpha, jj);
      
    }
    // ||||||||||||||||||||||||||||||||||| //
    
    bdenew=CE.max();
    if(bdenew<minbde){
      minbde=std::move(bdenew);
      
      tipState.cols(0, SMcol-1)=samplePath.row(CE.index_max());
      tipState.col(tipState.n_cols-1)=minbde;
      
    }
    
    c(0)=1; c(1)= ( (exp(bde))/(exp(bdenew)) );
    if( arma::randu<double>() < c.min()){
      sequ.swap(seqnew);
      bde=std::move(bdenew);
    }
    
    // ||||||||||||||||||||||||||||||||||| //
  } while (tt<tmax);

  return tipState;
}

// [[Rcpp::export]]
arma::mat TPestimate(const arma::mat& comb, arma::mat minset,
                     arma::rowvec alpha, arma::mat beta, 
                     const int& tmax=10000){
  
  const int& resrow=comb.n_rows;
  arma::mat res = arma::zeros(resrow, beta.n_cols+1);
  
  int n; int m;
  arma::rowvec ss1; arma::rowvec ss2; arma::rowvec tp;
  
  for(int i=0; i<resrow; ++i ){
    n=comb(i, 0);
    m=comb(i, 1);
    
    ss1=minset.row(n);
    ss2=minset.row(m);
      
    tp = FindingTippingpoint_cpp(ss1, ss2, alpha, beta, tmax);
    
    res.row(i) = tp;
  }
  
  return res;
}

//////////////////////////////////////////////////////////

// To stop calculation.
int checkIdent(const arma::mat& ss, const arma::rowvec& ysim)
{
  int rows=ss.n_rows;
  arma::vec res(rows);
  
  for(int l=0; l<rows; l++){
    res(l)=all(ysim==ss.row(l));
  }
  
  return any(res);
}

// Convert to decimal
int convDec(arma::rowvec v)
{
  int tmp;
  std::string str;
  for(int i=0; i< v.n_elem; i++){
    tmp=v(i);
    str+=std::to_string(tmp);
  }
  
  return std::stoi(str, 0, 2); 
  
}

// Convert to string
inline std::string convStr(const arma::rowvec& v)  {
  int tmp;
  std::string str;
  for(int i=0; i< v.n_elem; i++){
    tmp=v(i);
    str+=std::to_string(tmp);
  }
  
  return str; 
  
}

// Entropy numeric
inline double entropy(const arma::vec& v)
{
  arma::vec uniq=arma::unique(v);
  int N=uniq.n_elem;
  const double& total=v.n_elem;
  double ent=0;
  
  for(int i=0; i<N; i++){
    const double& prob=sum(v==uniq(i))/total;
    ent+=prob*log2(prob);
  }
  // Rcout << probV << endl;
  
  return ent*(-1);
}

// Entropy string
inline double entropy2(const Rcpp::StringVector& v)
{
  Rcpp::IntegerVector tab=table(v);
  const double& total=v.length();
  //NumericVector numtab=as<NumericVector>(tab);
  
  double ent=0;
  for(int i=0; i<tab.size(); i++){
    const double& prob=tab[i]/total;
    ent+=prob*log2(prob);
  }
  
  return ent*(-1);
}

// [[Rcpp::export]]
arma::mat SSentropy_cpp(arma::mat uoc, arma::mat ss,
                        arma::rowvec alpha, arma::mat beta, 
                        int seitr=1000, int convTime=10000){
 
  arma::mat entropyres=arma::zeros(uoc.n_rows, 3);
  for(int i=0; i < uoc.n_rows; i++){
    //int total_conv = 0;
    arma::mat logmat;
    //Rcpp::StringVector ssid(seitr);
    std::vector<std::string> ssid;
    arma::mat stable=arma::zeros(seitr, 2);
    
    for(int l=0; l < seitr; l++){
      int ident = 0;
      arma::rowvec ysim=uoc.row(i); 
      int tt = 0;

      do{
        tt+=1;
        logmat=LogmodelSimplecpp(ysim, alpha, beta);
        ysim=OnestepHBScpp(ysim, logmat);
        
        if(checkIdent(ss, ysim)==1){
          ident=1;
          //total_conv+=1;
          break;
        }
      } while ( tt<convTime );

      if (ident == 1) {
        ssid.push_back(convStr(ysim)); 
      }
      
      //ssid(l)=convStr(ysim);
      stable(l, 0)=tt;
      //stable(l, 1)=(tt==convTime);
      stable(l, 1) =ident;
      
    }
    //std::cout << "total_conv: " << total_conv << std::endl;
    // changed by Takano: 2025/08/14
    Rcpp::StringVector ssid_vec(ssid.begin(), ssid.end());
    entropyres(i,0)=entropy2(ssid_vec);
    entropyres(i,1)=mean(stable.col(0));
    entropyres(i,2)=sum(stable.col(1));
  }
  //return stable;
  return entropyres;
}

