#include <cassert>
#include "yorick.hpp"
#include "sphSmooth.hpp"
#include "mykdtree.hpp"
#include "miniargs.hpp"

using namespace std;
using namespace CosmoTool;

#define N_SPH 16

struct VCoord{
  float v[3];
};

template<int i>
ComputePrecision getVelocity(const VCoord& v)
{
  return v.v[i];
}

ComputePrecision getUnity(const VCoord& v)
{
  return 1.0;
}

typedef SPHSmooth<VCoord> MySmooth;
typedef MySmooth::SPHTree MyTree;
typedef MyTree::Cell MyCell;

int main(int argc, char **argv)
{
  char *fname1, *fname2;
  double rLimit, boxsize, rLimit2, cx, cy, cz;
  int Nres;

  MiniArgDesc args[] = {
    { "INPUT DATA1", &fname1, MINIARG_STRING },
    { "RADIUS LIMIT", &rLimit, MINIARG_DOUBLE },
    { "BOXSIZE", &boxsize, MINIARG_DOUBLE },
    { "RESOLUTION", &Nres, MINIARG_INT },
    { "CX", &cx, MINIARG_DOUBLE },
    { "CY", &cy, MINIARG_DOUBLE },
    { "CZ", &cz, MINIARG_DOUBLE },
    { 0, 0, MINIARG_NULL }
  };

  if (!parseMiniArgs(argc, argv, args))
    return 1;

  float *v1_data;
  uint32_t *dimList;
  uint32_t rank;
  uint32_t N1_points, N2_points;
  int *bins = new int[Nres*Nres*Nres];

  rLimit2 = rLimit*rLimit;

  loadArray(fname1, v1_data, dimList, rank);
  assert(rank == 2);
  assert(dimList[1] == 6);

  N1_points = dimList[0];
  delete[] dimList;
  
  cout << "Got " << N1_points << " in the first file." << endl;

  MyCell *allCells_1 = new MyCell[N1_points];
  
  for (long i = 0; i < Nres*Nres*Nres; i++)
    bins[i] = 0;

  cout << "Shuffling data in cells..." << endl;
  for (int i = 0 ; i < N1_points; i++)
    {
      for (int j = 0; j < 3; j++)
	allCells_1[i].coord[j] = v1_data[i*6 + j];
      for (int k = 0; k < 3; k++)
	allCells_1[i].val.pValue.v[k] = v1_data[i*6 + 3 + k];
      allCells_1[i].active = true;
      allCells_1[i].val.weight = 0.0;

      long rx = floor((allCells_1[i].coord[0]+cx)*Nres/boxsize+0.5);
      long ry = floor((allCells_1[i].coord[1]+cy)*Nres/boxsize+0.5);
      long rz = floor((allCells_1[i].coord[2]+cz)*Nres/boxsize+0.5);
      
      if (rx < 0 || rx >= Nres || ry < 0 || ry >= Nres || rz < 0 || rz >= Nres)
	  continue;
	  
      bins[rx + ry*Nres + rz*Nres*Nres]++;
    }
  delete[] v1_data;
   uint32_t dims[3] = { Nres, Nres, Nres } ;
   saveArray("num_in_cell.nc", bins, dims, 3);

  cout << "Building trees..." << endl;
  MyTree tree1(allCells_1, N1_points);

  cout << "Creating smoothing filter..." << endl;
  MySmooth smooth1(&tree1, N_SPH);

  uint32_t outDimList[3] = { Nres, Nres, Nres };
  uint32_t outDimList2[4] = { 3, Nres, Nres, Nres };
  ProgressiveOutput<float> out_den_1 = 
    ProgressiveOutput<float>::saveArrayProgressive("density.nc", outDimList, 3);
  ProgressiveOutput<float> out_vel_1 = 
    ProgressiveOutput<float>::saveArrayProgressive("v3d.nc", outDimList2, 4);
  ProgressiveOutput<float> out_rad_1 = 
    ProgressiveOutput<float>::saveArrayProgressive("rad.nc", outDimList, 3);
 
  cout << "Weighing..." << endl;
  for (int rz = 0; rz < Nres; rz++)
   {
      double pz = (rz)*boxsize/Nres-cz;

      cout << rz << " / " << Nres << endl;
      for (int ry = 0; ry < Nres; ry++)
	{
	  double py = (ry)*boxsize/Nres-cy;
	  for (int rx = 0; rx < Nres; rx++)
	    {
	      double px = (rx)*boxsize/Nres-cx;
	      
	      MyTree::coords c = { px, py, pz };

	      double r2 = c[0]*c[0]+c[1]*c[1]+c[2]*c[2];
	      if (r2 > rLimit2)
		{
		  continue;
		}

	      uint32_t numInCell = bins[rx + Nres*ry + Nres*Nres*rz];
	      if (numInCell > N_SPH)
		smooth1.fetchNeighbours(c, numInCell);
	      else
		smooth1.fetchNeighbours(c);
	      smooth1.addGridSite(c);
	    }
	}
    }

  
  

  cout << "Interpolating..." << endl;
  for (int rz = 0; rz < Nres; rz++)
    {
      double pz = (rz)*boxsize/Nres-cz;

      cout << rz << " / " << Nres << endl;
      for (int ry = 0; ry < Nres; ry++)
	{
	  double py = (ry)*boxsize/Nres-cy;
	  for (int rx = 0; rx < Nres; rx++)
	    {
	      double px = (rx)*boxsize/Nres-cx;
	      
	      MyTree::coords c = { px, py, pz };

	      double r2 = c[0]*c[0]+c[1]*c[1]+c[2]*c[2];
	      if (r2 > rLimit2)
		{
		  out_vel_1.put(0);
		  out_vel_1.put(0);
		  out_vel_1.put(0);
		  out_den_1.put(0);
		  out_rad_1.put(0);
		  continue;
		}

	      uint32_t numInCell = bins[rx + ry*Nres + rz*Nres*Nres];
	      if (numInCell > N_SPH)
		smooth1.fetchNeighbours(c, numInCell);
	      else
		smooth1.fetchNeighbours(c);

	      float val;

	      out_rad_1.put(smooth1.getSmoothingLen());
	      val = smooth1.computeSmoothedValue(c, getVelocity<0>);
	      out_vel_1.put(val);
	      val = smooth1.computeSmoothedValue(c, getVelocity<1>);
	      out_vel_1.put(val);
	      val = smooth1.computeSmoothedValue(c, getVelocity<2>);
	      out_vel_1.put(val);
	      val = smooth1.computeSmoothedValue(c, getUnity);
	      out_den_1.put(val);
	    }
	}
    }

  return 0;
};