#include <openvdb/openvdb.h>
#include <openvdb/tools/LevelSetSphere.h>
#include <openvdb/tools/LevelSetUtil.h>

#include <filesystem>
#include <iostream>
#include <string>

namespace fs = std::filesystem;

static void write_grid_file(const fs::path &path, const openvdb::GridPtrVec &grids)
{
  openvdb::io::File file(path.string());
  file.write(grids);
  file.close();
}

int main(int argc, char **argv)
{
  if (argc != 2) {
    std::cerr << "usage: vdb_fixture_generator OUTPUT_DIRECTORY\n";
    return 2;
  }

  try {
    openvdb::initialize();
    const fs::path output_dir = fs::absolute(argv[1]);
    fs::create_directories(output_dir);

    constexpr float voxel_size = 0.2f;
    auto surface = openvdb::tools::createLevelSetSphere<openvdb::FloatGrid>(
        3.0f, openvdb::Vec3f(0.0f), voxel_size, 3.0f, false);
    surface->setName("surface");

    auto density = surface->deepCopy();
    openvdb::tools::sdfToFogVolume(*density);
    density->setName("density");
    density->setGridClass(openvdb::GRID_FOG_VOLUME);

    auto temperature = openvdb::FloatGrid::create(0.0f);
    auto color = openvdb::Vec3SGrid::create(openvdb::Vec3f(0.0f));
    auto velocity = openvdb::Vec3SGrid::create(openvdb::Vec3f(0.0f));
    temperature->setTransform(density->transform().copy());
    color->setTransform(density->transform().copy());
    velocity->setTransform(density->transform().copy());
    temperature->setName("temperature");
    color->setName("color");
    velocity->setName("velocity");
    temperature->setGridClass(openvdb::GRID_FOG_VOLUME);
    color->setGridClass(openvdb::GRID_FOG_VOLUME);
    velocity->setGridClass(openvdb::GRID_STAGGERED);

    const auto density_accessor = density->getConstAccessor();
    auto temperature_accessor = temperature->getAccessor();
    auto color_accessor = color->getAccessor();
    auto velocity_accessor = velocity->getAccessor();
    for (int z = -18; z <= 18; ++z) {
      for (int y = -18; y <= 18; ++y) {
        for (int x = -18; x <= 18; ++x) {
          const openvdb::Coord coord(x, y, z);
          const float value = density_accessor.getValue(coord);
          if (value <= 0.0f) {
            continue;
          }
          temperature_accessor.setValueOn(coord, 300.0f + 1200.0f * value);
          color_accessor.setValueOn(
              coord,
              openvdb::Vec3f((x + 18.0f) / 36.0f, (y + 18.0f) / 36.0f, (z + 18.0f) / 36.0f));
          velocity_accessor.setValueOn(
              coord, openvdb::Vec3f(-0.02f * y, 0.02f * x, 0.1f * value));
        }
      }
    }

    write_grid_file(output_dir / "generated-smoke.vdb", {density, temperature, color, velocity});
    write_grid_file(output_dir / "generated-level-set.vdb", {surface});

    auto sparse = openvdb::FloatGrid::create(0.0f);
    sparse->setName("density");
    sparse->setGridClass(openvdb::GRID_FOG_VOLUME);
    auto sparse_accessor = sparse->getAccessor();
    sparse_accessor.setValueOn(openvdb::Coord(-100000, -100000, -100000), 0.25f);
    sparse_accessor.setValueOn(openvdb::Coord(0, 0, 0), 1.0f);
    sparse_accessor.setValueOn(openvdb::Coord(100000, 100000, 100000), 0.5f);
    write_grid_file(output_dir / "generated-large-bounds-sparse.vdb", {sparse});

    std::cout << "generated-vdb-fixtures output=" << output_dir.string()
              << " openvdb=" << openvdb::getLibraryVersionString() << " files=3\n";
    openvdb::uninitialize();
    return 0;
  }
  catch (const std::exception &error) {
    std::cerr << "VDB_FIXTURE_GENERATION_FAILED: " << error.what() << '\n';
    return 1;
  }
}