#include <boost/program_options.hpp>
#include <exodusII.h>
#include <mpi.h>

#include <iostream>
#include <iomanip>
#include <limits>
#include <sstream>
#include <vector>

#include "decomp.h"

namespace po = boost::program_options;

using integer_t = int;
using unsigned_integer_t = std::size_t;
using real_t = double;

template <typename T, typename Container>
auto write(
  int exo_id,
  ex_entity_id blk_id,
  const Container & cell_vertices,
  unsigned_integer_t num_verts_per_cell
) {

  std::string desc;

  switch(num_verts_per_cell) {
    case 4:
      desc = "quad4";
      break;
    case 8:
      desc = "hex8";
      break;
    default:
      std::cerr
        << "Unknown element type with " << num_verts_per_cell << " verts"
        << std::endl;
      return -1;
  }

  auto num_elem = cell_vertices.size() / num_verts_per_cell;
  auto status = ex_put_block( exo_id, EX_ELEM_BLOCK, blk_id, desc.c_str(),
    num_elem, num_verts_per_cell, 0, 0, 0 );
  if (status) return status;

  status = ex_put_name(exo_id, EX_ELEM_BLOCK, blk_id, "element block");
  if (status) return status;

  if ( std::is_same<typename Container::value_type, T>::value )
    status = ex_put_elem_conn( exo_id, blk_id, cell_vertices.data() );
  else {
    std::vector<T> vec(cell_vertices.begin(), cell_vertices.end());
    status = ex_put_elem_conn( exo_id, blk_id, vec.data() );
  }

  return status;
}

template <typename T, typename Container>
auto write_mapping(
  int exo_id,
  const Container & cell_mapping,
  const Container & vert_mapping
) {
  int status;
  if ( std::is_same<typename Container::value_type, T>::value ) {
    status = ex_put_elem_num_map( exo_id, cell_mapping.data() );
    if (status) return status;
    status = ex_put_node_num_map( exo_id, vert_mapping.data() );
    if (status) return status;
  }
  else {
    std::vector<T> vec(cell_mapping.begin(), cell_mapping.end());
    status = ex_put_elem_num_map( exo_id, vec.data() );
    if (status) return status;
    vec.assign(vert_mapping.begin(), vert_mapping.end());
    status = ex_put_node_num_map( exo_id, vec.data() );
    if (status) return status;
  }
  return status;
}

template <typename T, typename Container>
auto write_side_set(
    int ex_id,
    int side_id,
    const Container & elem_list,
    const Container & side_list)
{
  auto status = ex_put_side_set_param(ex_id, side_id, side_list.size(), 0);
  if (status) return status;

  if ( std::is_same<typename Container::value_type, T>::value ) {
    return ex_put_side_set(ex_id, side_id, elem_list.data(), side_list.data());
  }
  else {
    std::vector<T> elem(elem_list.begin(), elem_list.end());
    std::vector<T> side(side_list.begin(), side_list.end());
    return ex_put_side_set(ex_id, side_id, elem.data(), side.data());
  }
}

void subdivide(
    unsigned_integer_t nelem,
    unsigned_integer_t npart,
    std::vector<unsigned_integer_t> & dist )
{

  size_t quot = nelem / npart;
  size_t rem = nelem % npart;

  dist.clear();
  dist.reserve(npart + 1);
  dist.push_back(0);

  // Set the distributions for each rank. This happens on all ranks.
  // Each rank gets the average number of indices, with higher ranks
  // getting an additional index for non-zero remainders.
  for(size_t r(0); r < npart; ++r) {
    const size_t indices = quot + ((r >= (npart - rem)) ? 1 : 0);
    dist.push_back(dist[r] + indices);
  } // for
};

struct side_t {
  std::vector<unsigned_integer_t> elem_list;
  std::vector<unsigned_integer_t> side_list;
};

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

  bool is_verbose = false;
  bool force_int64 = false;

  MPI_Init(&argc, &argv);

  int comm_rank, comm_size;
  MPI_Comm_size( MPI_COMM_WORLD, &comm_size );
  MPI_Comm_rank( MPI_COMM_WORLD, &comm_rank );

  // Declare the supported options.
  po::options_description desc("Allowed options");
  desc.add_options()
      ("help", "produce help message")
      ("dimensions", po::value<std::vector<integer_t>>()->multitoken(), "Box dimensions.")
      ("partitions", po::value<std::vector<integer_t>>()->multitoken(), "Number partitions in each direction.")
      ("lower", po::value<std::vector<real_t>>()->multitoken(), "Box lower coordinates.")
      ("upper", po::value<std::vector<real_t>>()->multitoken(), "Box upper coordinates.")
      ("output-file", po::value<std::string>(), "output file")
      ("verbose", "Print extra debug info.")
      ("large-integer", "Use large integer support.")
  ;

  po::positional_options_description p;
  p.add("output-file,o", -1);

  po::variables_map vm;
  auto extras = po::command_line_style::unix_style ^ po::command_line_style::allow_short;
  po::store( po::command_line_parser(argc, argv).style(extras).options(desc).positional(p).run(), vm);

  if (!vm.count("dimensions")) {
    if (comm_rank == 0) {
      std::cout << "No dimensions provided!" << std::endl;
      std::cout << desc << "\n";
    }
    MPI_Finalize();
    return 1;
  }

  if (!vm.count("output-file")) {
    if (comm_rank == 0) {
      std::cout << "No output file provided!" << std::endl;
      std::cout << desc << "\n";
    }
    MPI_Finalize();
    return 1;
  }

  if (vm.count("help")) {
    if (comm_rank == 0) {
      std::cout << desc << "\n";
    }
    MPI_Finalize();
    return 1;
  }

  po::notify(vm);

  bool print_help = false;

  auto name = vm["output-file"].as<std::string>();

  const auto & args_dims = vm["dimensions"].as<std::vector<integer_t>>();
  std::vector<unsigned_integer_t> global_dims( args_dims.begin(), args_dims.end() );
  auto num_dims = global_dims.size();

  std::vector<real_t> global_lower(num_dims, -0.5);
  std::vector<real_t> global_upper(num_dims,  0.5);

  if (num_dims > 3 || num_dims <= 1) {
    if (comm_rank == 0) {
      std::cout << "Only two to three dimensions supported!" << std::endl;
    }
    print_help = true;
  }

  if (vm.count("lower")) {
    const auto & args = vm["lower"].as<std::vector<real_t>>();
    auto n = args.size();
    if ( n != num_dims ) {
      if (comm_rank == 0) {
        std::cout
          << "Number of lower box coordinates must match the number of "
          << "dimensions, " << n << " provided, " << num_dims << " expected."
          << std::endl;
      }
      print_help = true;
    }
    global_lower = args;
  }

  if (vm.count("upper")) {
    const auto & args = vm["upper"].as<std::vector<real_t>>();
    auto n = args.size();
    if ( n != num_dims ) {
      if (comm_rank == 0) {
        std::cout
          << "Number of upper box coordinates must match the number of "
          << "dimensions, " << n << " provided, " << num_dims << " expected."
          << std::endl;
      }
      print_help = true;
    }
    global_upper = args;
  }

  if (print_help) {
    if (comm_rank == 0) {
      std::cout << desc << "\n";
    }
    MPI_Finalize();
    return -1;
  }

  std::vector<unsigned_integer_t> block_sizes(num_dims, 1);

  if (vm.count("partitions")) {
    if (comm_rank==0) std::cout << "Using provided partitioning" << std::endl;
    const auto & args = vm["partitions"].as<std::vector<int>>();
    auto n = args.size();
    if ( n != num_dims ) {
      if ( comm_rank == 0 ) {
        std::cout
          << "Number of partition must match the number of "
          << "dimensions, " << n << " privided, " << num_dims << " expected."
          << std::endl;
      }
      MPI_Finalize();
      return -1;
    }
    unsigned_integer_t total_partitions = 1;
    for ( unsigned_integer_t i=0; i<num_dims; ++i ) {
      block_sizes[i] = args[i];
      total_partitions *= args[i];
    }
    if ( total_partitions != comm_size ) {
      if ( comm_rank == 0 ) {
        std::cout
          << "Total number of partitions does not match the total number of "
          << " mpi ranks" << std::endl;
      }
      MPI_Finalize();
      return -1;
    }
  }
  else {
    if (num_dims == 2) {
      auto decomp = cartmesh::decomp<2>({global_dims[0], global_dims[1]}, comm_size);
      for (unsigned_integer_t i = 0; i < 2; i++)
        block_sizes[i] = decomp[i];
    } else if (num_dims == 3) {
      auto decomp = cartmesh::decomp<3>({global_dims[0], global_dims[1], global_dims[2]}, comm_size);
      for (unsigned_integer_t i = 0; i < 3; i++)
        block_sizes[i] = decomp[i];
    } else {
      block_sizes[0] = comm_size;
    }

    if (comm_rank==0) {
      std::cout << "Automatically partitioned block sizes: ";
      for ( auto i : block_sizes ) std::cout << i << " ";
      std::cout << std::endl;
    }

  } // partitioning

  if (vm.count("large-integer")) {
    force_int64 = true;
    if (comm_rank == 0) {
      std::cout << "Forcing 64 bit integers" << std::endl;
    }
  }

  if (vm.count("verbose")) {
    is_verbose = true;
    if (comm_rank == 0) {
      std::cout << "Adding extra exodus debug." << std::endl;
    }
    ex_opts(EX_ABORT | EX_VERBOSE);
  }

  if (comm_rank == 0) {
    std::cout << num_dims << "-dimensional mesh selected." << std::endl;
  }

  // build a mesh
  std::vector<real_t> delta(num_dims);
  for ( unsigned_integer_t i=0; i<num_dims; ++i ) {
    if ( global_upper[i] < global_lower[i] ) {
      if (comm_rank == 0) {
        std::cout
          << "Bounding box is invalid: for i=" << i << ", " << global_lower[i]
          << " should be less than " << global_upper[i]
          << std::endl;
      }
      MPI_Finalize();
      return -1;
    }
    delta[i] =  (global_upper[i] - global_lower[i]) / global_dims[i];
  }

  std::vector<unsigned_integer_t> block_ijk(num_dims);
  auto temp = comm_rank;
  for ( int i=0; i<num_dims-1; ++i ) {
    block_ijk[i] = temp % block_sizes[i];
    block_ijk[i+1] = (temp - block_ijk[i]) / block_sizes[i];
    temp -= block_ijk[i];
    temp /= block_sizes[i];
  }

  if (is_verbose) {
    std::cout << std::flush;
    MPI_Barrier(MPI_COMM_WORLD);
    for ( unsigned_integer_t i=0; i<comm_size; ++i )
    {
      if ( comm_rank == i ) {
        std::cout << "Rank " << comm_rank << " has block index: ";
        for ( unsigned_integer_t j=0; j<num_dims; ++j ) std::cout << block_ijk[j] << " ";
        std::cout << std::endl;
      }
      MPI_Barrier(MPI_COMM_WORLD);
    }
  }

  std::vector< std::vector<unsigned_integer_t> > block_displ(num_dims);
  for ( int i=0; i<num_dims; ++i )
    subdivide( global_dims[i], block_sizes[i], block_displ[i] );

  unsigned_integer_t num_verts = 1;
  unsigned_integer_t num_cells = 1;
  unsigned_integer_t global_verts = 1;
  unsigned_integer_t global_cells = 1;
  unsigned_integer_t num_verts_per_cell = 1;

  std::vector<unsigned_integer_t> local_dims(num_dims);
  std::vector<real_t> local_lower(num_dims);
  std::vector<real_t> local_upper(num_dims);

  for ( unsigned_integer_t i=0; i<num_dims; ++i ) {
    global_verts *= global_dims[i]+1;
    global_cells *= global_dims[i];
    num_verts_per_cell *= 2;
    auto jblock = block_ijk[i];
    const auto & displ = block_displ[i];
    local_dims[i] = displ[jblock+1] - displ[jblock];
    num_cells *= local_dims[i];
    num_verts *= local_dims[i]+1;
    local_lower[i] = global_lower[i] + (displ[jblock])*delta[i] ;
    local_upper[i] = global_lower[i] + (displ[jblock+1])*delta[i];
  }

  MPI_Datatype mpi_type_t;
  if ( sizeof(unsigned_integer_t) == sizeof(unsigned int) )
    mpi_type_t = MPI_UNSIGNED;
  else if ( sizeof(unsigned_integer_t) == sizeof(unsigned long long) )
    mpi_type_t = MPI_UNSIGNED_LONG_LONG;
  else {
    if (comm_rank == 0) std::cout << "UNKNOWN MPI DATA TYPE." << std::endl;
    MPI_Finalize();
    return -1;
  }

  unsigned_integer_t temp_cells;
  MPI_Allreduce( &num_cells, &temp_cells, 1, mpi_type_t, MPI_SUM, MPI_COMM_WORLD );
  if ( temp_cells != global_cells ) {
    if ( comm_rank == 0) {
      std::cout << "Summed cell count should be " << global_cells
        << " got " << temp_cells << std::endl;
    }
    MPI_Finalize();
    return -1;
  }

  std::vector<real_t> coordinates(num_dims * num_verts);
  std::vector<unsigned_integer_t> cell_vertices(num_cells*num_verts_per_cell);
  std::vector<unsigned_integer_t> global_cell_id(num_cells);
  std::vector<unsigned_integer_t> global_vertex_id(num_verts);

  std::map<int,side_t> side_sets;

  if ( num_dims == 2 ) {

    auto vertex_local_id = [&](auto i, auto j) {
      return i + (local_dims[0]+1)*j;
    };
    auto vertex_global_id = [&](auto i, auto j) {
      auto iblk = block_ijk[0];
      auto istart = block_displ[0][ iblk ];
      auto iglobal = istart + i;

      auto jblk = block_ijk[1];
      auto jstart = block_displ[1][ jblk ];
      auto jglobal = jstart + j;
      return iglobal + (global_dims[0]+1)*jglobal;
    };

    for ( unsigned_integer_t j=0; j<local_dims[1]+1; ++j ) {
      for ( unsigned_integer_t i=0; i<local_dims[0]+1; ++i ) {
        auto id = vertex_local_id(i, j);
        coordinates[id              ] = local_lower[0] + i*delta[0];
        coordinates[id +   num_verts] = local_lower[1] + j*delta[1];
        global_vertex_id[id] = vertex_global_id(i,j) + 1;
      }
    }

    auto cell_local_id = [&](auto i, auto j) {
      return i + local_dims[0]*j;
    };
    auto cell_global_id = [&](auto i, auto j) {
      auto iblk = block_ijk[0];
      auto istart = block_displ[0][ iblk ];
      auto iglobal = istart + i;

      auto jblk = block_ijk[1];
      auto jstart = block_displ[1][ jblk ];
      auto jglobal = jstart + j;
      return iglobal + global_dims[0]*jglobal;
    };

    auto cell_is_bottom_boundary = [&](auto i, auto j)
    { return block_displ[1][block_ijk[1]] + j == 0; };
    auto cell_is_top_boundary = [&](auto i, auto j)
    { return block_displ[1][block_ijk[1]] + j == global_dims[1]-1; };

    auto cell_is_left_boundary = [&](auto i, auto j)
    { return block_displ[0][block_ijk[0]] + i == 0; };
    auto cell_is_right_boundary = [&](auto i, auto j)
    { return block_displ[0][block_ijk[0]] + i == global_dims[0]-1; };

    for ( unsigned_integer_t j=0, id=0; j<local_dims[1]; ++j ) {
      for ( unsigned_integer_t i=0; i<local_dims[0]; ++i ) {
        auto local_id = cell_local_id(i, j);
        global_cell_id[local_id] = cell_global_id(i, j) + 1;
        cell_vertices[id++] = vertex_local_id(i  ,   j) + 1;
        cell_vertices[id++] = vertex_local_id(i+1,   j) + 1;
        cell_vertices[id++] = vertex_local_id(i+1, j+1) + 1;
        cell_vertices[id++] = vertex_local_id(i  , j+1) + 1;
        if (cell_is_bottom_boundary(i, j)) {
          side_sets[3].elem_list.emplace_back(local_id+1);
          side_sets[3].side_list.emplace_back(1);
        }
        if (cell_is_right_boundary(i, j)) {
          side_sets[2].elem_list.emplace_back(local_id+1);
          side_sets[2].side_list.emplace_back(2);
        }
        if (cell_is_top_boundary(i, j)) {
          side_sets[4].elem_list.emplace_back(local_id+1);
          side_sets[4].side_list.emplace_back(3);
        }
        if (cell_is_left_boundary(i, j)) {
          side_sets[1].elem_list.emplace_back(local_id+1);
          side_sets[1].side_list.emplace_back(4);
        }
      }
    }

  } // two-dimensional

  else if ( num_dims == 3 ) {

    auto vertex_local_id = [&](auto i, auto j, auto k) {
      return i + (local_dims[0]+1)*(j + (local_dims[1]+1)*k);
    };
    auto vertex_global_id = [&](auto i, auto j, auto k) {
      auto iblk = block_ijk[0];
      auto istart = block_displ[0][ iblk ];
      auto iglobal = istart + i;

      auto jblk = block_ijk[1];
      auto jstart = block_displ[1][ jblk ];
      auto jglobal = jstart + j;

      auto kblk = block_ijk[2];
      auto kstart = block_displ[2][ kblk ];
      auto kglobal = kstart + k;
      return iglobal + (global_dims[0]+1)*(jglobal + (global_dims[1]+1)*kglobal);
    };

    auto cell_local_id = [&](auto i, auto j, auto k) {
      return i + local_dims[0]*(j + local_dims[1]*k);
    };
    auto cell_global_id = [&](auto i, auto j, auto k) {
      auto iblk = block_ijk[0];
      auto istart = block_displ[0][ iblk ];
      auto iglobal = istart + i;

      auto jblk = block_ijk[1];
      auto jstart = block_displ[1][ jblk ];
      auto jglobal = jstart + j;

      auto kblk = block_ijk[2];
      auto kstart = block_displ[2][ kblk ];
      auto kglobal = kstart + k;
      return iglobal + global_dims[0]*(jglobal + global_dims[1]*kglobal);
    };

    auto cell_is_left_boundary = [&](auto i, auto j, auto k)
    { return block_displ[0][block_ijk[0]] + i == 0; };
    auto cell_is_right_boundary = [&](auto i, auto j, auto k)
    { return block_displ[0][block_ijk[0]] + i == global_dims[0]-1; };

    auto cell_is_bottom_boundary = [&](auto i, auto j, auto k)
    { return block_displ[1][block_ijk[1]] + j == 0; };
    auto cell_is_top_boundary = [&](auto i, auto j, auto k)
    { return block_displ[1][block_ijk[1]] + j == global_dims[1]-1; };

    auto cell_is_front_boundary = [&](auto i, auto j, auto k)
    { return block_displ[2][block_ijk[2]] + k == 0; };
    auto cell_is_back_boundary = [&](auto i, auto j, auto k)
    { return block_displ[2][block_ijk[2]] + k == global_dims[2]-1; };

    for ( unsigned_integer_t k=0; k<local_dims[2]+1; ++k ) {
      for ( unsigned_integer_t j=0; j<local_dims[1]+1; ++j ) {
        for ( unsigned_integer_t i=0; i<local_dims[0]+1; ++i ) {
          auto id = vertex_local_id(i, j, k);
          coordinates[id              ] = local_lower[0] + i*delta[0];
          coordinates[id +   num_verts] = local_lower[1] + j*delta[1];
          coordinates[id + 2*num_verts] = local_lower[2] + k*delta[2];
          global_vertex_id[id] = vertex_global_id(i,j,k) + 1;
        }
      }
    }

    for ( unsigned_integer_t k=0, id=0; k<local_dims[2]; ++k ) {
      for ( unsigned_integer_t j=0; j<local_dims[1]; ++j ) {
        for ( unsigned_integer_t i=0; i<local_dims[0]; ++i ) {
          auto local_id = cell_local_id(i, j, k);
          global_cell_id[local_id] = cell_global_id(i, j, k) + 1;
          cell_vertices[id++] = vertex_local_id(i  ,   j, k) + 1;
          cell_vertices[id++] = vertex_local_id(i+1,   j, k) + 1;
          cell_vertices[id++] = vertex_local_id(i+1, j+1, k) + 1;
          cell_vertices[id++] = vertex_local_id(i  , j+1, k) + 1;

          cell_vertices[id++] = vertex_local_id(i  ,   j, k+1) + 1;
          cell_vertices[id++] = vertex_local_id(i+1,   j, k+1) + 1;
          cell_vertices[id++] = vertex_local_id(i+1, j+1, k+1) + 1;
          cell_vertices[id++] = vertex_local_id(i  , j+1, k+1) + 1;

          if (cell_is_bottom_boundary(i, j, k)) {
            side_sets[1].elem_list.emplace_back(local_id+1);
            side_sets[1].side_list.emplace_back(1);
          }
          if (cell_is_right_boundary(i, j, k)) {
            side_sets[2].elem_list.emplace_back(local_id+1);
            side_sets[2].side_list.emplace_back(2);
          }
          if (cell_is_top_boundary(i, j, k)) {
            side_sets[3].elem_list.emplace_back(local_id+1);
            side_sets[3].side_list.emplace_back(3);
          }
          if (cell_is_left_boundary(i, j, k)) {
            side_sets[4].elem_list.emplace_back(local_id+1);
            side_sets[4].side_list.emplace_back(4);
          }
          if (cell_is_front_boundary(i, j, k)) {
            side_sets[5].elem_list.emplace_back(local_id+1);
            side_sets[5].side_list.emplace_back(5);
          }
          if (cell_is_back_boundary(i, j, k)) {
            side_sets[6].elem_list.emplace_back(local_id+1);
            side_sets[6].side_list.emplace_back(6);
          }
        }
      }
    }

  } // three-dimensional

  // check if we need 64 bit integers
  constexpr auto max_int = std::numeric_limits<int>::max();
  auto is_int64 =
    force_int64 ||
    (cell_vertices.size() >= max_int/2) ||
    global_verts >= max_int ||
    global_cells >= max_int;

  // size of floating point variables used in app.
  int app_word_size = sizeof(real_t);

  // size of floating point to be stored in file.
  // change to float to save space
  int exo_word_size = sizeof(real_t);

  // determine the file creation mode
  auto cmode = EX_CLOBBER;
  if (is_int64) {
    cmode |= EX_ALL_INT64_DB | EX_ALL_INT64_API;
    if (comm_rank==0) std::cout << "Using 64-bit integers." << std::endl;
  }
  else
    if (comm_rank==0) std::cout << "Using 32-bit integers." << std::endl;

  // figure out mesh file name
  std::string filename;
  if (comm_size == 1 )
    filename = name;
  else {
    auto number = comm_size;
    unsigned int num_digits = 1;
    while (number) {
      number /= 10;
      num_digits++;
    }
    std::stringstream ss;
    ss
      << name
      << "."
      << std::setfill('0') << std::setw(num_digits) << comm_size
      << "."
      << std::setfill('0') << std::setw(num_digits) << comm_rank;
    filename = ss.str();
  }

  // create file
  auto exo_id =
      ex_create(filename.c_str(), cmode, &app_word_size, &exo_word_size);
  if (exo_id < 0) {
    if ( comm_rank == 0 ) {
      std::cerr << "Problem writing exodus file, ex_create() returned " << exo_id << std::endl;
    }
    MPI_Finalize();
    return exo_id;
  }
  else {
    if (comm_rank==0) std::cout << "Opened file for writing: " << name << std::endl;
  }

  ex_init_params exopar;
  std::strcpy(exopar.title, "Exodus II output from make_mesh.");
  exopar.num_dim = num_dims;
  exopar.num_nodes = num_verts;
  exopar.num_edge = 0;
  exopar.num_edge_blk = 0;
  exopar.num_face = 0;
  exopar.num_face_blk = 0;
  exopar.num_elem = num_cells;
  exopar.num_elem_blk = 1;
  exopar.num_node_sets = 0;
  exopar.num_edge_sets = 0;
  exopar.num_face_sets = 0;
  exopar.num_side_sets = side_sets.size();
  exopar.num_elem_sets = 0;
  exopar.num_node_maps = 0;
  exopar.num_edge_maps = 0;
  exopar.num_face_maps = 0;
  exopar.num_elem_maps = 0;

  auto status = ex_put_init_ext(exo_id, &exopar);
  if (status) {
    if (comm_rank==0)
      std::cerr << "Problem putting exodus file parameters, ex_put_init_ext() returned "
          << status << std::endl;
    MPI_Finalize();
    return status;
  }

  // exodus is kind enough to fetch the data in the real type we ask for
  status = ex_put_coord(
      exo_id, coordinates.data(), coordinates.data() + num_verts,
      coordinates.data() + 2 * num_verts);

  if (status) {
    if (comm_rank==0)
      std::cerr
          << "Problem putting vertex coordinates to exodus file, "
          << " ex_put_coord() returned " << status << std::endl;
    MPI_Finalize();
    return status;
  }
  else {
    if (comm_rank==0)
      std::cout << "Wrote " << num_verts << " vertex coordinates" << std::endl;
  }

  const char *coord_names[3] = {"x", "y", "z"};
  status = ex_put_coord_names (exo_id, const_cast<char**>(coord_names));
  if (status) {
    if (comm_rank==0)
      std::cerr
          << "Problem putting coordinate names to exodus file, "
          << " ex_put_coord_names() returned " << status << std::endl;
    MPI_Finalize();
    return status;
  }

  ex_entity_id elem_blk_id = 1;
  if (is_int64)
    status = write<long long>(exo_id, elem_blk_id, cell_vertices, num_verts_per_cell);
  else
    status = write<int>(exo_id, elem_blk_id, cell_vertices, num_verts_per_cell);
  if (status) {
    if (comm_rank==0)
      std::cerr
          << "Problem putting elements to exodus file, "
          << " ex_put_block() returned " << status << std::endl;
    MPI_Finalize();
    return status;
  }
  else {
    if (comm_rank==0) std::cout << "Wrote " << num_cells << " cells" << std::endl;
  }

  if (is_int64)
    status = write_mapping<long long>(exo_id, global_cell_id, global_vertex_id);
  else
    status = write_mapping<int>(exo_id, global_cell_id, global_vertex_id);
  if (status) {
    if (comm_rank==0)
      std::cerr
          << "Problem putting mapping to exodus file, "
          << " ex_put_num_map() returned " << status << std::endl;
    MPI_Finalize();
    return status;
  }
  else {
    if (comm_rank==0) std::cout << "Wrote mapping." << std::endl;
  }

  for (const auto & ss : side_sets) {
    const auto & elem_list = ss.second.elem_list;
    const auto & side_list = ss.second.side_list;
    auto sid = ss.first;
    if (is_int64)
      status = write_side_set<long long> (exo_id, sid, elem_list, side_list);
    else
      status = write_side_set<int> (exo_id, sid, elem_list, side_list);
    if (status) {
      if (comm_rank==0)
        std::cerr
            << "Problem putting side sets to exodus file, "
            << " ex_put_side_set() returned " << status << std::endl;
      MPI_Finalize();
      return status;
    }
    sid++;
  }

  status = ex_close(exo_id);
  if (status) {
    if (comm_rank==0)
      std::cerr
        << "Problem closing exodus file, ex_close() returned " << exo_id
        << std::endl;
    MPI_Finalize();
    return status;
  }
  else {
    if (comm_rank==0) std::cout << "Closed exodus file!" << std::endl;
  }

  return MPI_Finalize();

}


#include "MeshRenderer.h"
#include <bgfx.h>
#include <bx/timer.h>

#include "StrideDataFixedArray.h"
#include "Elements/ElementsContainer.h"
#include "Elements/IElement.h"

#include <ocornut-imgui/imgui.h>
#include <imgui/imgui.h>

#include <Eigen/Dense>
#include <glm/glm.hpp>
#include <glm/matrix.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <glm/gtx/rotate_vector.hpp>
#include <glm/gtx/quaternion.hpp>

#include "Signs.h"
#include "VertexDeclarations.h"

#include "TransformUtils.h"

#include "MaterialManager.h"

#include "Mesh.h"

#include "ScaleBar.h"

std::string meshViewOptionTitle = "Mesh view options";

void MeshRenderer::Init()
{
	m_renderWireframe = true;
	m_renderNodes = true;
	m_showMeshViewOption = true;
	m_renderNodesNumbers = false;
	m_renderElementsNumbers = false;
	m_renderUndeformedMesh = true;
	m_alphaUndeformedMesh = 0.3f;
	m_deformScale = 50.0f;
	m_min = 0.0f;
	m_max = 1.0f;
	m_intervals = 10.0f;
	m_needResize = false;

	NodeVertex::init();
	LegendMesh::Init();
	LoadSign::Init();
	ConstSign::Init();

	m_materialManager.Init();
	m_materialManager.CreateDefaultMaterials();
}

void MeshRenderer::ShowMeshViewOptions()
{
	if (!m_showMeshViewOption)
	{
		m_showMeshViewOption = true;
		ImGui::SetWindowFocus(meshViewOptionTitle.c_str());
		ResetMeshViewOptionsWindowPosition(true);
	}
}

void MeshRenderer::ResetMeshViewOptionsWindowPosition(bool force)
{
	ImVec2 vpos;
	int width = 300;
	vpos.x = static_cast<float>(m_width - width);
	vpos.y = 21;
	ImGuiSetCond_ cond = force ? ImGuiSetCond_Always : ImGuiSetCond_Appearing;
	ImGui::SetWindowPos(meshViewOptionTitle.c_str(), vpos, cond);
	ImVec2 vsize;
	vsize.x = static_cast<float>(width);
	vsize.y = 300;
	ImGui::SetWindowSize(meshViewOptionTitle.c_str(), vsize, cond);
}

void MeshRenderer::Render(int width, int height, Camera& camera)
{
	RenderContext cntxt;
	cntxt.m_camera = &camera;
	cntxt.m_height = height;
	cntxt.m_width = width;
	cntxt.m_materialManager = &m_materialManager;
	cntxt.m_renderNodes = m_renderNodes;
	cntxt.m_renderWireframe = m_renderWireframe;

	m_width = width;

#ifdef WIN32
	_set_output_format(_TWO_DIGIT_EXPONENT);
#endif

	if (m_needResize)
	{
		float m_position[2];
		m_position[0] = (m_max_bound.x() + m_min_bound.x()) / 2.0f;
		m_position[1] = (m_max_bound.y() + m_min_bound.y()) / 2.0f;
		float fovx = m_max_bound.x() - m_min_bound.x();
		float fovy = m_max_bound.y() - m_min_bound.y();
		float fov = fovx > fovy ? fovx : fovy;
		camera.SetPosition(m_position);
		camera.SetFOV(fov * 1.5f);
		m_needResize = false;
	}

	if (m_showMeshViewOption)
	{
		if (ImGui::Begin(meshViewOptionTitle.c_str(), &m_showMeshViewOption))
		{
			ImGui::Text("Mesh and lables");
			ImGui::Checkbox("Render wireframe", &m_renderWireframe);
			ImGui::Checkbox("Render nodes", &m_renderNodes);
			ImGui::Checkbox("Render elements lables", &m_renderElementsNumbers);
			ImGui::Checkbox("Render nodes lables", &m_renderNodesNumbers);
			ImGui::Separator();
			ImGui::Text("Undeformed mesh");
			ImGui::Checkbox("Render undeformed mesh", &m_renderUndeformedMesh);
			ImGui::PushItemWidth(100);
			ImGui::SliderFloat("Undeformed mesh\n transparency", &m_alphaUndeformedMesh, 0.0f, 1.0f);
			ImGui::PopItemWidth();
			ImGui::Separator();
			ImGui::Text("Deformations");
			ImGui::PushItemWidth(280);
			ImGui::SliderFloat("", &m_deformScale, 0.0f, 1000.0f, "Deformations x%.0f", 3.0f);
			ImGui::PopItemWidth();
			ResetMeshViewOptionsWindowPosition(false);
		}
		ImGui::End();
	}

	m_materialManager.SetIntervals(m_min, m_max, m_intervals);
	m_materialManager.SetRenderPlotEnabled(false);

	m_materialManager.SetDeformsScale(m_deformScale);

	cntxt.m_alpha = 1.0f;
	cntxt.m_greyMode = false;
	cntxt.m_state = 0;

	for (int i = 0, l = m_meshesList.size(); i < l; ++i)
	{
		m_meshesList[i].Render(cntxt);
	}

	m_materialManager.SetDeformsScale(0);

	cntxt.m_alpha = m_alphaUndeformedMesh;
	cntxt.m_greyMode = true;
	cntxt.m_state = BGFX_STATE_BLEND_EQUATION_ADD | BGFX_STATE_BLEND_ALPHA;

	if (m_renderUndeformedMesh)
	{
		for (int i = 0, l = m_meshesList.size(); i < l; ++i)
		{
			if (m_meshesList[i].HasDeforms())
			{
				m_meshesList[i].Render(cntxt);
			}
		}
	}

	cntxt.m_alpha = 1.0f;
	cntxt.m_greyMode = false;
	cntxt.m_state = 0;
	RenderScaleBar(cntxt, m_max, m_min, m_intervals);

	/*
	for (int i = 0, l = m_meshesList.m_objectList.size(); i < l; ++i)
	{
		if (m_renderElementsNumbers)
		{
			RenderElementsNumbers(m_meshesList.m_objectList[i], width, height, camera);
		}
		RenderBC(m_meshesList.m_objectList[i]);
	}

	if (m_renderNodesNumbers)
	{
		RenderNodesNumbers(width, height, camera);
	}
	*/
}

void MeshRenderer::Clear()
{
	/*
	if (m_meshesList.m_vbh.idx != uint16_t(- 1) )
	{
		bgfx::destroyVertexBuffer(m_meshesList.m_vbh);
		m_meshesList.m_vbh.idx = -1;
	}
	m_meshesList.m_nodesCopy.Init(0, 0);
	for (std::vector<Object>::iterator it = m_meshesList.m_objectList.begin(); it != m_meshesList.m_objectList.end(); ++it)
	{
		bgfx::destroyIndexBuffer(it->m_ibhTriang);
		bgfx::destroyIndexBuffer(it->m_ibhLines);
		delete it->m_elementsCenters;
	}
	m_meshesList.m_objectList.clear();
	*/
}

/*
void MeshRenderer::RenderElementsNumbers(const Object& mesh, int width, int height, const Camera& camera)
{
	for (int j = 0; j < mesh.m_elementsCenters->GetCount(); ++j)
	{
		float x, y;
		TransformToWorldSpace(
			(*mesh.m_elementsCenters)(j, 0),
			(*mesh.m_elementsCenters)(j, 1),
			0.0f, width, height, camera, x, y);
	char label[32];
	sprintf(label, "%d", j);
	imguiDrawText(static_cast<int>(x), static_cast<int>(y), ImguiTextAlign::Right, label, 0xFFFFFF00);
	}
}

void MeshRenderer::RenderNodesNumbers(int width, int height, const Camera& camera)
{
	int dof = m_meshesList.m_nodesCopy.GetDof();
	for (int j = 0; j < m_meshesList.m_nodesCopy.GetCount(); ++j)
	{
		float x, y;
		if (m_meshesList.m_hasDeforms)
		{
			TransformToWorldSpace(
				m_meshesList.m_nodesCopy(j, 0) + m_meshesList.m_nodesDeformation[dof * j + 0] * m_deformScale,
				m_meshesList.m_nodesCopy(j, 1) + m_meshesList.m_nodesDeformation[dof * j + 1] * m_deformScale,
				dof == 3 ? (m_meshesList.m_nodesCopy(j, 2) + m_meshesList.m_nodesDeformation[dof * j + 2] * m_deformScale) : 0.0f,
				width, height, camera, x, y);
		}
		else
		{
			TransformToWorldSpace(
				m_meshesList.m_nodesCopy(j, 0),
				m_meshesList.m_nodesCopy(j, 1),
				dof == 3 ? m_meshesList.m_nodesCopy(j, 2) : 0.0f,
				width, height, camera, x, y);
		}
		char label[32];
		sprintf(label, "%d", j);
		imguiDrawText(static_cast<int>(x), static_cast<int>(y), ImguiTextAlign::Right, label, 0xFF00FFFF);
	}
}

void MeshRenderer::RenderBC(const Object& mesh)
{
	int dof = m_meshesList.m_nodesCopy.GetDof();
	{
		Eigen::Vector4f color(1.0, 0.0, 0.0, 1.0);
		bgfx::setUniform(u_color, color.data(), 1);
	}
	for (int j = 0; j < mesh.m_nodalForceList->GetCount(); ++j)
	{
		int n = (*mesh.m_nodalForceList)[j];
		float direction[3];
		direction[0] = (*mesh.m_nodalForceList)(j, 0);
		direction[1] = (*mesh.m_nodalForceList)(j, 1);
		direction[2] = dof == 3 ? ((*mesh.m_nodalForceList)(j, 2)): 0.0f;
		RenderSign(mesh, n, LoadSign::vbh, LoadSign::ibh, direction);
	}
	{
		Eigen::Vector4f color(1.0, 1.0, 1.0, 1.0);
		bgfx::setUniform(u_color, color.data(), 1);
	}
	for (NodalConstraintList::const_iterator it = mesh.m_nodalConstraintList.begin(); it != mesh.m_nodalConstraintList.end(); ++it)
	{
		int n = it->node;
		for (int i = 0; i < 3; ++i)
		{
			if (it->type & i)
			{
				float direction[3] = { 0, 0, 0 };
				switch (it->type & i)
				{
				case NodalConstraint::UX:
					direction[0] = 1;
					break;
				case NodalConstraint::UY:
					direction[1] = 1;
					break;
				case NodalConstraint::UZ:
					direction[2] = 1;
					break;
				}
				RenderSign(mesh, n, ConstSign::vbh, ConstSign::ibh, direction);
			}
		}
	}
}
*/

void MeshRenderer::SetRanges(float min, float max)
{
	m_min = min;
	m_max = max;
	if (m_max - m_min < (m_min + m_max) / 2.0f * 1e-3f)
	{
		m_max = m_min + (m_min + m_max) / 2.0f * 1e-3f;
	}
}

void MeshRenderer::SetIntervals(int intervals)
{
	m_intervals = intervals;
}


#if !defined(BALSA_EIGEN_SPAN_HPP)
#define BALSA_EIGEN_SPAN_HPP

#include <Eigen/Dense>
#include <span>
#include "balsa/eigen/concepts/shape_types.hpp"

namespace balsa::eigen {
// Returns a span type
template<concepts::PlainObjectBaseDerived Matrix>
auto eigen2span(Matrix &&A) {
    using Mat = std::decay_t<Matrix>;
    constexpr static int Rows = Mat::RowsAtCompileTime;
    constexpr static int Cols = Mat::ColsAtCompileTime;

    using Scalar = typename Mat::Scalar;
    using ConstCorrectScalar = std::conditional_t<std::is_const_v<typename std::remove_reference_t<Matrix>>, const Scalar, Scalar>;
    if constexpr (Rows == Eigen::Dynamic || Cols == Eigen::Dynamic) {
        return std::span<ConstCorrectScalar>(A.data(), A.size());
    } else {
        constexpr static int Size = Rows * Cols;
        return std::span<ConstCorrectScalar, Size>(A.data(), Size);
    }
}
}// namespace balsa::eigen
#endif


#ifndef ALEPH_TOPOLOGY_SIMPLICIAL_COMPLEX_HH__
#define ALEPH_TOPOLOGY_SIMPLICIAL_COMPLEX_HH__

#include <boost/multi_index_container.hpp>

#include <boost/multi_index/indexed_by.hpp>
#include <boost/multi_index/mem_fun.hpp>
#include <boost/multi_index/ordered_index.hpp>
#include <boost/multi_index/random_access_index.hpp>

#include <algorithm>
#include <functional>
#include <initializer_list>
#include <iterator>
#include <ostream>
#include <limits>
#include <set>
#include <type_traits>
#include <vector>

namespace aleph
{

namespace topology
{

template <class Simplex> class SimplicialComplex
{
public:

  // Typedefs and tags for the multi-index container --------------------

  struct index_t {};
  struct lexicographical_t {};
  struct dimension_t {};

  using simplex_container_t
    =  boost::multi_index::multi_index_container<
        Simplex,
        boost::multi_index::indexed_by<
          boost::multi_index::random_access< boost::multi_index::tag<index_t> >,
          boost::multi_index::ordered_unique< boost::multi_index::tag<lexicographical_t>, boost::multi_index::identity<Simplex> >,
          boost::multi_index::ordered_non_unique< boost::multi_index::tag<dimension_t>, boost::multi_index::const_mem_fun<Simplex, std::size_t, &Simplex::dimension>        >
       >
  >;

  using const_iterator                 = typename simplex_container_t::template index<index_t>::type::const_iterator;
  using iterator                       = typename simplex_container_t::template index<index_t>::type::iterator;

  using lexicographical_iterator       = typename simplex_container_t::template index<lexicographical_t>::type::const_iterator;
  using const_lexicographical_iterator = typename simplex_container_t::template index<lexicographical_t>::type::iterator;

  using const_dimension_iterator       = typename simplex_container_t::template index<dimension_t>::type::const_iterator;
  using dimension_iterator             = typename simplex_container_t::template index<dimension_t>::type::iterator;

  // STL-like typedefs -------------------------------------------------

  // This simplifies writing algorithms that do not have any internal knowledge
  // about objects stored in this container.
  using value_type = Simplex;
  using ValueType  = value_type;

  // Constructors ------------------------------------------------------

  /** Creates an empty simplicial complex. */
  SimplicialComplex()
  {
  }

  /**
    Creates a simplicial complex from an initializer list of simplices.

    @param simplices Simplices to insert into the simplicial complex
  */

  SimplicialComplex( std::initializer_list<Simplex> simplices )
    : _simplices( simplices.begin(), simplices.end() )
  {
  }

  /**
    Creates a simplicial complex from a given range of simplices. For
    flexibility reasons, the range type is left unspecified. Note that the
    complex will store _copies_ of the simplices for better memory management.

    @param begin  Iterator pointing to begin of range
    @param end    Iterator pointing to end of range
  */

  template <class InputIterator> SimplicialComplex( InputIterator begin, InputIterator end )
    : _simplices( begin, end )
  {
  }

  // Simplex container modification ------------------------------------

  /** Clears the simplicial complex and removes all its simplices. */
  void clear()
  {
    _simplices.clear();
  }

  /**
    Given a range of simplices represented by two arbitrary input iterators,
    inserts the simplices into the simplicial complex.

    @param begin Iterator to begin of input range
    @param end   Iterator to end of input range
  */

  template <class InputIterator> void insert( InputIterator begin, InputIterator end )
  {
    _simplices.insert( _simplices.end(),
                       begin, end );
  }

  /**
    Inserts a new simplex into the simplicial complex. Note that the simplex is
    appended to the current filtration order, so the simplicial complex should
    be sorted again afterwards.

    @param simplex Simplex to insert into simplicial complex

    @warning If push_back() is called repeatedly, its performance will be worse
    than simply calling insert() for a range.
  */

  void push_back( const Simplex& simplex )
  {
    _simplices.push_back( simplex );
  }

  /**
    Rearranges the simplices in the simplicial complex using an external view
    that contains each element exactly once. This function is useful when being
    applied to a view that has been modified by some STL algorithm (as the
    container only provides \i const iterators).

    @param first Input iterator to the beginning of the view
  */

  template <typename InputIterator> void rearrange( InputIterator first )
  {
    _simplices.rearrange( first );
  }

  /**
    Replaces a simplex stored in the simplicial complex (described by an
    iterator) by another simplex. Note that this function is the only way
    for modifying a fixed simplicial complex.

    @param position Iterator describing the simplex that is to be replaced
    @param simplex  Simplex to replace the simplex with

    @returns true if the replacement took place, else false.
  */

  bool replace( iterator position, const Simplex& simplex )
  {
    return _simplices.replace( position, simplex );
  }

  // Simplex container access ------------------------------------------

  /**
    @returns Iterator to begin of simplices in current filtration order. A
    filtration may be applied by sorting the simplicial complex.

    @see SimplicialComplex::sort()
  */

  iterator begin()
  {
    return _simplices.template get<index_t>().begin();
  }

  /** @overload begin() */
  const_iterator begin() const
  {
    return _simplices.template get<index_t>().begin();
  }

  /**
    @returns Iterator to end of simplices in current filtration order. A
    filtration may be applied by sorting the simplicial complex.

    @see SimplicialComplex::sort()
  */

  iterator end()
  {
    return _simplices.template get<index_t>().end();
  }

  /** @overload end() */
  const_iterator end() const
  {
    return _simplices.template get<index_t>().end();
  }

  /**
    @param   index Simplex index
    @returns Simplex at corresponding index position. Invalid indices will not
    be caught.
  */

  const Simplex& operator[]( std::size_t index ) const
  {
    return _simplices.template get<index_t>().operator[]( index );
  }

  /**
    @param   index Simplex index
    @returns Simplex at corresponding index position
    @throws  std::runtime_error for invalid indices
  */

  const Simplex& at( std::size_t index ) const
  {
    return _simplices.template get<index_t>().at( index );
  }

  /** @returns Iterator to begin of simplices in lexicographical order. */
  lexicographical_iterator begin_lexicographical()
  {
    return _simplices.template get<lexicographical_t>().begin();
  }

  /** @overload begin_lexicographical() */
  const_lexicographical_iterator begin_lexicographical() const
  {
    return _simplices.template get<lexicographical_t>().begin();
  }

  /** @returns Iterator to end of simplices in lexicographical order. */
  lexicographical_iterator end_lexicographical()
  {
    return _simplices.template get<lexicographical_t>().end();
  }

  /** @overload end_lexicographical() */
  const_lexicographical_iterator end_lexicographical() const
  {
    return _simplices.template get<lexicographical_t>().end();
  }

  /** @returns Iterator to begin of simplices in dimensional order. */
  dimension_iterator begin_dimension()
  {
    return _simplices.template get<dimension_t>().begin();
  }

  /** @overload begin_dimension() */
  const_dimension_iterator begin_dimension() const
  {
    return _simplices.template get<dimension_t>().begin();
  }

  /** @returns Iterator to end of simplices in dimensional order. */
  dimension_iterator end_dimension()
  {
    return _simplices.template get<dimension_t>().end();
  }

  /** @overload end_dimension() */
  const_dimension_iterator end_dimension() const
  {
    return _simplices.template get<dimension_t>().end();
  }

  /**
    Given an output iterator, calculates the vertex set of the simplicial
    complex. The vertex set contains all vertices that occur in at least
    one simplex that is stored in the simplicial complex. The vertices are
    guaranteed to be reported in ascending order.

    @param result Output iterator for storing the result
  */

  template <class OutputIterator> void vertices( OutputIterator result ) const
  {
    // Using a set as an intermediate storage has the advantage that all of the
    // vertices are guaranteed to be sorted, regardless of the current order of
    // the simplicial complex.

    std::set<typename Simplex::vertex_type> vertices;

    const_dimension_iterator d_it;
    const_dimension_iterator d_it_end;

    for( std::tie( d_it, d_it_end ) = this->range( 0 );
         d_it != d_it_end;
         ++d_it )
    {
      vertices.insert( *( d_it->begin() ) );
    }

    std::copy( vertices.begin(), vertices.end(), result );
  }

  /**
    Checks whether the simplicial complex contains a given simplex. Note that
    this only checks for vertex equality. The (optional) user data is not taken
    into account. This is by design, because a simplicial complex cannot
    contain two simplices with equal vertex sets, regardless of their user data
    values.

    @param simplex Simplex whose existence is checked
    @returns true if the simplicial complex contains the simplex, else false.
  */

  bool contains( const Simplex& simplex ) const
  {
    return _simplices.template get<lexicographical_t>().find( simplex ) != _simplices.template get<lexicographical_t>().end();
  }

  /**
    Searches the simplicial complex for a given simplex and, if found, returns
    an iterator to it. If the complex does not contain the simplex, an iterator
    to the end of the simplices is returned. No (optional) user data is taken
    into account here because this makes the query more flexible: The client
    may take \i any simplex as the input parameter --- without knowing the user
    data --- but the result simplex (if valid) contains known user data that
    may only be stored within the simplicial complex.

    @param simplex Simplex to query complex for

    @returns Iterator to simplex, or an iterator to the end of the simplicial
    complex if the complex does not contain the given simplex.
  */

  const_iterator find( const Simplex& simplex ) const
  {
    auto&& it
      = _simplices.template get<lexicographical_t>().find( simplex );

    if( it != this->end_lexicographical() )
      return _simplices.template project<index_t>( it );
    else
      return this->end();
  }

  /** @overload find() */
  iterator find( const Simplex& simplex )
  {
    auto&& it
      = _simplices.template get<lexicographical_t>().find( simplex );

    if( it != this->end_lexicographical() )
      return _simplices.template project<index_t>( it );
    else
      return this->end();
  }

  /**
    Given a simplex contained by the simplicial complex, looks up its index in
    the current filtration order. If the simplex is not part of the complex, an
    exception will be thrown.

    @param simplex Simplex

    @returns Index of simplex in current filtration

    @throws std::runtime_error if the simplex is not part of the simplicial
    complex.
  */

  std::size_t index( const Simplex& simplex ) const
  {
    auto&& itSimplex = this->find( simplex );

    if( itSimplex != this->end() )
      return static_cast<std::size_t>( std::distance( this->begin(), itSimplex ) );
    else
      throw std::runtime_error( "Queried simplex does not exist" );
  }

  /** @returns Number of simplices stored in simplicial complex */
  std::size_t size() const
  {
    return _simplices.size();
  }

  /**
    @returns true if the simplicial is empty, i.e. if it does not contain any
    simplices.
  */

  bool empty() const
  {
    return _simplices.empty();
  }

  /** @returns Maximum dimension of simplices stored in simplicial complex */
  std::size_t dimension() const
  {
    if( !this->empty() )
      return _simplices.template get<dimension_t>().rbegin()->dimension();
    else
      throw std::runtime_error( "Unable to query dimensionality of empty simplicial complex" );
  }

  // Range queries -----------------------------------------------------

  /**
    Given a dimension, extracts all simplices whose dimension matches the
    user-specified one, and returns a pair of iterators for this range.

    @param dimension Dimension to extract simplices from

    @returns Pair of iterators describing the range of simplices matching the
    dimension. Note that the range is allowed to be empty.
  */

  std::pair<const_dimension_iterator, const_dimension_iterator> range( std::size_t dimension ) const
  {
    return this->range( [&] ( std::size_t d ) { return d >= dimension; },
                        [&] ( std::size_t d ) { return d <= dimension; } );
  }

  /**
    Given predicates describing the lower and upper bounds of a range, returns
    a pair of iterators for this range, assuming it is contained within the
    simplicial complex.

    @param lower Predicate describing lower bound of range
    @param upper Predicate describing upper bound of range

    @returns Pair of iterators describing the requested range.
   */

  template <typename LowerBounder, typename UpperBounder>
  std::pair<const_dimension_iterator, const_dimension_iterator> range( LowerBounder lower,
                                                                       UpperBounder upper ) const
  {
    return _simplices.template get<dimension_t>().range( lower, upper );
  }

  // Filtration modification -------------------------------------------

  /**
    Allows changing the current order of simplices, i.e. applying a certain
    simplicial filtration. This function only requires a simplex comparison
    function as its input.

    See the aleph::topology::filtrations namespace for admissible functors.

    @param comparison Simplex comparison object (or function)
  */

  template <class Comparison> void sort( Comparison&& comparison )
  {
    _simplices.sort( std::ref( comparison ) );
  }

  /** Sorts simplices according to their builtin comparison function */
  void sort()
  {
    _simplices.sort();
  }

  // -------------------------------------------------------------------

  /**
    Uses a range of vertex weights to recalculate all weights in the simplicial
    complex. Each higher-dimensional simplex is assigned the maximum of the
    weights of its lower-dimensional faces.

    @param begin Input iterator to begin of range
    @param end   Input iterator to end of range
  */

  template <class InputIterator> void recalculateWeights( InputIterator begin,
                                                          InputIterator end )
  {
    using data_type_  = typename std::iterator_traits<InputIterator>::value_type;
    using data_type   = typename Simplex::data_type;
    using vertex_type = typename Simplex::vertex_type;

    static_assert( std::is_same<data_type_, data_type>::value, "Data types must agree" );

    std::vector<data_type> weights( begin, end );

    // Reset all weights -----------------------------------------------

    for( auto itSimplex = this->begin(); itSimplex != this->end(); ++itSimplex )
    {
      Simplex newSimplex( *itSimplex );
      newSimplex.setData( std::numeric_limits<data_type>::max() );

      this->replace( itSimplex, newSimplex );
    }

    // Assign 0-dimensional weights ------------------------------------

    for( auto itSimplex = this->begin_dimension();
         itSimplex != this->end_dimension();
         ++itSimplex )
    {
      if( itSimplex->dimension() == 0 )
      {
        vertex_type v = *( itSimplex->begin() );

        Simplex newSimplex = ( *itSimplex );
        newSimplex.setData( weights.at(v) );

        this->replace( _simplices.template project<index_t>( itSimplex ),
                       newSimplex );
      }
      else
        break;
    }

    this->recalculateWeights();
  }

  // -------------------------------------------------------------------

  /**
    Recalculates simplex weights by assigning each simplex the maximum
    or minimum weight of its faces. Note that 0-dimensional simplices,
    i.e. vertices, are _always_ skipped by this function.

    @param useMaximum If set, uses the maximum data assigned to a face of a
    simplex in order to assign its final weight.

    @param skipOneDimensionalSimplices If set, skips both 0-dimensional and
    1-dimensional simplices and accepts their weights as the given truth.
  */

  void recalculateWeights( bool useMaximum = true, bool skipOneDimensionalSimplices = false )
  {
    // Assign weights based on lower-dimensional simplices -------------

    for( auto itSimplex = this->begin_dimension();
         itSimplex != this->end_dimension();
         ++itSimplex )
    {
      if(    ( itSimplex->dimension() == 0 )
          || ( skipOneDimensionalSimplices && itSimplex->dimension() == 1 ) )
      {
        continue;
      }

      using DataType = typename Simplex::DataType;

      DataType weight
        = useMaximum ? std::numeric_limits<DataType>::lowest()
                     : std::numeric_limits<DataType>::max();

      for( auto itBoundary = itSimplex->begin_boundary();
           itBoundary != itSimplex->end_boundary();
           ++itBoundary )
      {
        const_iterator itPos = this->find( *itBoundary );
        if( itPos != this->end() )
        {
          weight = useMaximum ? std::max( weight, itPos->data() )
                              : std::min( weight, itPos->data() );
        }

        // The if-branch above ignores missing boundaries. This is useful when a
        // filtration is only partially defined, e.g. only up to the 2-simplices
        // and higher-dimensional simplices.
      }

      Simplex newSimplex = ( *itSimplex );
      newSimplex.setData( weight );

      this->replace( _simplices.template project<index_t>( itSimplex ),
                     newSimplex );
    }
  }

  // Container modification --------------------------------------------

  /**
    Allows simplex removal by value. The simplicial complex will check whether
    the given simplex exists. If so, it will be erased. Note that erasing a
    simplex may trigger the removal of co-faces of the simplex in order to
    remain validity.

    @param simplex Simplex to remove
  */

  void remove( const Simplex& simplex )
  {
    _simplices.template get<index_t>().remove( simplex );

    bool foundInvalidSimplex = false;

    do
    {
      foundInvalidSimplex = false;

      for( auto itSimplex = this->begin();
           itSimplex != this->end(); )
      {
        if( this->checkValidity( *itSimplex ) == false )
        {
          foundInvalidSimplex = true;
          itSimplex           = _simplices.template get<index_t>().erase( itSimplex );
        }
        else
          ++itSimplex;
      }
    }
    while( foundInvalidSimplex );
  }

  /**
    Permits simplex removal by value. The function does *not* check
    whether additional co-faces will have to be removed. Hence, the
    complex may be invalid afterwards.
  */

  void remove_without_validation( const Simplex& simplex )
  {
    _simplices.template get<index_t>().remove( simplex );
  }

  /**
    Creates all missing faces of the current simplicial complex. When
    this function is finished, all of the faces for all simplices are
    part of the simplicial complex.

    This function is useful if constructing a simplicial complex from
    a partial range of simplices. The typical application scenario is
    the processing of triangulations of manifolds, which are commonly
    only specified by their top-level simplices. There is no need for
    clients to calculate these faces manually.
  */

  void createMissingFaces()
  {
    for( auto itSimplex = this->begin(); itSimplex != this->end(); itSimplex++ )
      this->checkAndRestoreValidity( *itSimplex );
  }

  // Comparison --------------------------------------------------------

  /**
    Checks two simplicial complexes for equality. This operator will
    check simplicial complexes for equality but only with respect to
    their current filtration order.
  */

  bool operator==( const SimplicialComplex<Simplex>& other ) const
  {
    if( this->size() != other.size() )
      return false;

    for( auto it1 = this->begin(), it2 = other.begin();
         it1 != this->end() && it2 != other.end();
         ++it1, ++it2 )
    {
      if( *it1 != *it2 )
        return false;
    }

    return true;
  }

  /**
    Checks whether two simplicial complexes differ by at least one
    simplex with respect to their current filtration order.
  */

  bool operator!=( const SimplicialComplex<Simplex>& other ) const
  {
    return !this->operator==( other );
  }

private:

  /**
    Checks and restores validity of the simplicial complex after adding a
    single simplex. This means that the simplicial complex will check whether
    it contains \i all faces of the given simplex, as well as the simplex
    itself.

    If a face is found to be missing, it will be created by the function.

    @param simplex Simplex whose addition to the simplicial complex caused the
    validity check to be triggered.
  */

  void checkAndRestoreValidity( const Simplex& simplex )
  {
    for( auto itFace = simplex.begin_boundary();
         itFace != simplex.end_boundary();
         ++itFace )
    {
      const_lexicographical_iterator itStoredFace
          = _simplices.template get<lexicographical_t>().find( *itFace );

      if( itStoredFace == _simplices.template get<lexicographical_t>().end() )
      {
        // The new simplex shall contain the same vertices as the "face
        // simplex", but the data from its parent simplex. This ensures that
        // the data of a coface is always greater than or equal to the data of
        // its faces (assuming that the data type is comparable).
        _simplices.push_back( Simplex( *itFace,
                                       simplex.data() ) );
      }
    }
  }

  /**
    Checks validity of a single simplex. A simplex in the simplicial complex is
    deemed valid if all of its faces can be found in the complex.

    @param simplex Simplex whose validity is to be checked
    @returns true if the simplex is valid, else false
  */

  bool checkValidity( const Simplex& simplex )
  {
    for( auto itFace = simplex.begin_boundary();
         itFace != simplex.end_boundary();
         ++itFace )
    {
      // Check whether an unknown face has been found. If so, the simplex is
      // invalid.
      const_lexicographical_iterator itStoredFace
          = _simplices.template get<lexicographical_t>().find( *itFace );

      if( itStoredFace == _simplices.template get<lexicographical_t>().end() )
        return false;
    }

    return true;
  }

  /**
    Simplex container. boost::multi_index is used to provide different "views"
    to the simplicial data set. The first view uses the current sorting order
    of simplices. It corresponds to a given "filtration" order, in which faces
    precede cofaces. The second view allows a lexicographical traversal of the
    current simplicial complex. This view is used to provide fast index
    queries.
  */

  simplex_container_t _simplices;
};

// ---------------------------------------------------------------------

/**
  Adds information about a simplicial complex to an output stream. This
  is useful for debugging purposes or intensive logging.

  @param o ostream to add simplicial complex to
  @param S Simplicial complex to stream to ostream

  @returns ostream with information about simplicial complex
*/

template <class Simplex> std::ostream& operator<<( std::ostream& o,
                                                   const topology::SimplicialComplex<Simplex>& S )
{
  if( S.empty() )
    return o;

  o << std::string( 80, '-' ) << "\n";

  for( auto it = S.begin(); it != S.end(); ++it )
    o << *it << "\n";

  o << std::string( 80, '-' ) << "\n";

  return o;
}

// ---------------------------------------------------------------------

} // namespace topology

} // namespace aleph

#endif


#include <fstream>
#include <csignal>
#include <cstdio>
#include <sstream>
#include <cstdlib>
#include <iostream>
#include <queue>
#include <limits>
#include <mutex>
// ROS
#include "ros/ros.h"
#include <nav_msgs/Odometry.h>
#include "sensor_msgs/Image.h"
#include <message_filters/time_synchronizer.h>
#include <message_filters/subscriber.h>
#include <message_filters/synchronizer.h>
#include <message_filters/sync_policies/approximate_time.h>
#include <sensor_msgs/point_cloud_conversion.h>
// Opencv
#include <cv_bridge/cv_bridge.h>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/core.hpp>
//#include "opencv2/core/eigen.hpp"
// PCL
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/point_cloud.h>
#include <pcl/io/ply_io.h>
#include <pcl/point_types.h>
#include <pcl/octree/octree.h>
#include <pcl/filters/passthrough.h>
#include <pcl/filters/project_inliers.h>
#include <pcl/ModelCoefficients.h>
#include <pcl/common/transforms.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/filters/approximate_voxel_grid.h>
// Eigen
#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/Geometry>
using namespace cv;
using namespace Eigen;
using namespace std;

typedef message_filters::sync_policies::ApproximateTime<nav_msgs::Odometry,sensor_msgs::Image> slamSyncPolicy;
struct color_data
{
    /* data */
    sensor_msgs::Image::ConstPtr image;
    nav_msgs::Odometry::ConstPtr pose;
};

class ir_lidar_color_pointcloud
{
private:
    /* pcl data */
    pcl::PointCloud<pcl::PointXYZ>::Ptr map_cloud_,show_cloud_;
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr color_cloud;
    std::queue<color_data> msgs;
    std::mutex que_mtx;
    /* data */
    ros::NodeHandle nh_;
    ros::Publisher globalmap_pub,sim_pub_;
    ros::Subscriber gray_sub,motor_sub;
    message_filters::Subscriber<nav_msgs::Odometry>* odom_sub_ ;
    message_filters::Subscriber<sensor_msgs::Image>* img_sub_;
    message_filters::Synchronizer<slamSyncPolicy>* sync_;
    void imageCallback(const sensor_msgs::Image::ConstPtr& data);
    void poseCallBack(const nav_msgs::Odometry::ConstPtr& data);
    void syncCallback(const nav_msgs::Odometry::ConstPtr& data, const sensor_msgs::Image::ConstPtr& msg);
    void loadMap(std::string &fileName);
public:
    bool is_save = true;
    ir_lidar_color_pointcloud(const ros::NodeHandle &nh = ros::NodeHandle("~"));
    ~ir_lidar_color_pointcloud();
    Eigen::Affine3f world2laser(nav_msgs::Odometry::Ptr odom);
    Eigen::Affine3f laser2cam();
    Eigen::Vector2f cam2pixel(Eigen::Vector3f pc);
    Eigen::Vector3f get_color(const cv::Mat image, const int row,const int col);
    void rawToJet(const cv::Mat &raw,cv::Mat *color);
    //void read_calibrationfile(std::string _choose_file,cv::Mat &out_RT);
    //void read_camerafile(std::string& file);
    void update();
    static void* poll(void* ptr);
};


// Copyright (c) 2021 Graphcore Ltd. All rights reserved.
#define BOOST_TEST_MODULE Test_Ir_executeOpNTimesEveryMTimesTests

#include <boost/algorithm/string/predicate.hpp>
#include <boost/test/unit_test.hpp>

#include <popart/dataflow.hpp>
#include <popart/graph.hpp>
#include <popart/ir.hpp>
#include <popart/ndarraywrapper.hpp>
#include <popart/op/add.hpp>
#include <popart/op/atan.hpp>
#include <popart/op/histogram.hpp>
#include <popart/op/identity.hpp>
#include <popart/op/sin.hpp>
#include <popart/op/sinh.hpp>
#include <popart/session.hpp>
#include <popart/tensorinfo.hpp>
#include <popart/testdevice.hpp>
#include <popart/transforms/automaticlossscaling.hpp>

using namespace popart;

namespace {

template <typename Ex>
std::function<bool(const Ex &)>
checkErrorMsgHasPrefixFn(const std::string &prefix) {
  return [=](const Ex &ex) -> bool {
    return boost::algorithm::starts_with(ex.what(), prefix);
  };
}

} // namespace

BOOST_AUTO_TEST_CASE(TestFrequency1input1output2Ops) {
  // in -> sinhOp -> sinOp -> out
  auto ir  = std::make_unique<Ir>();
  Graph &g = ir->getMainGraph();

  const TensorId t0   = "t0";
  const TensorId tout = "tout";
  std::vector<float> t0Data(1, 1);
  const TensorInfo t0_ti{DataType::FLOAT, Shape{}};
  const TensorInfo tout_ti{DataType::FLOAT, Shape{}};

  g.getTensors().addVarInit(t0, t0_ti, t0Data.data());

  Op::Settings gSettings(g, "op", {});
  Op *sinhOp = g.createConnectedOp<SinhInplaceOp>(
      {{SinhInplaceOp::getInIndex(), t0}},
      {{SinhInplaceOp::getOutIndex(), ir->createIntermediateTensorId(t0)}},
      gSettings.copy("SinhInplaceOp"));

  Op::Settings gSettingsSin(g, "sin", {});
  Op *sinOp = g.createConnectedOp<SinOp>(
      {{SinOp::getInIndex(), sinhOp->outTensor(0)->id}},
      {{SinOp::getOutIndex(), tout}},
      Onnx::Operators::Sin_7,
      gSettingsSin.copy("SinOp"));
  sinOp->pruneable = false;

  std::map<InIndex, OutIndex> identityInputToOutputIndiciesMapping{{0, 0}};
  std::map<OutIndex, float> outputIndiciesAndValues;
  constexpr int bps = 18;
  ir->setDataFlow(DataFlow{bps, {{tout, AnchorReturnType("All")}}});
  sinhOp = AutomaticLossScale::executeOpNTimesEveryMTimes(
      sinhOp,
      2,
      6,
      identityInputToOutputIndiciesMapping,
      outputIndiciesAndValues);

  ir->updateVertices();
  ir->setIsPrepared();

  const auto session = InferenceSession::createFromIr(
      std::move(ir), createTestDevice(TEST_TARGET));
  session->prepareDevice();

  const TensorInfo toutND_ti{DataType::FLOAT, Shape{bps * tout_ti.nelms()}};
  std::vector<float> outHost(bps * tout_ti.nelms());
  NDArrayWrapper<float> outWrapper(outHost.data(), toutND_ti);
  std::map<TensorId, IArray &> anchors = {{tout, outWrapper}};

  StepIO stepio({}, anchors);
  session->weightsFromHost();
  session->run(stepio);

  // n=2, m=6
  float v1 = 0.922767401;
  float v2 = 0.841470957;
  std::vector<float> expected{
      v1, v1, v2, v2, v2, v2, v1, v1, v2, v2, v2, v2, v1, v1, v2, v2, v2, v2};

  BOOST_TEST(expected.size() == outHost.size());
  BOOST_TEST(expected == outHost);
}

BOOST_AUTO_TEST_CASE(TestExecuteOpNTimesEveryMTimesShapes) {

  auto test = [](const std::map<InIndex, OutIndex>
                     &identityInputToOutputIndiciesMapping,
                 const std::map<OutIndex, float> &outputIndiciesAndValues,
                 float t0Value,
                 float tcValue,
                 Shape testShape,
                 const std::vector<float> &expected) {
    // t0    tc
    //  \   /
    //   Add
    //    |
    // IdentityOp
    //    |

    auto ir  = std::make_unique<Ir>();
    Graph &g = ir->getMainGraph();

    const TensorId t0   = "t0";
    const TensorId tc   = "tc";
    const TensorId tout = "tout";
    const TensorInfo t0_ti{DataType::FLOAT, testShape};
    const TensorInfo tc_ti{DataType::FLOAT, testShape};
    const TensorInfo tout_ti{DataType::FLOAT, testShape};
    std::vector<float> t0Data(t0_ti.nelms(), t0Value);

    g.getTensors().addVarInit(t0, t0_ti, t0Data.data());
    std::vector<float> constData(t0_ti.nelms(), tcValue);
    g.getTensors().addConstInit(tc, tc_ti, constData.data());

    Op::Settings gSettings(g, "op", {});
    Op *addOp = g.createConnectedOp<AddOp>(
        {{AddOp::getArg0InIndex(), t0}, {AddOp::getArg1InIndex(), tc}},
        {{AddOp::getOutIndex(), ir->createIntermediateTensorId(t0)}},
        Onnx::Operators::Add_7,
        gSettings.copy("Add"));

    Op::Settings gSettingsIdentityOp(g, "IdentityOp", {});
    Op *identityOp = g.createConnectedOp<IdentityOp>(
        {{IdentityOp::getInIndex(), addOp->outTensor(0)->id}},
        {{IdentityOp::getOutIndex(), tout}},
        Onnx::Operators::Identity_1,
        gSettingsIdentityOp.copy("IdentityOp"));

    identityOp->pruneable = false;

    constexpr int bps = 18;
    ir->setDataFlow(DataFlow{bps, {{tout, AnchorReturnType("All")}}});
    addOp = AutomaticLossScale::executeOpNTimesEveryMTimes(
        addOp,
        2,
        6,
        identityInputToOutputIndiciesMapping,
        outputIndiciesAndValues);

    ir->updateVertices();
    ir->setIsPrepared();

    const auto session = InferenceSession::createFromIr(
        std::move(ir), createTestDevice(TEST_TARGET));
    session->prepareDevice();

    const TensorInfo toutND_ti{DataType::FLOAT, Shape{bps * t0_ti.nelms()}};
    std::vector<float> outHost(bps * t0_ti.nelms());
    NDArrayWrapper<float> outWrapper(outHost.data(), toutND_ti);
    std::map<TensorId, IArray &> anchors = {{tout, outWrapper}};

    StepIO stepio({}, anchors);
    session->weightsFromHost();
    session->run(stepio);

    // n=2, m=6
    BOOST_TEST(expected.size() == outHost.size());
    BOOST_TEST(expected == outHost);
  };

  // Check t0 is connected by nop to output of Add op in no compute subgraph.
  std::map<InIndex, OutIndex> identityInToOutIndicies1{{0, 0}};
  std::map<OutIndex, float> outputIndiciesAndValues1;
  float t0Value                = 1;
  float tcValue                = 5;
  std::vector<float> expected1 = {
      6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1};
  test(identityInToOutIndicies1,
       outputIndiciesAndValues1,
       t0Value,
       tcValue,
       Shape{},
       expected1);

  // Check tc is connected by nop to output of Add op in no compute subgraph.
  std::map<InIndex, OutIndex> identityInToOutIndicies2{{1, 0}};
  std::map<OutIndex, float> outputIndiciesAndValues2;
  std::vector<float> expected2 = {
      6, 6, 5, 5, 5, 5, 6, 6, 5, 5, 5, 5, 6, 6, 5, 5, 5, 5};
  test(identityInToOutIndicies2,
       outputIndiciesAndValues2,
       t0Value,
       tcValue,
       Shape{},
       expected2);

  // Check outputIndiciesAndValues for empty subgraph.
  std::map<InIndex, OutIndex> identityInToOutIndiciesTest3;
  float outputValues = 11;
  std::map<OutIndex, float> outputIndiciesAndValues3{{0, outputValues}};
  std::vector<float> expected3 = {
      6, 6, 11, 11, 11, 11, 6, 6, 11, 11, 11, 11, 6, 6, 11, 11, 11, 11};
  test(identityInToOutIndiciesTest3,
       outputIndiciesAndValues3,
       t0Value,
       tcValue,
       Shape{},
       expected3);

  // Check outputIndiciesAndValues for empty subgraph and non scalar Shape.
  std::map<InIndex, OutIndex> identityInToOutIndicies4;
  float outputValues4 = 11;
  std::map<OutIndex, float> outputIndiciesAndValues4{{0, outputValues4}};
  std::vector<float> expected4 = {
      6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  11, 11, 11, 11, 11, 11,
      11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11,
      6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  11, 11, 11, 11, 11, 11,
      11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11,
      6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  11, 11, 11, 11, 11, 11,
      11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11};
  test(identityInToOutIndicies4,
       outputIndiciesAndValues4,
       t0Value,
       tcValue,
       Shape{2, 3},
       expected4);

  // Check t0 is connected by nop to output of Add op in no compute subgraph.
  std::map<InIndex, OutIndex> identityInToOutIndicies5{{0, 0}};
  std::map<OutIndex, float> outputIndiciesAndValues5;
  std::vector<float> expected5 = {
      6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
      1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 6, 6, 6, 6, 6, 6, 6, 6,
      6, 6, 6, 6, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
      1, 1, 1, 1, 1, 1, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 1, 1, 1, 1,
      1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1};
  test(identityInToOutIndicies5,
       outputIndiciesAndValues5,
       t0Value,
       tcValue,
       Shape{2, 3},
       expected5);
}

BOOST_AUTO_TEST_CASE(
    TestExecuteOpNTimesEveryMTimes1input1outputDifferentTypes) {
  //    ts
  //    |
  // Histogram
  //    |
  // IdentityOp
  //    |

  auto ir  = std::make_unique<Ir>();
  Graph &g = ir->getMainGraph();

  const TensorId ts   = "ts";
  const TensorId tout = "tout";
  const TensorInfo ts_ti{DataType::FLOAT, Shape{2, 2}};
  const TensorInfo tout_ti{DataType::UINT32, Shape{2}};

  g.getTensors().addStream(ts, ts_ti);

  Op::Settings gSettings(g, "op", {});
  std::vector<float> levels = {0.5};
  bool absoluteOfInput      = true;

  Op *histogramOp = g.createConnectedOp<HistogramOp>(
      {{HistogramOp::getInIndex(), ts}},
      {{HistogramOp::getOutIndex(), ir->createIntermediateTensorId(ts)}},
      Onnx::CustomOperators::Histogram,
      levels,
      absoluteOfInput,
      gSettings.copy("HistogramOp"));

  Op::Settings gSettingsIdentityOp(g, "IdentityOp", {});
  Op *identityOp = g.createConnectedOp<IdentityOp>(
      {{IdentityOp::getInIndex(), histogramOp->outTensor(0)->id}},
      {{IdentityOp::getOutIndex(), tout}},
      Onnx::Operators::Identity_1,
      gSettingsIdentityOp.copy("IdentityOp"));

  identityOp->pruneable = false;

  std::map<InIndex, OutIndex> identityInputToOutputIndiciesMapping;
  std::map<OutIndex, float> outputIndiciesAndValues{{0, 11}};

  histogramOp = AutomaticLossScale::executeOpNTimesEveryMTimes(
      histogramOp,
      2,
      6,
      identityInputToOutputIndiciesMapping,
      outputIndiciesAndValues);

  ir->updateVertices();
  ir->setIsPrepared();
  constexpr int bps = 20;
  ir->setDataFlow(DataFlow{bps, {{tout, AnchorReturnType("All")}}});

  const auto session = InferenceSession::createFromIr(
      std::move(ir), createTestDevice(TEST_TARGET));
  session->prepareDevice();

  const TensorInfo toutND_ti{DataType::UINT32, Shape{bps * tout_ti.nelms()}};
  std::vector<uint32_t> outHost(bps * tout_ti.nelms());
  NDArrayWrapper<uint32_t> outWrapper(outHost.data(), toutND_ti);
  std::map<TensorId, IArray &> anchors = {{tout, outWrapper}};

  std::vector<float> ts_init(bps * ts_ti.nelms());
  for (int i = 0; i < bps * ts_ti.nelms(); i++) {
    ts_init[i] = 2.0 * i;
  }

  const TensorInfo tsWraper_ti{DataType::FLOAT, Shape{bps, 2, 2}};
  popart::NDArrayWrapper<float> ts_wrapper(ts_init.data(), tsWraper_ti);
  std::map<popart::TensorId, popart::IArray &> inputs = {{ts, ts_wrapper}};

  StepIO stepio(inputs, anchors);
  session->weightsFromHost();
  session->run(stepio);

  // n=2, m=6
  std::vector<uint32_t> expected{1,  3,  0,  4,  11, 11, 11, 11, 11, 11,
                                 11, 11, 0,  4,  0,  4,  11, 11, 11, 11,
                                 11, 11, 11, 11, 0,  4,  0,  4,  11, 11,
                                 11, 11, 11, 11, 11, 11, 0,  4,  0,  4};

  BOOST_TEST(expected.size() == outHost.size());
  BOOST_TEST(expected == outHost);
}

template <typename T1>
void testType(
    int n,
    int m,
    const std::map<InIndex, OutIndex> &identityInputToOutputIndiciesMapping,
    const std::map<OutIndex, float> &outputIndiciesAndValues,
    T1 t0Value,
    T1 tcValue,
    const Shape &testShape,
    const std::vector<T1> &expected,
    DataType T2,
    bool enableGradientAccumulation,
    int bps,
    int accumulationFactor = 1) {

  // t0    tc
  //  \   /
  //   Add
  //    |
  // IdentityOp
  //    |

  auto ir  = std::make_unique<Ir>();
  Graph &g = ir->getMainGraph();

  SessionOptions &sessionOptions = ir->getSessionOptions();
  if (enableGradientAccumulation) {
    sessionOptions.enableGradientAccumulation = true;
    sessionOptions.accumulationFactor         = accumulationFactor;
  }

  const TensorId t0   = "t0";
  const TensorId tc   = "tc";
  const TensorId tout = "tout";
  TensorInfo t0_ti;
  t0_ti = TensorInfo(T2, testShape);
  const TensorInfo tc_ti{T2, testShape};
  const TensorInfo tout_ti{T2, testShape};
  std::vector<T1> t0Data(t0_ti.nelms(), t0Value);

  g.getTensors().addVarInit(t0, t0_ti, t0Data.data());
  std::vector<T1> constData(t0_ti.nelms(), tcValue);
  g.getTensors().addConstInit(tc, tc_ti, constData.data());

  Op::Settings gSettings(g, "op", {});
  Op *addOp = g.createConnectedOp<AddOp>(
      {{AddOp::getArg0InIndex(), t0}, {AddOp::getArg1InIndex(), tc}},
      {{AddOp::getOutIndex(), ir->createIntermediateTensorId(t0)}},
      Onnx::Operators::Add_7,
      gSettings.copy("Add"));

  Op::Settings gSettingsIdentityOp(g, "IdentityOp", {});
  Op *identityOp = g.createConnectedOp<IdentityOp>(
      {{IdentityOp::getInIndex(), addOp->outTensor(0)->id}},
      {{IdentityOp::getOutIndex(), tout}},
      Onnx::Operators::Identity_1,
      gSettingsIdentityOp.copy("IdentityOp"));

  identityOp->pruneable = false;

  ir->setDataFlow(DataFlow{bps, {{tout, AnchorReturnType("All")}}});

  addOp = AutomaticLossScale::executeOpNTimesEveryMTimes(
      addOp,
      n,
      m,
      identityInputToOutputIndiciesMapping,
      outputIndiciesAndValues);

  ir->updateVertices();
  ir->setIsPrepared();

  const auto session = InferenceSession::createFromIr(
      std::move(ir), createTestDevice(TEST_TARGET));

  session->prepareDevice();

  const TensorInfo toutND_ti{T2, Shape{bps * t0_ti.nelms()}};
  std::vector<T1> outHost(bps * t0_ti.nelms());
  NDArrayWrapper<T1> outWrapper(outHost.data(), toutND_ti);
  std::map<TensorId, IArray &> anchors = {{tout, outWrapper}};

  StepIO stepio({}, anchors);
  session->weightsFromHost();
  session->run(stepio);

  BOOST_TEST(expected.size() == outHost.size());
  BOOST_TEST(expected == outHost);
}

BOOST_AUTO_TEST_CASE(TestExecuteOpNTimesEveryMTimesTypes) {

  std::map<InIndex, OutIndex> identityInToOutIndicies1{{0, 0}};
  std::map<OutIndex, float> outputIndiciesAndValues1;
  float t0Value1               = 1;
  float tcValue1               = 5;
  std::vector<float> expected1 = {
      6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1};
  testType<float>(2,
                  6,
                  identityInToOutIndicies1,
                  outputIndiciesAndValues1,
                  t0Value1,
                  tcValue1,
                  Shape{},
                  expected1,
                  DataType::FLOAT,
                  false,
                  18);

  std::map<InIndex, OutIndex> identityInToOutIndicies2{{0, 0}};
  std::map<OutIndex, float> outputIndiciesAndValues2;
  int32_t t0Value2               = 1;
  int32_t tcValue2               = 5;
  std::vector<int32_t> expected2 = {
      6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1};
  testType<int32_t>(2,
                    6,
                    identityInToOutIndicies2,
                    outputIndiciesAndValues2,
                    t0Value2,
                    tcValue2,
                    Shape{},
                    expected2,
                    DataType::INT32,
                    false,
                    18);

  std::map<InIndex, OutIndex> identityInToOutIndicies3;
  int32_t t0Value3    = 1;
  int32_t tcValue3    = 5;
  float outputValues3 = 11;
  std::map<OutIndex, float> outputIndiciesAndValues3{{0, outputValues3}};
  std::vector<int32_t> expected3 = {
      6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  11, 11, 11, 11, 11, 11,
      11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11,
      6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  11, 11, 11, 11, 11, 11,
      11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11,
      6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  6,  11, 11, 11, 11, 11, 11,
      11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11};
  testType<int32_t>(2,
                    6,
                    identityInToOutIndicies3,
                    outputIndiciesAndValues3,
                    t0Value3,
                    tcValue3,
                    Shape{2, 3},
                    expected3,
                    DataType::INT32,
                    false,
                    18);

  // Check when identityInputToOutputIndiciesMapping and outputIndiciesAndValues
  // contain the same index.
  std::map<InIndex, OutIndex> identityInToOutIndicies4{{0, 0}};
  std::map<OutIndex, float> outputIndiciesAndValues4{{0, 11}};
  float t0Value4               = 1;
  float tcValue4               = 5;
  std::vector<float> expected4 = {
      6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1};

  const auto checkErrorFn4 = checkErrorMsgHasPrefixFn<error>(
      "[AutomaticLossScale transform][[executeOpNTimesEveryMTimes] "
      "Incorrect frequency modifier settings for 100 (ai.onnx.Add:7) operator "
      "Identity output and default output can not have the same index 0.");

  BOOST_REQUIRE_EXCEPTION(testType<float>(2,
                                          6,
                                          identityInToOutIndicies4,
                                          outputIndiciesAndValues4,
                                          t0Value4,
                                          tcValue4,
                                          Shape{},
                                          expected4,
                                          DataType::FLOAT,
                                          false,
                                          18),
                          error,
                          checkErrorFn4);

  // Check when identityInputToOutputIndiciesMapping contain an invalid index.
  std::map<InIndex, OutIndex> identityInToOutIndicies5{{3, 0}};
  std::map<OutIndex, float> outputIndiciesAndValues5;
  float t0Value5               = 1;
  float tcValue5               = 5;
  std::vector<float> expected5 = {
      6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1};

  const auto checkErrorFn5 = checkErrorMsgHasPrefixFn<error>(
      "[AutomaticLossScale transform][[executeOpNTimesEveryMTimes] "
      "identityInputToOutputIndiciesMapping has invalid input index 3 "
      "for operator 100 (ai.onnx.Add:7).");

  BOOST_REQUIRE_EXCEPTION(testType<float>(2,
                                          6,
                                          identityInToOutIndicies5,
                                          outputIndiciesAndValues5,
                                          t0Value5,
                                          tcValue5,
                                          Shape{},
                                          expected5,
                                          DataType::FLOAT,
                                          false,
                                          18),
                          error,
                          checkErrorFn5);

  // Check when identityInputToOutputIndiciesMapping contain an invalid index.
  std::map<InIndex, OutIndex> identityInToOutIndicies6{{1, 3}};
  std::map<OutIndex, float> outputIndiciesAndValues6;
  float t0Value6               = 1;
  float tcValue6               = 5;
  std::vector<float> expected6 = {
      6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1};

  const auto checkErrorFn6 = checkErrorMsgHasPrefixFn<error>(
      "[AutomaticLossScale transform][[executeOpNTimesEveryMTimes] "
      "identityInputToOutputIndiciesMapping has invalid output index 3 "
      "for operator 100 (ai.onnx.Add:7).");

  BOOST_REQUIRE_EXCEPTION(testType<float>(2,
                                          6,
                                          identityInToOutIndicies6,
                                          outputIndiciesAndValues6,
                                          t0Value6,
                                          tcValue6,
                                          Shape{},
                                          expected6,
                                          DataType::FLOAT,
                                          false,
                                          18),
                          error,
                          checkErrorFn6);

  // Check when outputIndiciesAndValues contain an invalid index.
  std::map<InIndex, OutIndex> identityInToOutIndicies7;
  std::map<OutIndex, float> outputIndiciesAndValues7{{3, 11}};
  float t0Value7               = 1;
  float tcValue7               = 5;
  std::vector<float> expected7 = {
      6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1};

  const auto checkErrorFn7 = checkErrorMsgHasPrefixFn<error>(
      "[AutomaticLossScale transform][[executeOpNTimesEveryMTimes] "
      "outputIndiciesAndValues has invalid output index 3 "
      "for operator 100 (ai.onnx.Add:7).");

  BOOST_REQUIRE_EXCEPTION(testType<float>(2,
                                          6,
                                          identityInToOutIndicies7,
                                          outputIndiciesAndValues7,
                                          t0Value7,
                                          tcValue7,
                                          Shape{},
                                          expected7,
                                          DataType::FLOAT,
                                          false,
                                          18),
                          error,
                          checkErrorFn7);

  // Check M.
  std::map<InIndex, OutIndex> identityInToOutIndicies8{{0, 0}};
  std::map<OutIndex, float> outputIndiciesAndValues8;
  float t0Value8               = 1;
  float tcValue8               = 5;
  std::vector<float> expected8 = {
      6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1};

  const auto checkErrorFn8 = checkErrorMsgHasPrefixFn<error>(
      "[AutomaticLossScale transform][[executeOpNTimesEveryMTimes]."
      "Argument M of executeOpNTimesEveryMTimes has inconsistent value 6. "
      "Operation 100 (ai.onnx.Add:7) is in the Normal execution context and "
      "gradient accumulation is enabled hence M should be a factor of "
      "gradient accumulation factor 10.");

  BOOST_REQUIRE_EXCEPTION(testType<float>(4,
                                          6,
                                          identityInToOutIndicies8,
                                          outputIndiciesAndValues8,
                                          t0Value8,
                                          tcValue8,
                                          Shape{},
                                          expected8,
                                          DataType::FLOAT,
                                          true,
                                          18,
                                          10),
                          error,
                          checkErrorFn8);

  // Check M and batchesPerStep.
  std::map<InIndex, OutIndex> identityInToOutIndicies9{{0, 0}};
  std::map<OutIndex, float> outputIndiciesAndValues9;
  float t0Value9               = 1;
  float tcValue9               = 5;
  std::vector<float> expected9 = {
      6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1, 6, 6, 1, 1, 1, 1};

  const auto checkErrorFn9 = checkErrorMsgHasPrefixFn<error>(
      "[AutomaticLossScaletransform][[executeOpNTimesEveryMTimes]. "
      "Argument M of executeOpNTimesEveryMTimes has inconsistent value 6. "
      "M should be a factor of batches per step 10.");

  BOOST_REQUIRE_EXCEPTION(testType<float>(4,
                                          6,
                                          identityInToOutIndicies9,
                                          outputIndiciesAndValues9,
                                          t0Value9,
                                          tcValue9,
                                          Shape{},
                                          expected9,
                                          DataType::FLOAT,
                                          false,
                                          10),
                          error,
                          checkErrorFn9);
}