//============================================================================
//  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.
//
//============================================================================

#include <fides/DataSetReader.h>
#include <fides/DataSetWriter.h>

#include <vtkm/VectorAnalysis.h>
#include <vtkm/cont/ArrayHandleCartesianProduct.h>
#include <vtkm/cont/ArrayHandleUniformPointCoordinates.h>
#include <vtkm/cont/ArrayPortalToIterators.h>
#include <vtkm/cont/CellSetStructured.h>
#include <vtkm/cont/DataSetBuilderExplicit.h>
#include <vtkm/cont/DataSetBuilderRectilinear.h>
#include <vtkm/cont/DataSetBuilderUniform.h>

#include <numeric>
#include <random>
#include <string>
#include <vector>

#ifdef FIDES_USE_MPI
#include <mpi.h>
#endif

#include <string>
#include <unordered_map>
#include <vector>


class RandomNumbers
{
public:
  RandomNumbers()
    : Gen(0)
    , Dist(0, 1)
  {
  }

  vtkm::FloatDefault RandomFloat(vtkm::FloatDefault v0, vtkm::FloatDefault v1)
  {
    auto r = this->Dist(this->Gen);
    return v0 + (v1 - v0) * r;
  }

  template <typename T>
  T RandomVal(T v0, T v1)
  {
    vtkm::FloatDefault t = this->Dist(this->Gen);
    vtkm::FloatDefault dV = static_cast<vtkm::FloatDefault>(v1 - v0);

    return v0 + static_cast<T>(t * dV);
  }

  template <typename T>
  std::vector<T> RandomVec(T v0, T v1, std::size_t n)
  {
    std::vector<T> vec(n);
    do
    {
      for (auto& val : vec)
      {
        val = this->RandomVal(v0, v1);
      }
    } while (std::accumulate(vec.begin(), vec.end(), 0) == 0);
    //We don't want a vector full of 0, so if we get that, try again.

    return vec;
  }

private:
  std::default_random_engine Gen;
  std::uniform_real_distribution<vtkm::FloatDefault> Dist;
};

vtkm::cont::ArrayHandle<vtkm::FloatDefault> MakeArray(std::size_t numValues,
                                                      RandomNumbers& randomNumbers)
{
  std::vector<vtkm::FloatDefault> arr;

  arr.resize(numValues);
  for (auto& v : arr)
  {
    v = randomNumbers.RandomFloat(-100, 100);
  }

  return vtkm::cont::make_ArrayHandle(arr, vtkm::CopyFlag::On);
}

std::vector<vtkm::Vec3f> MakeVecVector(std::size_t numValues,
                                       vtkm::FloatDefault minVal,
                                       vtkm::FloatDefault maxVal,
                                       RandomNumbers& randomNumbers)
{
  std::vector<vtkm::Vec3f> arr;
  arr.resize(numValues);
  for (auto& v : arr)
  {
    v[0] = randomNumbers.RandomFloat(minVal, maxVal);
    v[1] = randomNumbers.RandomFloat(minVal, maxVal);
    v[2] = randomNumbers.RandomFloat(minVal, maxVal);
  }
  return arr;
}

vtkm::cont::ArrayHandle<vtkm::Vec3f> MakeVecArray(std::size_t numValues,
                                                  RandomNumbers& randomNumbers)
{
  auto arr = MakeVecVector(numValues, -100, 100, randomNumbers);
  return vtkm::cont::make_ArrayHandle(arr, vtkm::CopyFlag::On);
}

vtkm::cont::PartitionedDataSet DataSetsByRank(const std::vector<vtkm::cont::DataSet>& dataSets,
                                              int rank,
                                              const std::vector<int>& blocksPerRank)
{
  int offset = 0;
  for (std::size_t i = 0; i < static_cast<size_t>(rank); i++)
  {
    offset += blocksPerRank[i];
  }

  vtkm::cont::PartitionedDataSet pds;
  for (int i = 0; i < blocksPerRank[rank]; i++)
  {
    pds.AppendPartition(dataSets[offset + i]);
  }
  return pds;
}

void AddFields(vtkm::cont::DataSet& ds, RandomNumbers& randomNumbers)
{
  std::size_t numPoints = static_cast<std::size_t>(ds.GetNumberOfPoints());
  std::size_t numCells = static_cast<std::size_t>(ds.GetNumberOfCells());

  // Add scalar fields.
  ds.AddPointField("varPoint", MakeArray(numPoints, randomNumbers));
  ds.AddCellField("varCell", MakeArray(numCells, randomNumbers));

  // Add vector fields.
  ds.AddPointField("vecPoint", MakeVecArray(numPoints, randomNumbers));
  ds.AddCellField("vecCell", MakeVecArray(numCells, randomNumbers));
}

std::vector<vtkm::cont::PartitionedDataSet> MakeUniformData(const std::vector<int>& numBlocks,
                                                            int rank,
                                                            RandomNumbers& randomNumbers)
{
  int totalNumBlocks = std::accumulate(numBlocks.begin(), numBlocks.end(), 0);

  std::vector<vtkm::cont::DataSet> dataSets;
  for (int i = 0; i < totalNumBlocks; i++)
  {
    vtkm::Vec3f spacing, origin;
    vtkm::Id3 dim;

    for (int j = 0; j < 3; j++)
    {
      spacing[j] = randomNumbers.RandomFloat(0.1, 10.);
      origin[j] = randomNumbers.RandomFloat(-100, 100);
      dim[j] = randomNumbers.RandomVal<int>(2, 10);
    }

    auto ds = vtkm::cont::DataSetBuilderUniform::Create(dim, origin, spacing);
    AddFields(ds, randomNumbers);
    dataSets.push_back(ds);
  }

  std::vector<vtkm::cont::PartitionedDataSet> ret;
  ret.push_back(DataSetsByRank(dataSets, rank, numBlocks));
  return ret;
}

int main(int argc, char** argv)
{
#ifdef FIDES_USE_MPI
  MPI_Init(&argc, &argv);
#endif

  int retVal = 0;
  RandomNumbers randomNumbers;

  std::cout << __LINE__ << std::endl;
  std::string fname = "out.bp";
  fides::io::DataSetAppendWriter writer(fname);
  std::cout << __LINE__ << std::endl;

  int step = 0;
  while (true)
  {
    auto pds = MakeUniformData({ 1 }, 0, randomNumbers);
    std::cout << "Write step: " << step << std::endl;
    //pds[0].PrintSummary(std::cout);
    //writer.Write(pds[0], "BPFile");
    writer.Write(pds[0], "SST");
    step++;
    sleep(2);
  }

  writer.Close();

#ifdef FIDES_USE_MPI
  MPI_Finalize();
#endif

  return retVal;
}
