//============================================================================
//  Copyright (c) Kitware, Inc.
//  All rights reserved.
//  See LICENSE.txt for details.
//
//  This software is distributed WITHOUT ANY WARRANTY; without even
//  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
//  PURPOSE.  See the above copyright notice for more information.
//============================================================================
#ifndef vtk_m_filter_FBSampling_hxx
#define vtk_m_filter_FBSampling_hxx

#include <vector>
#include <vtkm/cont/DataSet.h>
//#include <vtkm/worklet/NDimsHistogram.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 <vtkm/worklet/FieldHistogram.h>
#include <vtkm/cont/cuda/DeviceAdapterCuda.h>

namespace vtkm
{
namespace filter
{

inline VTKM_CONT FBSampling::FBSampling()
{
}

void FBSampling::SetFieldId(vtkm::Id fId)
{
  this->fieldId = fId;
}

void FBSampling::SetBinCount(vtkm::Id numBins)
{
  this->binCount = numBins;
}

void FBSampling::SetSampleRate(vtkm::Float32 sample_rate)
{
  this->sample_percent = sample_rate;
}


void FBSampling::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 FBSampling::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;
    }
  }
  
}



template <typename Policy>
inline VTKM_CONT vtkm::cont::DataSet FBSampling::DoExecute(const vtkm::cont::DataSet& inData,
                                                            vtkm::filter::PolicyBase<Policy> policy)
{
  //get the name of the field to extract data from for sampling
  std::string fieldName = inData.GetField(this->fieldId).GetName(); 
  
  // then copyto a concrete handle
  vtkm::cont::ArrayHandle <vtkm::Float64> dataArray;
  vtkm::cont::VariantArrayHandle dataHandle = inData.GetField(this->fieldId).GetData();
  dataHandle.CopyTo(dataArray);
  vtkm::Int32 tot_points = dataArray.GetNumberOfValues();
  
  
  vtkm::filter::Gradient gradientFilter;
  gradientFilter.SetActiveField(fieldName);
  gradientFilter.SetComputePointGradient(true);
  vtkm::cont::DataSet gradient = gradientFilter.Execute(inData);

  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 <vtkm::Float64> grad_dataArray;
  grad_mag_Handle.CopyTo(grad_dataArray);
  

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

  
  //Extract the marginal distributions for both the variables
  std::vector<vtkm::cont::ArrayHandle<vtkm::Id>> marginalBinIds;
  vtkm::cont::ArrayHandle<vtkm::Id> marginalFreqs;
  this->getMarginalHistogram(histogram, true, false, this->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

  
  //reconstruct the bins for the marginal scalar value distribution
  vtkm::cont::ArrayHandle<vtkm::Id> counts;
  std::vector<vtkm::Id> countsVec(this->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);


  //Extract the grad magnitude histogram information
  std::vector<vtkm::cont::ArrayHandle<vtkm::Id>> grad_marginalBinIds;
  vtkm::cont::ArrayHandle<vtkm::Id> grad_marginalFreqs;
  this->getMarginalHistogram(histogram, false, true, this->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

  //reconstruct the bins for the marginal gradient distribution
  vtkm::cont::ArrayHandle<vtkm::Id> grad_counts;
  std::vector<vtkm::Id> grad_countsVec(this->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);
  

  vtkm::cont::ArrayHandleIndex indexArray (this->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();

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

  vtkm::Float32 remainingSamples = this->sample_percent*tot_points;
  vtkm::Float32 remainingBins = this->binCount;
  std::vector<vtkm::Float32> targetSamples(this->binCount, 0.0);

  for (int i = 0; i < this->binCount; ++i)
  {
    vtkm::Float32 targetNeededSamples = remainingSamples/(1.0*remainingBins);
    vtkm::Float32 curCount = (vtkm::Float32)binPortal.Get(i).first;
    vtkm::Float32 samplesTaken;
    
    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; 
  }


  // vtkm::Float32 tot_samp_taken = 0.0;
  // for(int i=0; i<this->binCount; i++){
  //   tot_samp_taken += targetSamples[i];
  // }
  // std::cout << "total samples taken = " << tot_samp_taken << std::endl;



  // Start calculating the 2d acceptance probability histogram
  std::vector<std::vector<vtkm::Float32>> acceptProbVec_2d;
  //initialize
  for(int i=0; i<this->binCount; i++){
    std::vector<vtkm::Float32> tempVec(this->binCount, 0.0);
    acceptProbVec_2d.push_back(tempVec);
  }

  // for each bin in the val histogram, identify how many samples to pick from each bin of the corresponding marginal gradient histogram bin
  for(int i=0; i<this->binCount; i++){
    vtkm::Float32 current_bin_target_samples = targetSamples[i];
    std::vector<vtkm::Id> gradCondFrequency(this->binCount, 0);
    this->getConditionalGradientHistogram(histogram, this->binCount, i, gradCondFrequency);

    //Approach 1: Start picking samples from the highest gradient bin
    for(int j=this->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(this->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(this->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) * this->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<this->binCount; i++){
    for(int j=0; j<this->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);

  

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

  vtkm::cont::DataSet out_stencil_ds;
  //out_stencil_ds.AddField(vtkm::cont::make_FieldPoint("stencil_info", stencilBool_ah));
  out_stencil_ds.AddField(vtkm::cont::make_FieldPoint("sample_vals", output));


  return out_stencil_ds;

  
}

//-----------------------------------------------------------------------------
template <typename T, typename StorageType, typename DerivedPolicy>
inline VTKM_CONT bool FBSampling::DoMapField(vtkm::cont::DataSet&,
                                              const vtkm::cont::ArrayHandle<T, StorageType>&,
                                              const vtkm::filter::FieldMetadata&,
                                              vtkm::filter::PolicyBase<DerivedPolicy>)
{
  return false;
}
}
}
#endif
