#include <Rcpp.h>
#include <climits>
using namespace Rcpp;

// FUNCTIONS FOR EXACT SAMPLING //

// Inner function to resize list:
List resize( const List& x, int n ){
    int oldsize = x.size() ;
    List y(n) ;
    for( int i=0; i<oldsize; i++) y[i] = x[i] ;
    return y ;
}

// Inner function to simulate random uniforms in a matrix:
NumericMatrix RandMat(int nrow, int ncol)
 {
  int N = nrow * ncol;
  NumericMatrix Res(nrow,ncol);
  NumericVector Rands  = runif(N);
   for (int i = 0; i < N; i++) 
  {
    Res[i] = Rands[i];
  }
  return(Res);
 }

// Resolve the optional 'delta' (single-ion / Blume-Capel quadratic term)
// argument to a concrete length-N vector. R_NilValue (NULL) -> all zeros (pure
// Ising). A single value is recycled to all nodes; a length-N vector is used
// as-is. Other lengths are an error.
NumericVector resolveDelta(Nullable<NumericVector> delta, int N)
{
  NumericVector out(N); // zero-initialized
  if (delta.isNotNull())
  {
    NumericVector d(delta);
    if (d.size() == N)
    {
      for (int i=0;i<N;i++) out[i] = d[i];
    } else if (d.size() == 1)
    {
      for (int i=0;i<N;i++) out[i] = d[0];
    } else if (d.size() != 0)
    {
      stop("'delta' must have length 1 or the number of nodes.");
    }
  }
  return out;
}

// Computes maximal and minimal probability of node flipping:
NumericVector PplusMinMax(int i, NumericMatrix J, NumericVector s, NumericVector h, double beta, NumericVector responses)
{
  // The function computes the probability that node i is in Response 1 instead of 0, given all other nodes, which might be missing.
  // Output: minimal and maximal probablity

  NumericVector H0(2, h[i] * responses[0]); // relevant part of the Hamiltonian for state = 0
  NumericVector H1(2, h[i] * responses[1]); // relevant part of the Hamiltonian for state = 1

  NumericVector Res(2);

  int N = J.nrow();
  NumericVector TwoOpts(2);

  for (int j=0; j<N; j++)
  {
    if (i != j)
    {
      if (!ISNA(s[j]))
      {
       H0[0] += J(i,j) * responses[0] * s[j];
       H0[1] += J(i,j) * responses[0] * s[j];
       H1[0] += J(i,j) * responses[1] * s[j];
       H1[1] += J(i,j) * responses[1] * s[j]; 
      } else 
      {
               
        TwoOpts[0] = J(i,j) * responses[1] * responses[0];
        TwoOpts[1] = J(i,j) * responses[1] * responses[1];

        if (TwoOpts[1] > TwoOpts[0])
        {
          H1[0] += TwoOpts[0];
          H1[1] += TwoOpts[1];
          
          H0[0] += J(i,j) * responses[0] * responses[0];
          H0[1] += J(i,j) * responses[0] * responses[1];
        } else 
        {
          H1[0] += TwoOpts[1];
          H1[1] += TwoOpts[0];          
          
          H0[0] += J(i,j) * responses[0] * responses[1];
          H0[1] += J(i,j) * responses[0] * responses[0];
        }
      }
    }
  }

  Res[0] = exp(beta * H1[0]) / ( exp(beta * H0[0]) + exp(beta * H1[0]) );
  Res[1] = exp(beta * H1[1]) / ( exp(beta * H0[1]) + exp(beta * H1[1]) );
  
  
  return(Res);
}
       
// Inner function:
NumericVector IsingEx(NumericMatrix graph, NumericVector thresholds, double beta, int nIter, NumericVector responses, bool exact,
NumericVector constrain)
{
  // Parameters and results vector:
  int N = graph.nrow();
  NumericVector state(N, NA_REAL);
  double u;
  NumericVector P(2);
  int maxChain = 100;
  List U(1);
  int minT = 0;
  bool anyNA = true;
    
  do
  { 
    // Resize U if needed:
    if (minT > 0)
    {
      U = resize(U, minT+1);
    }
    
    // Generate new random numbers:
    U[minT] = RandMat(nIter, N);
    
    // Initialize states. Constrained (clamped) nodes are fixed to their
    // constrained value and never updated; CFTP on the remaining (free) nodes
    // then yields exact samples from the conditional distribution given the
    // clamped nodes (clamped neighbors are simply known values in the
    // monotone-coupling bounds of PplusMinMax, never NA):
    for (int i=0; i<N; i++)
    {
      if (!ISNA(constrain[i]))
      {
        state[i] = constrain[i];
      } else if (exact)
      {
        state[i] = NA_REAL;
      } else
      {
        state[i] = ifelse(runif(1) < 0.5, responses[1], responses[0])[0];
      }
    }

    // START ALGORITHM
    for (int t=minT; t > -1;  t--)
    {
      for (int it=0;it<nIter;it++)
      {
        NumericMatrix Ucur = U[t];
        for (int node=0;node<N;node++)
        {
          // Clamped nodes are never updated:
          if (!ISNA(constrain[node])) continue;
          u = Ucur(it, node);
          P = PplusMinMax(node, graph, state, thresholds, beta, responses);
          if (u < P[0])
          {
            state[node] = responses[1];
          } else if (u >= P[1])
          {
            state[node] = responses[0];
          } else
          {
            state[node] = NA_REAL;
          }
        }
      }
    }
    
    anyNA = false;
    if (exact)
    {
      if (minT < maxChain)
      {
       for (int i=0; i<N; i++)
       {
        if (ISNA(state[i]))
        {
          anyNA = true;
        }
        }
      } 
    }    
    minT++;
    
  } while (anyNA);

  // Rf_PrintValue(wrap(minT));
  return(state);
}


// FUNCTIONS FOR METROPOLIS SAMPLER //
// Computes the full-conditional categorical distribution of node i over all
// response options, given the (complete) states of all other nodes.
// Returns a probability vector of length responses.size().
// For two response options (and delta = 0) this reduces exactly to the binary
// Ising conditional. 'delta' is the per-node Blume-Capel quadratic term.
NumericVector Pcat(int i, NumericMatrix J, NumericVector s, NumericVector h, double beta, NumericVector responses, NumericVector delta)
{
  int N = J.nrow();
  int K = responses.size();

  // Local field acting on node i: F = h_i + sum_{j != i} J_ij s_j.
  // The part of -H depending on s_i is r * F - delta_i * r^2 (H contains the
  // term +delta_i * s_i^2), so the conditional depends only on F and delta_i.
  double F = h[i];
  for (int j=0; j<N; j++)
  {
    if (i != j)
    {
      F += J(i,j) * s[j];
    }
  }

  // Softmax over beta * (responses[k] * F - delta_i * responses[k]^2), shifted by
  // the maximum for numerical stability (prevents overflow when beta or the
  // field is large).
  NumericVector P(K);
  double maxH = beta * (responses[0] * F - delta[i] * responses[0] * responses[0]);
  for (int k=1; k<K; k++)
  {
    double Hk = beta * (responses[k] * F - delta[i] * responses[k] * responses[k]);
    if (Hk > maxH) maxH = Hk;
  }
  double sum = 0;
  for (int k=0; k<K; k++)
  {
    P[k] = exp(beta * (responses[k] * F - delta[i] * responses[k] * responses[k]) - maxH);
    sum += P[k];
  }
  for (int k=0; k<K; k++) P[k] /= sum;

  return(P);
}

// Draws a response index from categorical probabilities P given uniform u.
// Iterates from the last response downward so that, for two response options,
// the mapping is identical to the original binary sampler (u < P[1] -> response 1).
int drawResponse(NumericVector P, double u)
{
  int K = P.size();
  double cum = 0;
  for (int k=K-1; k>0; k--)
  {
    cum += P[k];
    if (u < cum) return(k);
  }
  return(0);
}

// Initializes a random state with each node drawn uniformly from all response
// options. For two response options this reproduces the original binary
// initialization (uniform over the two responses); for more options it avoids
// biasing the starting state toward the first two responses, which otherwise
// skews the realized mode at high beta when the chain mixes slowly.
NumericVector randomState(int N, NumericVector responses)
{
  int K = responses.size();
  NumericVector state(N);
  if (K == 2)
  {
    state = ifelse(runif(N) < 0.5, responses[1], responses[0]);
  } else
  {
    for (int i=0; i<N; i++) state[i] = responses[(int) floor(R::runif(0, K))];
  }
  return(state);
}


NumericVector IsingMet(NumericMatrix graph, NumericVector thresholds, double beta, int nIter, NumericVector responses,
NumericVector constrain, NumericVector delta)
{
  // Parameters and results vector:
  int N = graph.nrow();
  NumericVector state = randomState(N, responses);
  for (int i=0; i<N; i++)
  {
    if (!ISNA(constrain[i]))
    {
      state[i] = constrain[i];
    }
  }
  double u;
  NumericVector P;

    // START ALGORITHM
    for (int it=0;it<nIter;it++)
    {
      for (int node=0;node<N;node++)
      {
        if (ISNA(constrain[node]))
        {
         u = runif(1)[0];
         P = Pcat(node, graph, state, thresholds, beta, responses, delta);
         state[node] = responses[drawResponse(P, u)];
        }
      }
    }

  return(state);
}


///ISING PROCESS SAMPLER:
// [[Rcpp::export]]
NumericMatrix IsingProcess(int nSample, NumericMatrix graph, NumericVector thresholds, double beta, NumericVector responses, Nullable<NumericVector> delta = R_NilValue)
{
  // Parameters and results vector:
  int N = graph.nrow();
  NumericVector deltaVec = resolveDelta(delta, N);
  NumericVector state = randomState(N, responses);
  double u;
  NumericVector P;
  NumericMatrix Res(nSample,N);
  int node;

    // START ALGORITHM
    for (int it=0;it<nSample;it++)
    {
      node = floor(R::runif(0,N));
        u = runif(1)[0];
        P = Pcat(node, graph, state, thresholds, beta, responses, deltaVec);
        state[node] = responses[drawResponse(P, u)];
        for (int k=0; k<N; k++) Res(it,k) = state[k];
    }
   
  return(Res);
}

// OVERAL FUNCTION //
// [[Rcpp::export]]
NumericMatrix IsingSamplerCpp(int n, NumericMatrix graph, NumericVector thresholds, double beta, int nIter, NumericVector responses, bool exact,
NumericMatrix constrain, Nullable<NumericVector> delta = R_NilValue)
{
  int Ni = graph.nrow();
  NumericMatrix Res(n,Ni);
  NumericVector state(Ni);
  NumericVector constrainVec(Ni);
  NumericVector deltaVec = resolveDelta(delta, Ni);
  if (exact)
  {
    // Exact sampling (CFTP) is only used with two response options (enforced in
    // R). For two response options the on-site term delta_i * s_i^2 is an affine
    // function of s_i and folds exactly into the thresholds:
    //   delta_i * s_i^2 = delta_i*(r0+r1)*s_i - delta_i*r0*r1,
    // and the constant cancels in the conditional. Since H uses -h_i*s_i and
    // +delta_i*s_i^2, the effective threshold is h_i - delta_i*(r0+r1). This keeps
    // CFTP correct without changing the coupling routines.
    NumericVector threshEff(Ni);
    double rsum = responses[0] + responses[1];
    for (int i=0;i<Ni;i++) threshEff[i] = thresholds[i] - deltaVec[i] * rsum;
    for (int s=0;s<n;s++)
    {
      for (int i=0;i<Ni;i++) constrainVec[i] = constrain(s,i);
      state = IsingEx(graph, threshEff, beta, nIter, responses, exact, constrainVec);
      for (int i=0;i<Ni;i++) Res(s,i) = state[i];
    }
  } else
  {
    for (int s=0;s<n;s++)
    {
      for (int i=0;i<Ni;i++) constrainVec[i] = constrain(s,i);
      state = IsingMet(graph, thresholds, beta, nIter, responses, constrainVec, deltaVec);
      for (int i=0;i<Ni;i++) Res(s,i) = state[i];
    }
  }

  return(Res);
}


// HELPER FUNCTIONS //
// Hamiltonian:
// H = - sum_i h_i s_i - sum_{i<j} J_ij s_i s_j + sum_i delta_i s_i^2
// The last (Blume-Capel single-ion / quadratic) term is optional; delta = NULL
// reproduces the pure Ising Hamiltonian.
// [[Rcpp::export]]
double H(NumericMatrix J, NumericVector s, NumericVector h, Nullable<NumericVector> delta = R_NilValue)
{
  double Res = 0;
  int N = J.nrow();
  NumericVector d = resolveDelta(delta, N);
  for (int i=0;i<N;i++)
  {
    Res -= h[i] * s[i];
    Res += d[i] * s[i] * s[i];
    for (int j=i;j<N;j++)
    {
      if (j!=i) Res -= J(i,j) * s[i] * s[j];
    }
  }
  return(Res);
}


// Likelihood without Z
// [[Rcpp::export]]
double f(NumericMatrix Y, NumericMatrix J, NumericVector h, Nullable<NumericVector> delta = R_NilValue)
{
  double Res = 1;
  int Np = Y.nrow();
  int Ni = J.ncol();
  NumericVector s(Ni);
  for (int p=0;p<Np;p++)
  {
    for (int i=0;i<Ni;i++)
    {
      s[i] = Y(p,i);
    }
    Res *= exp(-1.0 * H(J, s, h, delta));
  }
  return(Res);
}


// VECTOR VERSIONS //


// Hamiltonian:
// [[Rcpp::export]]
double Hvec(IntegerVector s, NumericVector Theta, int N)
{
  double Res = 0;
  int c=0;
  for (int i=0;i<N;i++)
  {
    Res -= Theta[c] * s[i];
    c++;
  }
  for (int i=0;i<N;i++)
  {
    for (int j=i; j<N;j++)
    {
     if (j!=i) 
     {
       Res -= Theta[c] * s[i] * s[j]; 
       c++;
     }
    }
  }
  return(Res);
}

// Likelihood without Z
double fvec(IntegerMatrix Y, NumericVector Theta)
{
  double Res = 1;
  int Np = Y.nrow();
  int Ni = Y.ncol();
  IntegerVector s(Ni);
  for (int p=0;p<Np;p++)
  {
    for (int i=0;i<Ni;i++)
    {
      s[i] = Y(p,i);
    }
    Res *= exp(-1.0 * Hvec(s, Theta, Ni));
  }
  return(Res);
}

// Log-Likelihood without Z
double fveclog(IntegerMatrix Y, NumericVector Theta)
{
  double Res = 1;
  int Np = Y.nrow();
  int Ni = Y.ncol();
  IntegerVector s(Ni);
  for (int p=0;p<Np;p++)
  {
    for (int i=0;i<Ni;i++)
    {
      s[i] = Y(p,i);
    }
    Res -= Hvec(s, Theta, Ni);
  }
  return(Res);
}
