#include <iostream>
#include <valarray>
#include <dmtk/matrix.h>
#include <LAPACK/dsyev.h>

#define IBITS(n,i) (((n) & ((long long int)1 << i)) >> i)

enum
{
  RIGHT,
  TOP,
  LEFT,
  BOTTOM,
};

class site
{
  public:
    int nn[4];
};

void write_state(int n, int state);
void neighbors(int n, int bc, std::valarray<site> &lattice);
void build_tables(std::valarray<double> &hdiag, std::valarray<double> &hflip);
int build_basis(int n, int nup, std::valarray<int> &basis);
int bisect(int state, std::valarray<int> &basis, int dim);
void build_hami(dmtk::Matrix<double> &h, std::valarray<int> &basis, int dim, int n, std::valarray<site> &lattice, std::valarray<double> &hdiag, std::valarray<double> &hflip);

int main()
{

  int n, nup;
  int max_dim = 1000;
  int dim; 
  std::valarray<int> basis(max_dim);
  std::valarray<site> lattice(36);
  std::valarray<double> hdiag(4), hflip(4);
  dmtk::Matrix<double> hami(max_dim,max_dim);
 
  int bc;

  std::cout << "Input number of sites (N) : " << std::endl;
  std::cin >> n; 
  std::cout << "Input number of up spins (nup) : " << std::endl;
  std::cin >> nup; 

  bc = 1; // 0: OBC - 1: PBC
  neighbors(n, bc, lattice);

// Initialization

  build_tables(hdiag, hflip);

// Basis

  dim = build_basis(n, nup, basis);

// Hamiltonian matrix

  hami.resize(dim,dim);
  build_hami(hami, basis, dim, n, lattice, hdiag, hflip);
//  diagonalize(h);

  for(int i = 0; i < dim; i++){
    for(int j = 0; j < dim; j++){
      std::cout << i << " "<< j << " " << hami(i,j) << std::endl;
    }
  }

// Diagonalization

  double w[dim];
  double work[3*dim];
  int info;
  dsyev_('V', 'L', dim, hami.array(), dim, w, work, 3*dim, info); 

  std::cout << info << std::endl;

  for(int i = 0; i < dim; i++)
    std::cout << "Eigenvalue " << i << " " << w[i] << std::endl;

// Measurement

  for(int l = 0; l < dim; l++){
// We measure O for each eigenvector l

    std::cout << "=======================================\n";
    std::cout << "Eigenvector " << l << std::endl;
    std::cout << "=======================================\n";

    double mtot = 0; // Stag. mag. for eigenvector l
    for(int i = 0; i < dim; i++) {
// We measure O for the state i
      int state = basis[i];
      write_state(n,state); std::cout << " " << hami(l,i) << " ";
      int m = 0; // Stag. Mag. for configuration i
      for(int site = 0; site < n; site++){
        m += (-2*IBITS(site,0)+1)*(2*IBITS(state, site)-1);
      }
      std::cout << m << std::endl;
      mtot += hami(l,i)*hami(l,i)*abs(m);
    }
    mtot *= 0.5;

    std::cout << "Stag. Mag. " << l << " " << mtot << std::endl;

  }

}

void
build_hami(dmtk::Matrix<double> &h, std::valarray<int> &basis, int dim, int n, std::valarray<site> &lattice, std::valarray<double> &hdiag, std::valarray<double> &hflip)
{

  h = double(0);
  for(int i = 0; i  < dim; i++){
    int state = basis[i];

    // Diagonal term
    for(int site_i = 0; site_i < n; site_i++){
      int site_j = lattice[site_i].nn[RIGHT]; 
      if(site_j != -1){
        int two_sites  = IBITS(state,site_i) | (IBITS(state,site_j) << 1);   
        double value = hdiag[two_sites];
        h(i,i) += value;  
      }
    }
  }
//////////////////////////////////////////////////////////////////

//  for(int i = 0; i < dim; i++){
//    for(int j = 0; j < dim; j++){
//      std::cout << i << " "<< j << " " << h(i,j) << std::endl;
//    }
//  }

//////////////////////////////////////////////////////////////////
  for(int i = 0; i  < dim; i++){
    int state = basis[i];
    // Off-diagonal term
    for(int site_i = 0; site_i < n; site_i++){
      int site_j = lattice[site_i].nn[RIGHT]; 
      if(site_j != -1){
        int mask = (1 << site_i) | (1 << site_j);
        int two_sites  = IBITS(state,site_i) | (IBITS(state,site_j) << 1);   
        double value = hflip[two_sites];
        if(value != 0.){
          int new_state = (state ^ mask);
          int j = bisect(new_state, basis, dim);
          h(i,j) += value;  
        }
      }
    }

  }

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

}


void
neighbors(int n, int bc, std::valarray<site> &lattice)
{
  // 1D only    N-1 , N-2, ...... , 1, 0.
  for(int i = 0; i < n; i++) {
    lattice[i].nn[RIGHT] = i-1;
    lattice[i].nn[LEFT] = i+1;
  }
  if(bc == 0){ // Open Boundary Conditions
    lattice[0].nn[RIGHT] = -1; // This means error
    lattice[n-1].nn[LEFT] = -1; 
  } else {    // Periodic Boundary Conditions
    lattice[0].nn[RIGHT] = n-1; // We close the ring 
    lattice[n-1].nn[LEFT] = 0; 
  }
}  

void
build_tables(std::valarray<double> &hdiag, std::valarray<double> &hflip)
{
  hdiag[0] = +0.25;
  hdiag[1] = -0.25;
  hdiag[2] = -0.25;
  hdiag[3] = +0.25;

  hflip[0] = 0.;
  hflip[1] = 0.5;
  hflip[2] = 0.5;
  hflip[3] = 0.;
}  

int
build_basis(int n, int nup, std::valarray<int> &basis)
{
  int nstates = 0;
  for(int state = 0; state < (1<<n); state++){
    int n_ones = 0;
    for(int bit = 0; bit < n; bit++){
      if((state & (1 << bit)) != 0) n_ones++;
    } 
    if(n_ones == nup) {
      basis[nstates] = state;
      nstates++;
      std::cout << nstates << " "; write_state(n,state); std::cout << "\n";
    }
  }

  return nstates;
}

void
write_state(int n, int state)
{
  for(int bit = 0; bit < n; bit++){
    std::cout << IBITS(state,bit);
  } 
}

int
bisect(int state, std::valarray<int> &basis, int dim)
{
  int ret_val = -1;

  int origin = 0;
  int end = dim-1;
  int middle = (origin+end)/2;

  while(true){
    int index_old = middle;

    middle = (origin+end)/2;

    if(state < basis[middle])
      end = middle;
    else
      origin = middle; 

    if(basis[middle] == state) break;

    if(middle == index_old){
      if(middle == end)
        end = end - 1;
      else
        origin = origin + 1;
    }

  }

  ret_val = middle;
  return ret_val;
}