#include <vtkm/io/reader/VTKDataSetReader.h>
#include <vtkm/filter/Gradient.h>
#include <vtkm/filter/VectorMagnitude.h>
#include <vtkm/io/writer/VTKDataSetWriter.h>
#include <vtkm/cont/Initialize.h>
#include <vtkm/filter/NDHistogram.h>

#include <vtkm/worklet/NDimsHistMarginalization.h>

//#include <DataSet.h>
#include <vtkm/worklet/FieldHistogram.h>
#include <vtkm/cont/cuda/DeviceAdapterCuda.h>

typedef vtkm::Float64 CURRENT_DTYPE;

//this is particularly for this example sampling code
struct VariableCondition
{
  VariableCondition(vtkm::cont::ArrayHandle<bool> mv):marginalVariable(mv){}
  VTKM_EXEC
  bool operator()(vtkm::Id var, vtkm::Id binId) const
  {
    //find the id of the non-marginal variable
    vtkm::Id vid = 1;
    if(!marginalVariable.GetPortalConstControl().Get(0)) vid = 0;
    if (var == vid)
    {
      //if (1 < binId && binId < 4)
        return true;
    }

    return false;
  }

  vtkm::cont::ArrayHandle<bool> marginalVariable;
};

struct equalsK1
{
  equalsK1(vtkm::cont:: ArrayHandle <vtkm::Float32 > sample_probs):sample_probabilites(sample_probs){}
  equalsK1(){};
  VTKM_EXEC_CONT  bool operator() ( vtkm:: Pair<vtkm::Id,vtkm::Float32> x) 
  const 
  {
    vtkm::cont::ArrayHandle<vtkm::Float32>::PortalConstControl probPortal = sample_probabilites.GetPortalConstControl();
    //cout<<x.first<<" "<<x.second<<" "<<probPortal.Get(x.first)<<endl;
    return x.second<probPortal.Get(x.first); 
  }

  vtkm::cont:: ArrayHandle <vtkm::Float32 > sample_probabilites;
};

struct equalsK1_2d
{
  equalsK1_2d(vtkm::cont:: ArrayHandle <vtkm::Float32 > sample_probs):sample_probabilites(sample_probs){}
  equalsK1_2d(){};
  VTKM_EXEC_CONT  bool operator() ( vtkm:: Pair<vtkm::Id,vtkm::Float32> x) 
  const 
  {
    vtkm::cont::ArrayHandle<vtkm::Float32>::PortalConstControl probPortal = sample_probabilites.GetPortalConstControl();
    //cout<<x.first<<" "<<x.second<<" "<<probPortal.Get(x.first)<<endl;
    return x.second<probPortal.Get(x.first); 
  }

  vtkm::cont:: ArrayHandle <vtkm::Float32 > sample_probabilites;
};

struct randArraygen
{
  VTKM_EXEC_CONT  vtkm::Float32 operator() ( vtkm::UInt32 seed)  const  // RAND_MAX assumed to be 32767 
  { 
      seed = seed * 1103515245 + 12345; 
      return ((unsigned int)(seed/65536) % 32768)/32767.0; 
  } 
};

void display2DHistogram(vtkm::cont::DataSet& ndhist){
  vtkm::Id nonSparseBins = ndhist.GetField(0).GetNumberOfValues();

  vtkm::cont::ArrayHandle<vtkm::Id> binId0;
  ndhist.GetField(0).GetData().CopyTo(binId0);
  vtkm::cont::ArrayHandle<vtkm::Id> binId1;
  ndhist.GetField("magnitude").GetData().CopyTo(binId1);  

  std::vector<vtkm::cont::ArrayHandle<vtkm::Id>> binIds;

  binIds.push_back(binId0);
  binIds.push_back(binId1);

  vtkm::cont::ArrayHandle<vtkm::Id> freqs;
  ndhist.GetField("Frequency").GetData().CopyTo(freqs);

  for(int i = 0; i < nonSparseBins; i++)
  {
    vtkm::Id idx0 = binIds[0].GetPortalControl().Get(i);
    vtkm::Id idx1 = binIds[1].GetPortalControl().Get(i);
    vtkm::Id f = freqs.GetPortalControl().Get(i);

    std::cout << "(" << idx0 << "," << idx1 << ") --> " << f << std::endl; 
  }

}

void displayMarginalHistogram(vtkm::cont::DataSet& ndhist, bool var1_flag, bool var2_flag, vtkm::Id binCount){
  
  vtkm::Id nonSparseBins = ndhist.GetField(0).GetNumberOfValues();

  vtkm::cont::ArrayHandle<vtkm::Id> binId0;
  ndhist.GetField(0).GetData().CopyTo(binId0);
  vtkm::cont::ArrayHandle<vtkm::Id> binId1;
  ndhist.GetField("magnitude").GetData().CopyTo(binId1);  

  std::vector<vtkm::cont::ArrayHandle<vtkm::Id>> binIds;

  binIds.push_back(binId0);
  binIds.push_back(binId1);

  vtkm::cont::ArrayHandle<vtkm::Id> freqs;
  ndhist.GetField("Frequency").GetData().CopyTo(freqs);

  bool marginalVariableAry[2] = {var1_flag, var2_flag};
  vtkm::cont::ArrayHandle<bool> marginalVariable =  vtkm::cont::make_ArrayHandle(marginalVariableAry, 2);

  std::vector<vtkm::Id> numberOfBinsVec(2, binCount);
  vtkm::cont::ArrayHandle<vtkm::Id> numberOfBins = vtkm::cont::make_ArrayHandle(numberOfBinsVec);

  // calculate marginal histogram by the setup (return sparse represetnation)
  std::vector<vtkm::cont::ArrayHandle<vtkm::Id>> marginalBinIds;
  vtkm::cont::ArrayHandle<vtkm::Id> marginalFreqs;
  vtkm::worklet::NDimsHistMarginalization ndHistMarginalization;
  ndHistMarginalization.Run(binIds, freqs,  numberOfBins, marginalVariable, VariableCondition(marginalVariable),  marginalBinIds, marginalFreqs);

  vtkm::Id marginal_nbins = marginalBinIds[0].GetPortalControl().GetNumberOfValues();

  std::cout << "Marginal Histogram: " << std::endl;

  for(int i = 0; i < marginal_nbins; i++)
  {
    vtkm::Id idx0 = marginalBinIds[0].GetPortalControl().Get(i);
    vtkm::Id f = marginalFreqs.GetPortalControl().Get(i);

    std::cout << "(" << idx0 << ") --> " << f << std::endl; 
  }
}

void getMarginalHistogram(vtkm::cont::DataSet& ndhist, bool var1_flag, bool var2_flag, vtkm::Id binCount, 
                            std::vector<vtkm::cont::ArrayHandle<vtkm::Id>>& marginalBinIds,
                            vtkm::cont::ArrayHandle<vtkm::Id>& marginalFreqs){
  
  vtkm::Id nonSparseBins = ndhist.GetField(0).GetNumberOfValues();

  vtkm::cont::ArrayHandle<vtkm::Id> binId0;
  ndhist.GetField(0).GetData().CopyTo(binId0);
  vtkm::cont::ArrayHandle<vtkm::Id> binId1;
  ndhist.GetField("magnitude").GetData().CopyTo(binId1);  

  std::vector<vtkm::cont::ArrayHandle<vtkm::Id>> binIds;

  binIds.push_back(binId0);
  binIds.push_back(binId1);

  vtkm::cont::ArrayHandle<vtkm::Id> freqs;
  ndhist.GetField("Frequency").GetData().CopyTo(freqs);

  bool marginalVariableAry[2] = {var1_flag, var2_flag};
  vtkm::cont::ArrayHandle<bool> marginalVariable =  vtkm::cont::make_ArrayHandle(marginalVariableAry, 2);

  std::vector<vtkm::Id> numberOfBinsVec(2, binCount);
  vtkm::cont::ArrayHandle<vtkm::Id> numberOfBins = vtkm::cont::make_ArrayHandle(numberOfBinsVec);

  // calculate marginal histogram by the setup (return sparse represetnation)
  //std::vector<vtkm::cont::ArrayHandle<vtkm::Id>> marginalBinIds;
  //vtkm::cont::ArrayHandle<vtkm::Id> marginalFreqs;
  vtkm::worklet::NDimsHistMarginalization ndHistMarginalization;
  ndHistMarginalization.Run(binIds, freqs,  numberOfBins, marginalVariable, VariableCondition(marginalVariable),  marginalBinIds, marginalFreqs);

  //vtkm::Id marginal_nbins = marginalBinIds[0].GetPortalControl().GetNumberOfValues();
  
}

void displayConditionalGradientHistogram(vtkm::cont::DataSet& ndhist, vtkm::Id binCount, vtkm::Id dataBinId){
  vtkm::Id nonSparseBins = ndhist.GetField(0).GetNumberOfValues();

  vtkm::cont::ArrayHandle<vtkm::Id> binId0;
  ndhist.GetField(0).GetData().CopyTo(binId0);
  vtkm::cont::ArrayHandle<vtkm::Id> binId1;
  ndhist.GetField("magnitude").GetData().CopyTo(binId1);  

  std::vector<vtkm::cont::ArrayHandle<vtkm::Id>> binIds;

  binIds.push_back(binId0);
  binIds.push_back(binId1);


  vtkm::cont::ArrayHandle<vtkm::Id> freqs;
  ndhist.GetField("Frequency").GetData().CopyTo(freqs);


  //store the gradient magnitude frequencies (in a vector of size binCount) for conditioned on dataBinId for data distribution
  std::vector<vtkm::Id> gradCondFrequency(binCount, 0);

  for(int i=0; i < nonSparseBins; i++){
    vtkm::Id idx0 = binIds[0].GetPortalControl().Get(i);
    if(idx0 == dataBinId){
      vtkm::Id idx1 = binIds[1].GetPortalControl().Get(i);
      vtkm::Id f = freqs.GetPortalControl().Get(i);
      gradCondFrequency[idx1] = f;
    }
  }

  for(int i=0; i<binCount; i++){
    std::cout << gradCondFrequency[i] << ", " ;

  }
  std::cout << std::endl;
}

void getConditionalGradientHistogram(vtkm::cont::DataSet& ndhist, vtkm::Id binCount, vtkm::Id dataBinId, std::vector<vtkm::Id>& gradCondFrequency){
  vtkm::Id nonSparseBins = ndhist.GetField(0).GetNumberOfValues();

  vtkm::cont::ArrayHandle<vtkm::Id> binId0;
  ndhist.GetField(0).GetData().CopyTo(binId0);
  vtkm::cont::ArrayHandle<vtkm::Id> binId1;
  ndhist.GetField("magnitude").GetData().CopyTo(binId1);  

  std::vector<vtkm::cont::ArrayHandle<vtkm::Id>> binIds;

  binIds.push_back(binId0);
  binIds.push_back(binId1);


  vtkm::cont::ArrayHandle<vtkm::Id> freqs;
  ndhist.GetField("Frequency").GetData().CopyTo(freqs);


  //store the gradient magnitude frequencies (in a vector of size binCount) for conditioned on dataBinId for data distribution
  //std::vector<vtkm::Id> gradCondFrequency(binCount, 0);

  for(int i=0; i < nonSparseBins; i++){
    vtkm::Id idx0 = binIds[0].GetPortalControl().Get(i);
    if(idx0 == dataBinId){
      vtkm::Id idx1 = binIds[1].GetPortalControl().Get(i);
      vtkm::Id f = freqs.GetPortalControl().Get(i);
      gradCondFrequency[idx1] = f;
    }
  }

  
  
}

std::vector<std::string> readFileToVector(const std::string& filename)
{
    std::ifstream source;
    source.open(filename);
    std::vector<std::string> lines;
    std::string line;
    while (std::getline(source, line))
    {
        lines.push_back(line);
    }
    return lines;
}


int main(int argc, char *argv[])
{

  vtkm::cont::Initialize(argc, argv, vtkm::cont::InitializeOptions::None);

  //currently hard coded
  //remember to change the datatype typedef above when changing datafiles nyx:float64, asteroid:float32
  // vtkm::Id fieldId = 0; //which field in the vtk file to operate on nyx:2 asteriod:0
  // vtkm::Id binCount = 10;
  // vtkm::Float32 sample_percent = 0.01;
  // std::string input_fname = "/home/shazarika/projects/VTKM/data/nyx_converted_ascii.vtk";
  // std::string output_fname = "/home/shazarika/projects/VTKM/data/nyx_2d_sampling.stencil_info";

  // std::string input_fname = "/home/shazarika/projects/VTKM/data/yA31_v02_300x300x300_99.vtk";
  // std::string output_fname = "/home/shazarika/projects/VTKM/data/asteriod_2d_sampling_0_01.stencil_info";


  
  std::string input_config_fname(argv[1]);
  std::vector<std::string> input_arg_vector = readFileToVector(input_config_fname);


  std::string input_fname = input_arg_vector[0];
  std::string output_fname = input_arg_vector[1];

  vtkm::Id fieldId = std::stoi(input_arg_vector[2]);
  vtkm::Id binCount = std::stoi(input_arg_vector[3]);
  vtkm::Float32 sample_percent = std::stoi(input_arg_vector[4])/100.0;

  // std::cout << input_fname << std::endl;
  // std::cout << output_fname << std::endl;
  // std::cout << fieldId << std::endl;
  // std::cout << binCount << std::endl;
  // std::cout << sample_percent << std::endl;
  




  vtkm::io::reader::VTKDataSetReader reader(input_fname);
  vtkm::cont::DataSet dataSet = reader.ReadDataSet();
  
  std::string fieldName = dataSet.GetField(fieldId).GetName(); 
  std::cout << "Operating sampling on " << fieldName << " field." << std::endl;

  // then copyto a concrete handle
  vtkm::cont:: ArrayHandle <CURRENT_DTYPE> dataArray;
  vtkm::cont::VariantArrayHandle dataHandle = dataSet.GetField(fieldId).GetData();
  dataHandle.CopyTo(dataArray);
  // variantHandle.CopyTo(concreteHandle );
  vtkm::Int32 tot_points = dataArray.GetNumberOfValues();
  std::cout << "total values=" << tot_points << std::endl;

  vtkm::filter::Gradient gradientFilter;
  gradientFilter.SetActiveField(fieldName);
  gradientFilter.SetComputePointGradient(true);
  vtkm::cont::DataSet gradient = gradientFilter.Execute(dataSet);

  vtkm::filter::VectorMagnitude magnitudeFilter;
  magnitudeFilter.SetActiveField("Gradients");
  vtkm::cont::DataSet magnitude = magnitudeFilter.Execute(gradient);

  vtkm::cont::VariantArrayHandle grad_mag_Handle = magnitude.GetField("magnitude").GetData();
  vtkm::cont:: ArrayHandle <CURRENT_DTYPE> grad_dataArray;
  grad_mag_Handle.CopyTo(grad_dataArray);
  std::cout << "grad_mag fields total values=" << grad_dataArray.GetNumberOfValues() << std::endl;
  

  vtkm::filter::NDHistogram ndHistogram;
  // TODO: NDHistogram should internally come up with a unambiguous name for
  // the bins.
  ndHistogram.AddFieldAndBin(fieldName, binCount);
  ndHistogram.AddFieldAndBin("magnitude", binCount);
  vtkm::cont::DataSet histogram = ndHistogram.Execute(magnitude);

  std::cout << "NDhistogram Done!" << std::endl;

  //display2DHistogram(histogram);
  //displayMarginalHistogram(histogram, true, false, binCount);
  //displayMarginalHistogram(histogram, false, true, binCount);
  //displayConditionalGradientHistogram(histogram, binCount, 0);

  

  //Extract the marginal distributions for both the variables
  std::vector<vtkm::cont::ArrayHandle<vtkm::Id>> marginalBinIds;
  vtkm::cont::ArrayHandle<vtkm::Id> marginalFreqs;
  getMarginalHistogram(histogram, true, false, binCount, marginalBinIds, marginalFreqs);
  
  vtkm::Range range = ndHistogram.GetDataRange(0); // get the range of the 1st variable
  vtkm::Float32 delta = ndHistogram.GetBinDelta(0); // get the delta of the bins for 1st variable

  std::cout << "[marginal] range = [" << range.Min << ", " << range.Max << "]" << std::endl;
  std::cout << "[marginal] delta = " << delta << std::endl;

  vtkm::cont::ArrayHandle<vtkm::Id> counts;
  std::vector<vtkm::Id> countsVec(binCount,0);
  vtkm::Id marginal_nbins = marginalBinIds[0].GetPortalControl().GetNumberOfValues();
  for(int i=0; i<marginal_nbins; i++){
    vtkm::Id idx0 = marginalBinIds[0].GetPortalControl().Get(i);
    vtkm::Id f = marginalFreqs.GetPortalControl().Get(i);
    countsVec[idx0] = f;
  }
  counts = vtkm::cont::make_ArrayHandle(countsVec);
  //debug print
  std::cout << "freqencies of val marginal hostogram : ";
  for(int i=0; i<countsVec.size(); i++){
    std::cout << countsVec[i] << ", " ;
  }
  std::cout << std::endl;





  //Extract the grad magnitude histogram information
  std::vector<vtkm::cont::ArrayHandle<vtkm::Id>> grad_marginalBinIds;
  vtkm::cont::ArrayHandle<vtkm::Id> grad_marginalFreqs;
  getMarginalHistogram(histogram, false, true, binCount, grad_marginalBinIds, grad_marginalFreqs);
  
  vtkm::Range grad_range = ndHistogram.GetDataRange(1); // get the range of the 2nd variable
  vtkm::Float32 grad_delta = ndHistogram.GetBinDelta(1); // get the delta of the bins for 2nd variable

  std::cout << "[grad marginal] range = [" << grad_range.Min << ", " << grad_range.Max << "]" << std::endl;
  std::cout << "[grad marginal] delta = " << grad_delta << std::endl;

  vtkm::cont::ArrayHandle<vtkm::Id> grad_counts;
  std::vector<vtkm::Id> grad_countsVec(binCount,0);
  vtkm::Id grad_marginal_nbins = grad_marginalBinIds[0].GetPortalControl().GetNumberOfValues();
  for(int i=0; i<grad_marginal_nbins; i++){
    vtkm::Id idx0 = grad_marginalBinIds[0].GetPortalControl().Get(i);
    vtkm::Id f = grad_marginalFreqs.GetPortalControl().Get(i);
    grad_countsVec[idx0] = f;
  }
  grad_counts = vtkm::cont::make_ArrayHandle(grad_countsVec);
  //debug print
  std::cout << "freqencies of grad marginal hostogram : ";
  for(int i=0; i<grad_countsVec.size(); i++){
    std::cout << grad_countsVec[i] << ", " ;
  }
  std::cout << std::endl;

  



  vtkm::cont::ArrayHandleIndex indexArray (binCount);
  vtkm::cont::ArrayHandle<vtkm::Id> indices;
  vtkm::cont::Algorithm::Copy(indexArray,indices);
  vtkm::cont::ArrayHandle<vtkm::Id> counts_copy;
  vtkm::cont::Algorithm::Copy(counts,counts_copy);
  vtkm::cont::ArrayHandleZip <vtkm::cont:: ArrayHandle <vtkm::Id >, vtkm::cont:: ArrayHandle <vtkm::Id >> zipArray(counts_copy , indices );
  //sorts the histogram bins based on their counts
  vtkm::cont:: Algorithm ::Sort(zipArray );
  // vtkm::cont::ArrayHandle<vtkm::Pair< vtkm::Id , vtkm::Id>> dataVals;
  // zipArray.CopyTo(dataVals);
  vtkm::cont::ArrayHandleZip <vtkm::cont::ArrayHandle <vtkm::Id >, vtkm::cont:: ArrayHandle <vtkm::Id >>::PortalConstControl binPortal = zipArray.GetPortalConstControl();

  for (int i = 0; i < binCount; ++i)
  {
    std::cout << binPortal.Get(i) << std::endl;
  }

  // create a float array now to hold the fractional counts for computing acceptance histogram
  vtkm::cont::ArrayHandle<vtkm::Float32> targetCounts;

  vtkm::Float32 remainingSamples = sample_percent*tot_points;
  //cout<<"sample_percent "<<sample_percent<<" tot_points "<<tot_points<<endl;
  vtkm::Float32 remainingBins = binCount;
  std::vector<vtkm::Float32> targetSamples(binCount, 0.0);

  for (int i = 0; i < binCount; ++i)
  {
    vtkm::Float32 targetNeededSamples = remainingSamples/(1.0*remainingBins);
    vtkm::Float32 curCount = (vtkm::Float32)binPortal.Get(i).first;
    vtkm::Float32 samplesTaken;
    // cout<<"curCount "<<curCount<<" targetNeededSamples "<<targetNeededSamples<<endl;

    if (curCount<targetNeededSamples)
    {
      samplesTaken = curCount;
    }
    else // for speed up, this else loop can be used to set the rest of the samples
    {
      samplesTaken = targetNeededSamples;
    }
    targetSamples[binPortal.Get(i).second] = samplesTaken;
    remainingBins = remainingBins-1;
    remainingSamples = remainingSamples - samplesTaken; 
  }


  std::cout << "Target Samples taken from each bins of the value distribution:" << std::endl;
  vtkm::Float32 tot_samp_taken = 0.0;
  for(int i=0; i<binCount; i++){
    std::cout << "[" << i << "] --> " << targetSamples[i] << "/" << counts.GetPortalConstControl().Get(i) << std::endl ;
    tot_samp_taken += targetSamples[i];
  }
  std::cout << "total samples taken = " << tot_samp_taken << std::endl;

  std::vector<std::vector<vtkm::Float32>> acceptProbVec_2d;
  //initialize
  for(int i=0; i<binCount; i++){
    std::vector<vtkm::Float32> tempVec(binCount, 0.0);
    acceptProbVec_2d.push_back(tempVec);
  }

  
  for(int i=0; i<binCount; i++){
    vtkm::Float32 current_bin_target_samples = targetSamples[i];

    std::vector<vtkm::Id> gradCondFrequency(binCount, 0);
    getConditionalGradientHistogram(histogram, binCount, i, gradCondFrequency);

    for(int j=binCount-1; j>=0 ; j--){

      if(current_bin_target_samples <= gradCondFrequency[j])
      {
        acceptProbVec_2d[i][j] = current_bin_target_samples;
        current_bin_target_samples = 0;
      }
      else{
        acceptProbVec_2d[i][j] = gradCondFrequency[j];
        current_bin_target_samples -= gradCondFrequency[j];
      }

      if(gradCondFrequency[j] > 0)
        acceptProbVec_2d[i][j] /= gradCondFrequency[j];
      else
        acceptProbVec_2d[i][j] = 0.0;

    }
  }

  //create field with bin ids for val distribution
  vtkm::cont::ArrayHandle<vtkm::Id> binIndex;
  binIndex.Allocate(tot_points);
  // Worklet to set the bin number for each data value
  vtkm::worklet::FieldHistogram::SetHistogramBin<vtkm::Float32> binWorklet(binCount, range.Min, delta);
  vtkm::worklet::DispatcherMapField <vtkm::worklet::FieldHistogram::SetHistogramBin<vtkm::Float32>> setHistogramBinDispatcher(binWorklet);
  setHistogramBinDispatcher.Invoke(dataArray, binIndex);

  //create field with bin ids for grad distribution
  vtkm::cont::ArrayHandle<vtkm::Id> grad_binIndex;
  grad_binIndex.Allocate(tot_points);
  // Worklet to set the bin number for each data value
  vtkm::worklet::FieldHistogram::SetHistogramBin<vtkm::Float32> grad_binWorklet(binCount, grad_range.Min, grad_delta);
  vtkm::worklet::DispatcherMapField <vtkm::worklet::FieldHistogram::SetHistogramBin<vtkm::Float32>> grad_setHistogramBinDispatcher(grad_binWorklet);
  grad_setHistogramBinDispatcher.Invoke(grad_dataArray, grad_binIndex);


  std::vector<vtkm::Id> combined_index_vec(tot_points, -1);
  //this is not an optimal way of updating, TODO: use functors
  for(int i=0; i<combined_index_vec.size(); i++){
    combined_index_vec[i] = binIndex.GetPortalConstControl().Get(i) * binCount + grad_binIndex.GetPortalConstControl().Get(i);
  }
  vtkm::cont::ArrayHandle<vtkm::Id> combined_binIndex = vtkm::cont::make_ArrayHandle(combined_index_vec);

  

  vtkm::cont::ArrayHandleImplicit<randArraygen> randArray(randArraygen(),tot_points);
  auto stencil = vtkm::cont:: make_ArrayHandleZip(combined_binIndex, randArray);

  //convert vec to ArrayHandle from acceptanceProbsVec
  std::vector<vtkm::Float32> probVec;
  for(int i=0; i<binCount; i++){
    for(int j=0; j<binCount; j++){
      probVec.push_back(acceptProbVec_2d[i][j]);
    }
  }
  vtkm::cont::ArrayHandle<vtkm::Float32> probArray = vtkm::cont::make_ArrayHandle(probVec);

  auto stencilBool = vtkm::cont::make_ArrayHandleTransform(stencil , equalsK1_2d(probArray));
  auto stencilPortal = stencilBool.GetPortalConstControl();
  std::vector<vtkm::Int32> stencilBool_vec;
  for(int i=0; i<stencilBool.GetNumberOfValues(); i++){
    vtkm::Int32 boolVal = stencilPortal.Get(i);
    stencilBool_vec.push_back(boolVal);
  }
  vtkm::cont::ArrayHandle<vtkm::Int32> stencilBool_ah = vtkm::cont::make_ArrayHandle(stencilBool_vec);

  //// WRITE OUT THE STENCIL VECTOR FOR VISUALIZATION
  std::ofstream fout(output_fname.c_str());
  for(auto const& x : stencilBool_vec){
    fout << x << '\n';
  }

  //// WORKAROUND TO WRITE STRUCTURED DATASETS USING VTK-M
  // vtkm::cont::DataSet out_stencil_ds;
  // out_stencil_ds.AddField(vtkm::cont::make_FieldPoint("stencil_info", stencilBool_ah));
  // std::cout << "new no. of fields [out_stencil_ds]=" << out_stencil_ds.GetNumberOfFields() << std::endl;
  // std::cout << "field name = " << out_stencil_ds.GetField(0).GetName() << std::endl;
  // std::cout << "No. of Coordinate Systems [out_stencil_ds] = " << out_stencil_ds.GetNumberOfCoordinateSystems() << std::endl;
  // //write this dataset out to vtk and then combine them with the original data in a sepearte python script
  // vtkm::io::writer::VTKDataSetWriter writer("/home/shazarika/projects/VTKM/data/nyx_2d_sampling_stencil_info.vtk");
  // writer.WriteDataSet(out_stencil_ds);


  //// FILTER OUT THE DATAARRAY BASED ON THE STENCIL 
  vtkm::cont::ArrayHandle <CURRENT_DTYPE> output;
  vtkm::cont::Algorithm::CopyIf(dataArray, stencilBool, output);
  std::cout << "final total sample counts=" << output.GetNumberOfValues() << std::endl;











  // //vtkm::io::writer::VTKDataSetWriter writer("grad_mag_hist.vtk");
  // //writer.WriteDataSet(histogram);
}