/*
 * Copyright(C) 1999-2021, 2023 National Technology & Engineering Solutions
 * of Sandia, LLC (NTESS).  Under the terms of Contract DE-NA0003525 with
 * NTESS, the U.S. Government retains certain rights in this software.
 *
 * See packages/seacas/LICENSE for details
 */
/**********************************************************************/
/* matlab mat file to exodus II.  This takes .mat files exactly as
   generated by the tool exo2mat written by mrtabba and converts them
   back to exodus II format

   rmnaeth. August 8, 2003

   modified by D. Todd Griffith on 12/09/2005
   * modifications include:
   1) writes global, nodal and element variable names
   2) writes global, nodal and elemnent variable results
   3) writes complete set of time steps (previous version
          skipped first step)
   4) writes complete node set information (node set numbers,
          dist. factors, etc)
   5) writes complete side set information (side set numbers,
          dist. factors, etc)

   modified by D. Todd Griffith on 12/16/2005
   * side set distribution factors now written as double (not int)

   modified by Greg Sjaardema, 07/05/2012 to use matio instead of matlab libraries.
*/

#include "add_to_log.h" // for add_to_log
#include "fmt/printf.h"
#include "matio.h"         // for matvar_t, Mat_VarFree, etc
#include "matio_pubconf.h" // for MATIO_VERSION
#include <SL_tokenize.h>   // for tokenize

#include <array>
#include <cstddef>    // for size_t
#include <cstring>    // for strtok, memcpy, strlen, etc
#include <exodusII.h> // for ex_put_variable_param, etc
#include <numeric>    // for accumulate
#include <smart_assert.h>
#include <string> // for char_traits, string
#include <vector> // for vector

#if MATIO_VERSION < 151
#error "MatIO Version 1.5.1 or greater is required"
#endif

/**********************************************************************/
mat_t *mat_file = nullptr; /* file for binary .mat input */

/**********************************************************************/
static std::array<std::string, 3> qainfo{"mat2exo", "2021/09/27", "4.06"};

/**********************************************************************/
void get_put_names(int exo_file, ex_entity_type entity, int num_vars, const std::string &name);
void get_put_user_names(int exo_file, ex_entity_type entity, int num_entity, const char *mname);
void get_put_attr_names(int exo_file, int seq, int id, int num_attr);
void get_put_vars(int exo_file, ex_entity_type type, const std::vector<int> &ids, int num_blocks,
                  int num_vars, int num_time_steps, const std::vector<int> &num_per_block,
                  const char *mname);

std::vector<std::string> matGetStr(const std::string &name);
int matGetDbl(const std::string &name, size_t n1, size_t n2, std::vector<double> &data);
int matGetInt(const std::string &name, size_t n1, size_t n2, std::vector<int> &data);
int matGetInt(const std::string &name);
int matArrNRow(const std::string &name);
int matArrNCol(const std::string &name);

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

  /* QA Info */
  fmt::printf("%s: %s, %s\n", qainfo[0], qainfo[2], qainfo[1]);

  /* usage message*/
  if (argc != 2) {
    fmt::printf("%s matlab_file_name.\n", argv[0]);
    fmt::printf("   the matlab_file_name is required\n");
    fmt::printf("%d", argc);
    exit(1);
  }

  /*open input file*/
  mat_file = Mat_Open(argv[1], MAT_ACC_RDONLY);
  if (mat_file == nullptr) {
    fmt::printf("Error opening matlab file %s\n", argv[1]);
    return (1);
  }

  /*open output file*/
  int cpu_word_size = sizeof(double);
  int io_word_size  = sizeof(double);

  ex_opts(EX_VERBOSE);

  const std::string ext{".exo"};
  std::string       line(argv[1]);
  line = line.substr(0, line.find("."));
  line += ext;
  int exo_file = ex_create(line.c_str(), EX_CLOBBER, &cpu_word_size, &io_word_size);
  if (exo_file < 0) {
    fmt::print(stderr, "MAT2EXO: error creating '{}'\n", line);
    exit(1);
  }

  fmt::print(stderr, "translating '{}' to '{}'\n", argv[1], line);

  int num_axes         = matGetInt("naxes");
  int num_nodes        = matGetInt("nnodes");
  int num_elements     = matGetInt("nelems");
  int num_blocks       = matGetInt("nblks");
  int num_node_sets    = matGetInt("nnsets");
  int num_side_sets    = matGetInt("nssets");
  int num_time_steps   = matGetInt("nsteps");
  int num_global_vars  = matGetInt("ngvars");
  int num_nodal_vars   = matGetInt("nnvars");
  int num_element_vars = matGetInt("nevars");
  int num_nodeset_vars = matGetInt("nnsvars");
  int num_sideset_vars = matGetInt("nssvars");

  ex_put_init(exo_file, line.c_str(), num_axes, num_nodes, num_elements, num_blocks, num_node_sets,
              num_side_sets);

  if (num_global_vars > 0) {
    ex_put_variable_param(exo_file, EX_GLOBAL, num_global_vars);
  }

  if (num_nodal_vars > 0) {
    ex_put_variable_param(exo_file, EX_NODAL, num_nodal_vars);
  }

  if (num_element_vars > 0) {
    ex_put_variable_param(exo_file, EX_ELEM_BLOCK, num_element_vars);
  }

  if (num_nodeset_vars > 0) {
    ex_put_variable_param(exo_file, EX_NODE_SET, num_nodeset_vars);
  }

  if (num_sideset_vars > 0) {
    ex_put_variable_param(exo_file, EX_SIDE_SET, num_sideset_vars);
  }

  /* nodal coordinates */
  {
    std::vector<double> x;
    std::vector<double> y;
    std::vector<double> z;
    matGetDbl("x0", num_nodes, 1, x);
    if (num_axes > 1) {
      matGetDbl("y0", num_nodes, 1, y);
    }
    if (num_axes > 2) {
      matGetDbl("z0", num_nodes, 1, z);
    }
    ex_put_coord(exo_file, x.data(), y.data(), z.data());
  }

  /* side sets */
  std::vector<int> num_sideset_sides(num_side_sets);
  if (num_side_sets > 0) {

    std::vector<int> ids;
    matGetInt("ssids", num_side_sets, 1, ids);
    matGetInt("nsssides", num_side_sets, 1, num_sideset_sides);
    std::vector<int> nssdfac(num_side_sets);
    matGetInt("nssdfac", num_side_sets, 1, nssdfac);

    std::vector<int>    elem_list;
    std::vector<int>    side_list;
    std::vector<double> dist_fact;
    for (int i = 0; i < num_side_sets; i++) {
      ex_put_set_param(exo_file, EX_SIDE_SET, ids[i], num_sideset_sides[i], nssdfac[i]);

      std::string name = fmt::sprintf("sselem%02d", i + 1);
      matGetInt(name, num_sideset_sides[i], 1, elem_list);

      name = fmt::sprintf("ssside%02d", i + 1);
      matGetInt(name, num_sideset_sides[i], 1, side_list);
      ex_put_set(exo_file, EX_SIDE_SET, ids[i], elem_list.data(), side_list.data());

      if (nssdfac[i] > 0) {
        name = fmt::sprintf("ssfac%02d", i + 1);
        matGetDbl(name, nssdfac[i], 1, dist_fact);
        ex_put_set_dist_fact(exo_file, EX_SIDE_SET, ids[i], dist_fact.data());
      }
    }

    get_put_user_names(exo_file, EX_SIDE_SET, num_side_sets, "ssusernames");
  }

  /* node sets */
  std::vector<int> num_nodeset_nodes;
  if (num_node_sets > 0) {

    std::vector<int> ids;
    matGetInt("nsids", num_node_sets, 1, ids);
    matGetInt("nnsnodes", num_node_sets, 1, num_nodeset_nodes);
    std::vector<int> ndfac;
    matGetInt("nnsdfac", num_node_sets, 1, ndfac);

    std::vector<double> dist_fact;
    std::vector<int>    node_list;
    for (int i = 0; i < num_node_sets; i++) {
      ex_put_set_param(exo_file, EX_NODE_SET, ids[i], num_nodeset_nodes[i], ndfac[i]);

      std::string name = fmt::sprintf("nsnod%02d", i + 1);
      matGetInt(name, num_nodeset_nodes[i], 1, node_list);
      ex_put_set(exo_file, EX_NODE_SET, ids[i], node_list.data(), nullptr);

      if (ndfac[i] > 0) {
        name = fmt::sprintf("nsfac%02d", i + 1);
        matGetDbl(name, ndfac[i], 1, dist_fact);
        ex_put_set_dist_fact(exo_file, EX_NODE_SET, ids[i], dist_fact.data());
      }
    }

    get_put_user_names(exo_file, EX_NODE_SET, num_node_sets, "nsusernames");
  }

  /* element blocks */
  std::vector<int> num_elem_in_block(num_blocks);
  {
    std::vector<int> ids;
    matGetInt("blkids", num_blocks, 1, ids);

    /* get elem block types */
    auto block_names = matGetStr("blknames");
    SMART_ASSERT(block_names.size() == (size_t)num_blocks);
    std::vector<int> connect;
    for (int i = 0; i < num_blocks; i++) {
      std::string name              = fmt::sprintf("blk%02d", i + 1);
      int         num_node_per_elem = matArrNRow(name);
      num_elem_in_block[i]          = matArrNCol(name);
      matGetInt(name, num_node_per_elem, num_elem_in_block[i], connect);

      name                  = fmt::sprintf("blk%02d_nattr", i + 1);
      int num_attr_per_elem = matGetInt(name);
      ex_put_block(exo_file, EX_ELEM_BLOCK, ids[i], block_names[i].c_str(), num_elem_in_block[i],
                   num_node_per_elem, 0, 0, num_attr_per_elem);
      ex_put_conn(exo_file, EX_ELEM_BLOCK, ids[i], connect.data(), nullptr, nullptr);

      if (num_attr_per_elem > 0) {
        get_put_attr_names(exo_file, i + 1, ids[i], num_attr_per_elem);
        std::vector<double> attr_data;
        for (int j = 0; j < num_attr_per_elem; j++) {
          name = fmt::sprintf("blk%02d_attr%02d", i + 1, j + 1);
          matGetDbl(name, num_elem_in_block[i], 1, attr_data);
          ex_put_one_attr(exo_file, EX_ELEM_BLOCK, ids[i], j + 1, attr_data.data());
        }
      }
    }
    get_put_user_names(exo_file, EX_ELEM_BLOCK, num_blocks, "blkusernames");
  }

  /* time values */
  if (num_time_steps > 0) {
    std::vector<double> times;
    matGetDbl("time", num_time_steps, 1, times);
    for (int i = 0; i < num_time_steps; i++) {
      ex_put_time(exo_file, i + 1, &times[i]);
    }
  }

  /* global variables */
  if (num_global_vars > 0) {
    get_put_names(exo_file, EX_GLOBAL, num_global_vars, "gnames");

    std::vector<double> var_vals(num_global_vars * num_time_steps);
    std::vector<double> temp(num_time_steps);
    for (int j = 0; j < num_global_vars; j++) {
      std::string name = fmt::sprintf("gvar%02d", j + 1);
      matGetDbl(name, num_time_steps, 1, temp);
      for (int i = 0; i < num_time_steps; i++) {
        var_vals[num_global_vars * i + j] = temp[i];
      }
    }
    for (int i = 0; i < num_time_steps; i++) {
      size_t offset = num_global_vars * i;
      ex_put_var(exo_file, i + 1, EX_GLOBAL, 1, 0, num_global_vars, &var_vals[offset]);
    }
  }

  /* nodal variables */
  if (num_nodal_vars > 0) {
    get_put_names(exo_file, EX_NODAL, num_nodal_vars, "nnames");

    std::vector<int> ids(1, 1);
    std::vector<int> node_block(1, num_nodes);

    get_put_vars(exo_file, EX_NODAL, ids, 1, num_nodal_vars, num_time_steps, node_block,
                 "nvar%02d");
  }

  /* element variables */
  if (num_element_vars > 0) {
    std::vector<int> ids;
    matGetInt("blkids", num_blocks, 1, ids);
    get_put_names(exo_file, EX_ELEM_BLOCK, num_element_vars, "enames");

    get_put_vars(exo_file, EX_ELEM_BLOCK, ids, num_blocks, num_element_vars, num_time_steps,
                 num_elem_in_block, "evar%02d");
  }

  /* nodeset variables */
  if (num_nodeset_vars > 0) {
    std::vector<int> ids;
    matGetInt("nsids", num_node_sets, 1, ids);
    get_put_names(exo_file, EX_NODE_SET, num_nodeset_vars, "nsnames");

    get_put_vars(exo_file, EX_NODE_SET, ids, num_node_sets, num_nodeset_vars, num_time_steps,
                 num_nodeset_nodes, "nsvar%02d");
  }

  /* sideset variables */
  if (num_sideset_vars > 0) {
    std::vector<int> ids;
    matGetInt("ssids", num_side_sets, 1, ids);
    get_put_names(exo_file, EX_SIDE_SET, num_sideset_vars, "ssnames");

    get_put_vars(exo_file, EX_SIDE_SET, ids, num_side_sets, num_sideset_vars, num_time_steps,
                 num_sideset_sides, "ssvar%02d");
  }

  /* node and element number maps */
  {
    std::vector<int> ids;
    if (matGetInt("node_num_map", num_nodes, 1, ids) == 0) {
      ex_put_id_map(exo_file, EX_NODE_MAP, ids.data());
    }
  }

  {
    std::vector<int> ids;
    if (matGetInt("elem_num_map", num_elements, 1, ids) == 0) {
      ex_put_id_map(exo_file, EX_ELEM_MAP, ids.data());
    }
  }

  ex_close(exo_file);
  Mat_Close(mat_file);

  fmt::printf("done.\n");

  add_to_log("mat2exo", 0);
  return (0);
}

/**********************************************************************/
std::vector<std::string> matGetStr(const std::string &name)
{
  matvar_t *matvar = Mat_VarRead(mat_file, name.c_str());
  if (matvar == nullptr) {
    return std::vector<std::string>();
  }

  if (matvar->dims[0] != 1) {
    fmt::printf("Error: Multiline string copy attempted\n");
  }

  size_t bytes = matvar->nbytes;
  if (matvar->data_size == 2 && matvar->data_type == MAT_T_UINT16) {
    // Data stored as 16bit, but we want 8bit (This is due to some UTF
    // strangeness in matio)
    char *data = reinterpret_cast<char *>(matvar->data);
    for (size_t i = 0, j = 0; i < matvar->nbytes; i += 2, j++) {
      data[j] = data[i];
    }
    bytes /= 2;
  }
  std::string mat_names(reinterpret_cast<char *>(matvar->data), bytes);
  auto        names = SLIB::tokenize(mat_names, "\n", true);
  Mat_VarFree(matvar);

  return names;
}

/**********************************************************************/
int matGetDbl(const std::string &name, size_t n1, size_t n2, std::vector<double> &data)
{
  matvar_t *matvar = Mat_VarRead(mat_file, name.c_str());
  if (matvar == nullptr) {
    return -1;
  }

  SMART_ASSERT(matvar->dims[0] == n1);
  SMART_ASSERT(matvar->dims[1] == n2);

  data.resize(n1 * n2);
  memcpy(data.data(), static_cast<int *>(matvar->data), n1 * n2 * sizeof(double));

  Mat_VarFree(matvar);
  return 0;
}

/**********************************************************************/
int matGetInt(const std::string &name, size_t n1, size_t n2, std::vector<int> &data)
{
  matvar_t *matvar = Mat_VarRead(mat_file, name.c_str());
  if (matvar == nullptr) {
    return -1;
  }

  SMART_ASSERT(matvar->dims[0] == n1);
  SMART_ASSERT(matvar->dims[1] == n2);

  data.resize(n1 * n2);
  memcpy(data.data(), static_cast<int *>(matvar->data), n1 * n2 * sizeof(int));

  Mat_VarFree(matvar);
  return 0;
}

/**********************************************************************/
int matGetInt(const std::string &name)
{
  matvar_t *matvar = Mat_VarRead(mat_file, name.c_str());
  if (matvar == nullptr) {
    return -1;
  }

  SMART_ASSERT(matvar->dims[0] == 1);
  SMART_ASSERT(matvar->dims[1] == 1);

  int data = static_cast<int *>(matvar->data)[0];

  Mat_VarFree(matvar);
  return data;
}

/**********************************************************************/
int matArrNRow(const std::string &name)
{
  matvar_t *matvar = Mat_VarRead(mat_file, name.c_str());
  if (matvar == nullptr) {
    return -1;
  }

  int nrow = matvar->dims[0];
  Mat_VarFree(matvar);
  return nrow;
}

/**********************************************************************/
int matArrNCol(const std::string &name)
{
  matvar_t *matvar = Mat_VarRead(mat_file, name.c_str());
  if (matvar == nullptr) {
    return -1;
  }

  int ncol = matvar->dims[1];
  Mat_VarFree(matvar);
  return ncol;
}

void get_put_names(int exo_file, ex_entity_type entity, int num_vars, const std::string &name)
{
  auto names = matGetStr(name);
  SMART_ASSERT(names.size() == (size_t)num_vars);
  std::vector<const char *> str2(num_vars);
  for (int i = 0; i < num_vars; i++) {
    str2[i] = names[i].c_str();
  }
  ex_put_variable_names(exo_file, entity, num_vars, const_cast<char **>(str2.data()));
}

void get_put_user_names(int exo_file, ex_entity_type entity, int num_entity, const char *mname)
{
  auto names = matGetStr(mname);
  SMART_ASSERT(names.size() == (size_t)num_entity)(names.size())(num_entity);
  std::vector<const char *> str2(num_entity);
  for (int i = 0; i < num_entity; i++) {
    str2[i] = names[i].c_str();
  }
  ex_put_names(exo_file, entity, const_cast<char **>(str2.data()));
}

void get_put_attr_names(int exo_file, int seq, int id, int num_attr)
{
  std::string str = fmt::sprintf("blk%02d_attrnames", seq);

  auto names = matGetStr(str);
  SMART_ASSERT(names.size() == (size_t)num_attr);
  std::vector<const char *> str2(num_attr);
  for (int i = 0; i < num_attr; i++) {
    str2[i] = names[i].c_str();
  }
  ex_put_attr_names(exo_file, EX_ELEM_BLOCK, id, const_cast<char **>(str2.data()));
}

void get_put_vars(int exo_file, ex_entity_type type, const std::vector<int> &ids, int num_blocks,
                  int num_vars, int num_time_steps, const std::vector<int> &num_per_block,
                  const char *mname)
{
  size_t num_entity = std::accumulate(num_per_block.begin(), num_per_block.end(), 0);

  for (int i = 0; i < num_vars; i++) {
    std::string name;
    name = fmt::sprintf(mname, i + 1);
    std::vector<double> var_vals;
    matGetDbl(name, num_entity, num_time_steps, var_vals);
    size_t n = 0;
    for (int j = 0; j < num_time_steps; j++) {
      for (int k = 0; k < num_blocks; k++) {
        ex_put_var(exo_file, j + 1, type, i + 1, ids[k], num_per_block[k], &var_vals[n]);
        n += num_per_block[k];
      }
    }
  }
}