////////////////////////////////////////////////////////////////////////////////
//
// gofr.C: compute radial distribution function from an xyz file
// 
// note: units of cell dimensions, rmax, dr and coordinates must be consistent
// The xyz file must contain a repeated sequence as:
// <natoms>
// <iteration> <cell_vec_a> <cell_vec_b> <cell_vec_c>
// atom x y z
// ...
// atom x y z
//
// where <cell_vec_x> are 3-vectors
//
////////////////////////////////////////////////////////////////////////////////

#include <iostream>
#include <stdlib.h>
#include "D3vector.h"
#include "UnitCell.h"
#include <string>
#include <vector>
using namespace std;
// D3vector pbc(D3vector r, D3vector cell);

int main(int argc, char** argv)
{
  UnitCell cell;
  string line, nm;
  string name1, name2;
  vector<string> name;
  vector<D3vector> tau;
  double omega;
  int na,nbin,nconfig=0;
  vector<int> count, isp1, isp2;
  int na1,na2;
  double rmax,dr;

  if ( argc != 5 )
  {
    cerr << " use: gofr name1 name2 rmax dr < file.xyz > g.dat" << endl;
    assert(0);
  }

  name1 = argv[1];
  name2 = argv[2];
  D3vector a0,a1,a2;
  rmax = atof(argv[3]);
  dr = atof(argv[4]);

  cin >> na;

  tau.resize(na);
  nbin = (int) (rmax/dr);
  count.resize(nbin);
  name.resize(na);
  isp1.resize(na);
  isp2.resize(na);
  na1 = 0;
  na2 = 0;

  // count[i] contains the number of atoms with distance 
  // in [(i-0.5)*dr,(i+0.5)*dr]
  for ( int i = 0; i < nbin; i++ )
  {
    count[i] = 0;
  }

  while ( !cin.eof() )
  {
    nconfig++;

    // read iter number and cell parameters on same line as na
    int iter;
    cin >> iter >> a0 >> a1 >> a2;
    while( !cin.eof() && cin.get() != '\n');
    cell.set(a0,a1,a2);
    omega = cell.volume();

    for ( int i = 0; i < na; i++ )
    {
      cin >> nm >> tau[i]; 
      while( !cin.eof() && cin.get() != '\n');
      if ( nconfig == 1 )
      {
        name[i] = nm;
        // name[i] now contains the name of the species
        // check if name1 substring is found at beginning of name[i]
        if ( name[i].find(name1) == 0 ) isp1[na1++] = i;
        if ( name[i].find(name2) == 0 ) isp2[na2++] = i;
      }
    }

    for ( int i = 0; i < na1; i++ )
    {
      for ( int j = 0; j < na2 ; j++ )
      {
        // D3vector t = pbc ( tau[isp1[i]] - tau[isp2[j]], cell );
        D3vector t = tau[isp1[i]] - tau[isp2[j]];
        cell.fold_in_ws(t);
        // find bin index in which length(t) falls
        int k = (int) (length(t/dr)+0.5);
        // length(t) is in bin k, i.e. length(t) in [k*dr-0.5*dr, k*dr+0.5*dr]
        if ( k < nbin ) count[k]++;
      }
    }
    cin >> na;
  }
  cerr << "nconfig=" << nconfig << endl;
  
  if ( na1 == 0 )
  {
    cerr << " no atoms " << name1 << " found" << endl;
    return 1;
  }
  if ( na2 == 0 )
  {
    cerr << " no atoms " << name2 << " found" << endl;
    return 1;
  }
  cerr << na1 << " " << name1 << " atoms" << endl;
  cerr << na2 << " " << name2 << " atoms" << endl;

  // normalization differs for same species vs different species
  int npairs;
  if ( name1 == name2 )
  {
    npairs = na1 * (na2-1);
  }
  else
  {
    npairs = na1 * na2;
  }

  // plot g(r)
  //cout << nbin-1 << " " << name1 << " " << name2 << " g(r)" << endl;
  cout << endl << endl;
  for ( int i = 1; i < nbin; i++ )
  {
    const double r = i*dr;
    const double rmin = (i - 0.5)*dr;
    const double rmax = (i + 0.5)*dr;
    const double vshell = ( 4.0 * M_PI / 3.0 ) * 
                    ( pow(rmax,3.0) - pow(rmin,3.0) );
    // count_id is the count in an ideal homogeneous system
    const double count_id = vshell * nconfig * npairs / omega;

    const double g = count[i]/count_id;
    cout << r << " " << g << endl;
  }
  
  // plot integrated density
  // cout << nbin-1 << " " << name1 << " " << name2 << " N(r)" << endl;
  cout << endl << endl;
  double n = 0.0;
  for ( int i = 1; i < nbin; i++ )
  {
    const double r = i*dr;
    n += ( (double) count[i] ) / ( na1 * nconfig );
    cout << r << " " << n << endl;
  }

  return 0;
}

#if 0
D3vector pbc(D3vector r, D3vector cell)
{
  double x = r.x;
  double y = r.y;
  double z = r.z;

  const double a = cell.x;
  const double b = cell.y;
  const double c = cell.z;

  while ( x >  0.5*a ) x -= a;
  while ( x < -0.5*a ) x += a;
  while ( y >  0.5*b ) y -= b;
  while ( y < -0.5*b ) y += b;
  while ( z >  0.5*c ) z -= c;
  while ( z < -0.5*c ) z += c;

  return D3vector(x,y,z);
}
#endif