// =======================================================================
//   Sample program to calculate density and sum of states of H2O
// =======================================================================

#include <iostream>
#include <iomanip>
#include <vector>
using namespace std;

// ----- direct count class (declaration) --------------------------------

class directCount {
 public:
  bool ready;                   // flag (ready for calculation)
  bool doneCount;               // flag (done counting)
  double grainSize;             // grain size [cm-1]
  vector<int> roundFreq;        // frequencies in 'grain size' unit
  vector<double> rho;           // density of states [cm]
  vector<double> W;             // sum of states [-]
 public:
  directCount();
  directCount(int numVib, double freq[], double grSiz, double maxE);
  void clear();
  void init(int numVib, double freq[], double grSiz, double maxE);
  bool count();
  ostream &print(ostream &os);
};
ostream &operator<<(ostream &os, directCount &dc);

// ===== main function ===================================================

int main() {
  double vibFreq[3] = {3657., 1595., 3756.};
  directCount dc(3, vibFreq, 100., 60000.);
  if (!dc.ready)
    { cout << " Failed to initialize directCount object.\n"; return 1; }
  dc.count();
  cout << dc;
  return 0;
}

// ----- direct count class (inplementation) -----------------------------

directCount::directCount() { clear(); }

directCount::directCount(int numVib, double freq[], double grSiz,
 double maxE) { init(numVib, freq, grSiz, maxE); }

void directCount::clear() {
  ready = false; doneCount = false;
  roundFreq.clear(); rho.clear(); W.clear();
}

void directCount::init(int numVib, double freq[], double grSiz,
 double maxE) {
  int iv, ig, ng;
  clear();
  if (numVib <= 0) { return; }
  for (iv = 0; iv < numVib; iv++) {
    roundFreq.push_back((int)(freq[iv] / grSiz + 0.5));
    if (roundFreq[iv] == 0) { return; }
  }
  grainSize = grSiz;
  ng = (int)(maxE / grSiz + 1.5);
  if (ng <= 1) { return; }
  for (ig = 0; ig < ng; ig++) {
    rho.push_back(0.);
    W.push_back(1.);
  }
  rho[0] = 1.;
  ready = true;
}

bool directCount::count() {
  int nv = roundFreq.size(), ng = rho.size(), iv, jfr, ig, jg;
  if (!ready) { return false; }
  for (iv = 0; iv < nv; iv++) {
    jfr = roundFreq[iv];
    for (ig = jfr; ig < ng; ig++) {
      jg = ig - jfr;
      rho[ig] += rho[jg];
      W[ig] += W[jg];
    }
  }
  for (ig = 0; ig < ng; ig++) { rho[ig] /= grainSize; }
  doneCount = true;
  return true;
}

ostream &directCount::print(ostream &os) {
  int ig, ng = rho.size();
  if (!doneCount) {
    cout << " No counting has been done properly.\n";
    return os;
  }
  os << "   E[cm-1],     rho[cm],           W\n";
  os << setiosflags(ios::right | ios::showpoint | ios::uppercase);
  for (ig = 0; ig < ng; ig++) {
    os << resetiosflags(ios::scientific) << setiosflags(ios::fixed)
       << setw(10) << setprecision(2) << ig * grainSize << ',';
    os << resetiosflags(ios::fixed) << setiosflags(ios::scientific)
       << setw(12) << setprecision(4) << rho[ig] << ',';
    os << setw(12) << W[ig] << endl;
  }
  return os;
}

ostream &operator<<(ostream &os, directCount &dc) { return dc.print(os); }