/*+
    VIDE -- Void IDentification and Examination -- ./c_tools/mock/prepObservation.cpp
    Copyright (C) 2010-2014 Guilhem Lavaux
    Copyright (C) 2011-2014 P. M. Sutter

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; version 2 of the License.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License along
    with this program; if not, write to the Free Software Foundation, Inc.,
    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+*/


#include <healpix_map.h>
#include <healpix_map_fitsio.h>
#include <pointing.h>
#include <boost/format.hpp>
#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include "prepObservation_conf.h"
#include "contour_pixels.hpp"
#include <netcdf>
#include <CosmoTool/fortran.hpp>
#include <gsl/gsl_integration.h>

#define LIGHT_SPEED 299792.458

using namespace std;
using boost::format;
using namespace CosmoTool;
using namespace netCDF;

struct NYU_Data
{
  int index;
  int sector;
  int region;
  double ra, dec;
  double cz;
  double fgotten;
  double phi_z;
  long uniqueID;
};

struct Position
{
  double xyz[3];
};

struct ParticleData
{
  vector<int> id_gal;
  vector<double> ra;
  vector<double> dec;
  vector<double> redshift;
  vector<double> catalogID;
  vector<long> uniqueID;
  // PMS
  int mask_index;
  // END PMS
  int edgeFlag = 0;
  vector<Position> pos;
  double box[3][2];
  double Lmax;
};

typedef vector<NYU_Data> NYU_VData;

// this defines the expansion function that we will integrate
// Laveaux & Wandelt (2012) Eq. 24
struct my_expan_params { double Om; double w0; double wa; };
double expanFun (double z, void * p) {
  struct my_expan_params * params = (struct my_expan_params *)p;
  double Om = (params->Om);
  double w0 = (params->w0);
  double wa = (params->wa);

  //const double h0 = 1.0;
  const double h0 = 0.71;
  double ez;

  double wz = w0 + wa*z/(1+z);

  ez = Om*pow(1+z,3) + (1.-Om);
  //ez = Om*pow(1+z,3) + pow(h0,2) * (1.-Om)*pow(1+z,3+3*wz);

  ez = sqrt(ez);
  //ez = sqrt(ez)/h0;

  ez = 1./ez;

  return  ez;
}

void loadData(const string& fname, NYU_VData & data)
{
  ifstream f(fname.c_str());
   
  while (!f.eof())
    {
      NYU_Data d;
      f >> d.index >> d.sector >> d.region >> d.ra >> d.dec >> d.cz >> d.fgotten >> d.phi_z;
      // Maubert - avoid double counting of last particle in array if EOF is after a blank line
      if (!f) 
      { 
        continue; 
      } 
      // End Maubert
      d.uniqueID = d.index;
      data.push_back(d);
    }
}

void placeGalaxiesInCube(NYU_VData& data, ParticleData& output_data, 
                            bool useComoving, double omegaM) {
  double d2r = M_PI/180;

  gsl_function expanF;
  expanF.function = &expanFun;
  struct my_expan_params expanParams;
  double result, error;
  size_t nEval;

  expanParams.Om = omegaM;
  expanParams.w0 = -1.0;
  expanParams.wa = 0.0;
  expanF.params = &expanParams;

  output_data.pos.resize(data.size());
  output_data.id_gal.resize(data.size());
  output_data.ra.resize(data.size());
  output_data.dec.resize(data.size());
  output_data.redshift.resize(data.size());
  output_data.uniqueID.resize(data.size());

  for (int j = 0; j < 3; j++)
    {
      output_data.box[j][0] = -INFINITY;
      output_data.box[j][1] = INFINITY;
    }

  for (int i = 0; i < data.size(); i++)
    {  
      double ra = data[i].ra*d2r, dec = data[i].dec*d2r;
      double Dc = data[i].cz;
      Position& p = output_data.pos[i];

      if (useComoving) {
        gsl_integration_qng(&expanF, 0.0, data[i].cz/LIGHT_SPEED,
                            1.e-6, 1.e-6, &result, &error, &nEval);
        Dc = result*LIGHT_SPEED;
      }

      p.xyz[0] = Dc*cos(ra)*cos(dec);
      p.xyz[1] = Dc*sin(ra)*cos(dec);
      p.xyz[2] = Dc*sin(dec);

      output_data.id_gal[i] = data[i].index;
      output_data.ra[i] = ra;
      output_data.dec[i] = dec;
      output_data.redshift[i] = data[i].cz;
      output_data.uniqueID[i] = data[i].uniqueID;

      for (int j = 0; j < 3; j++) {
	      if (p.xyz[j] > output_data.box[j][0])
	        output_data.box[j][0] = p.xyz[j];
	      if (p.xyz[j] < output_data.box[j][1])
	        output_data.box[j][1] = p.xyz[j];	  
	    }
    }

  // normalize the box volume
  float left = INFINITY;
  float right = -INFINITY;
  for (int j = 0; j < 3; j++) {
    if (output_data.box[j][1] < left) left = output_data.box[j][1];
    if (output_data.box[j][0] > right) right = output_data.box[j][0];
  }
  for (int j = 0; j < 3; j++) {
    output_data.box[j][1] = left;
    output_data.box[j][0] = right;
  }

  double Rmax = -1;
  for (int j = 0; j < 3; j++) {
    Rmax = max(Rmax, max(output_data.box[j][0], -output_data.box[j][1]));
  }
  output_data.Lmax = Rmax;

  cout << format("Galaxy position generated: %d galaxies") % output_data.pos.size() << endl;
  cout << format("box is %g < x < %g;  %g < y < %g;   %g < z < %g") 
    % (1e-2*output_data.box[0][1]) % (1e-2*output_data.box[0][0])
    % (1e-2*output_data.box[1][1]) % (1e-2*output_data.box[1][0])
    % (1e-2*output_data.box[2][1]) % (1e-2*output_data.box[2][0]) << endl;
 
}

void flagEdgeGalaxies(prepObservation_info& args ,
			 Healpix_Map<float>& mask, 
			 vector<int>& contourPixels, 
			 NYU_VData& data, 
			 ParticleData& output_data,
			 bool useComoving, 
			 double omegaM) {

  //Maubert - Needed for comobile distance in mock_sphere
  gsl_function expanF;
  expanF.function = &expanFun;
  struct my_expan_params expanParams;
  expanParams.Om = omegaM;
  expanParams.w0 = -1.0;
  expanParams.wa = 0.0;
  expanF.params = &expanParams;
  
  double result, error;
  size_t nEval;
  //End Maubert - Needed for comobile distance in mock_sphere

  // TODO - REMOVE THIS
  output_data.mask_index = output_data.id_gal.size();

  // write a small text file with galaxy position (for diagnostic purposes)
  FILE *fp;
  fp = fopen("galaxies.txt", "w");
  for (int i = 0; i < data.size(); i++) {  
    Position& p = output_data.pos[i];
    fprintf(fp, "%e %e %e\n", 
            (p.xyz[0]), 
            (p.xyz[1]), 
            (p.xyz[2]));
  }
  fclose(fp);

/* NOTE: temporarily moved to python for quick debugging. Will move back to
         here once it's all sorted

  // convert redshift boundaries to covmoving if necessary
  double rmin;
  double rmax; 
  if (useComoving) {
    gsl_integration_qng(&expanF, 0.0, args.zMin_arg, 1.e-6, 1.e-6, &result, 
                        &error, &nEval);
    rmin = result*LIGHT_SPEED;
        
    gsl_integration_qng(&expanF, 0.0, args.zMax_arg, 1.e-6, 1.e-6, &result, 
                        &error, &nEval);
    rmax = result*LIGHT_SPEED;
  } else {
    rmin = args.zMin_arg * LIGHT_SPEED;
    rmax = args.zMax_arg * LIGHT_SPEED;
  }
  
  double dx = args.meanPartSep_arg;
  int nSteps = floor( (rmax - rmin) / dx);
  cout << "Assumed resolution element: " << dx << " " << nSteps << endl;

  // flag galaxies near mask edges
  // using the "ray marching" algorithm: follow rays along lines of sight
  //   of all mask edges, flagging nearest neighbor galaxies as we go
  // TODO - replace this with faster kd-tree search
  cout << "Flagging galaxies on edges of survey..." << endl;

  // explore rays along mask contours
  for (int pixel : contourPixels) {
    cout << "Working with pixel " << pixel << endl;
    vec3 v = mask.pix2vec(pixel);
    //cout << v*rmin << " " << v*rmax << endl;

    // march along single ray and find nearest neighbors
    for (int n = 0; n <= nSteps; n++) {
      double r = n*dx + rmin;
      vec3 rayPos = v*r;
 
      double x = rayPos.x;
      double y = rayPos.y;
      double z = rayPos.z;

      //cout << "Step " << n << " " << rayPos << endl;
  
      // scan all galaxies
      double minDist = INFINITY;
      double dist = 0;
      int closestGal = -1;
      for (int i = 0; i < data.size(); i++) {  
        Position& galPos = output_data.pos[i];

        dist = pow(galPos.xyz[0] - x, 2) + 
               pow(galPos.xyz[1] - y, 2) +
               pow(galPos.xyz[2] - z, 2);

        if (dist < minDist) closestGal = i;
      } // galaxy search


    } // marching along one ray

  } // all contours


  // flag galaxies near redshift boundaries
  cout << "Flagging galaxies at redshift boundaries..." << endl;

*/

} // end flagEdgeGalaxies

void saveForZobov(ParticleData& pdata, const string& fname, 
                  const string& paramname)
{
  UnformattedWrite f(fname);
  static const char axis[] = { 'X', 'Y', 'Z' };
  double Lmax = pdata.Lmax;
  double r2d = 180./M_PI;

  f.beginCheckpoint();
  f.writeInt32(pdata.pos.size());
  f.endCheckpoint();


  for (int j = 0; j < 3; j++)
    {
      cout << format("Writing %c components...") % axis[j] << endl;
      f.beginCheckpoint();
      for (uint32_t i = 0; i < pdata.pos.size(); i++)
	{
	  f.writeReal32((pdata.pos[i].xyz[j]+Lmax)/(2*Lmax));
	}
      f.endCheckpoint();
    }

  cout << format("Writing RA...")  << endl;
  f.beginCheckpoint();
  for (uint32_t i = 0; i < pdata.pos.size(); i++) {
	  f.writeReal32(pdata.ra[i]*r2d);
	}
  f.endCheckpoint();
   
  cout << format("Writing Dec...")  << endl;
  f.beginCheckpoint();
  for (uint32_t i = 0; i < pdata.pos.size(); i++) {
	  f.writeReal32(pdata.dec[i]*r2d);
	}
  f.endCheckpoint();
    
  cout << format("Writing Redshift...")  << endl;
  f.beginCheckpoint();
  for (uint32_t i = 0; i < pdata.pos.size(); i++) {
	  f.writeReal32(pdata.redshift[i]);
	}
  f.endCheckpoint();
   
  cout << format("Writing Unique ID...")  << endl;
  f.beginCheckpoint();
  for (uint32_t i = 0; i < pdata.pos.size(); i++) {
	  f.writeInt64(pdata.uniqueID[i]);
	}
  f.endCheckpoint();
   
  NcFile fp(paramname.c_str(), NcFile::replace);

  fp.putAtt("range_x_min", ncDouble, -Lmax/100.);
  fp.putAtt("range_x_max", ncDouble, Lmax/100.);
  fp.putAtt("range_y_min", ncDouble, -Lmax/100.);
  fp.putAtt("range_y_max", ncDouble, Lmax/100.);
  fp.putAtt("range_z_min", ncDouble, -Lmax/100.);
  fp.putAtt("range_z_max", ncDouble, Lmax/100.);
  fp.putAtt("mask_index", ncInt, pdata.mask_index); // PMS
  fp.putAtt("is_observation", ncInt, 1); // PMS

  int nOutputPart = pdata.mask_index;
  //int nOutputPart = pdata.pos.size();

  NcDim NumPart_dim = fp.addDim("numpart_dim", nOutputPart);
  NcVar v = fp.addVar("particle_ids", ncInt, NumPart_dim);
  //NcVar v2 = fp.addVar("expansion", ncDouble, NumPart_dim);

  //double *expansion_fac = new double[pdata.pos.size()];

  //for (int i = 0; i <  pdata.pos.size(); i++)
  //  expansion_fac[i] = 1.0;

  v.putVar({0}, {size_t(nOutputPart)}, &pdata.id_gal[0]);
  //v2->put(expansion_fac, pdata.pos.size());

}

int main(int argc, char **argv)
{
  prepObservation_info args_info;
  prepObservation_conf_params args_params;
 
  prepObservation_conf_init(&args_info);
  prepObservation_conf_params_init(&args_params);
  
  args_params.check_required = 0;
  if (prepObservation_conf_ext (argc, argv, &args_info, &args_params))
    return 1;
  
  if (!args_info.configFile_given)
    {
      if (prepObservation_conf_required (&args_info, PREPOBSERVATION_CONF_PACKAGE))
        return 1;
    }
  else
    {
      args_params.check_required = 1;
      args_params.initialize = 0;
      if (prepObservation_conf_config_file (args_info.configFile_arg,
					 &args_info,
					 &args_params))
	return 1;
    }
  
  prepObservation_conf_print_version();

  cout << "Loading galaxy data " << args_info.catalog_arg << "..." << endl;
  vector<NYU_Data> data;
  vector<int> contourPixels;
  ParticleData output_data;

  loadData(args_info.catalog_arg, data);

  cout << "Loading mask..." << endl;
  Healpix_Map<float> mask;
  Healpix_Map<float> o_mask;

  int newNside = args_info.nsideForContour_arg;
  read_Healpix_map_from_fits(args_info.mask_arg, o_mask);

  if (newNside == -1) newNside = o_mask.Nside();

  mask.SetNside(newNside, RING);
  mask.Import(o_mask);

  computeContourPixels(mask, contourPixels);

  cout << "Placing galaxies..." << endl;
  placeGalaxiesInCube(data, output_data, args_info.useComoving_flag, 
                      args_info.omegaM_arg);


  //cout << "Flagging edge galaxies..." << endl;
  flagEdgeGalaxies(args_info, mask, contourPixels,
                   data, output_data,args_info.useComoving_flag, 
                   args_info.omegaM_arg);

  saveForZobov(output_data, args_info.output_arg, args_info.params_arg);
  
  // PMS - TODO REMOVE THIS
  FILE *fp = fopen("mask_index.txt", "w");
  fprintf(fp, "%d", output_data.mask_index);
  fclose(fp);

  fp = fopen("total_particles.txt", "w");
  fprintf(fp, "%d", output_data.pos.size());
  fclose(fp);
  printf("Done!\n");
  // END PMS - TODO REMOVE THIS

  return 0;
}