// clang-format off
/* ----------------------------------------------------------------------
   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
   https://www.lammps.org/, Sandia National Laboratories
   Steve Plimpton, sjplimp@sandia.gov

   Copyright (2003) Sandia Corporation.  Under the terms of Contract
   DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
   certain rights in this software.  This software is distributed under
   the GNU General Public License.

   See the README file in the top-level LAMMPS directory.
------------------------------------------------------------------------- */

/* ----------------------------------------------------------------------
   Contributing author: Norbert Podhorszki (Oak Ridge National Laboratory)
------------------------------------------------------------------------- */

#include "reader_adios.h"
#include "comm.h"
#include "error.h"
#include "memory.h"
#include <cmath>
#include <cstdlib>
#include <cstring>

#include "adios2.h"
#include "math_const.h"

using namespace LAMMPS_NS;

using namespace MathConst;

// also in read_dump.cpp

enum { ID, TYPE, X, Y, Z, VX, VY, VZ, Q, IX, IY, IZ, FX, FY, FZ };
enum { UNSET, NOSCALE_NOWRAP, NOSCALE_WRAP, SCALE_NOWRAP, SCALE_WRAP };

#define SMALL 1.0e-6

// true if the difference between two floats is "small".
// cannot use fabsf() since it is not fully portable.
static bool is_smalldiff(const float &val1, const float &val2)
{
    return (fabs(static_cast<double>(val1 - val2)) < SMALL);
}

namespace LAMMPS_NS
{
class ReadADIOSInternal
{

public:
    ReadADIOSInternal() {};
    ~ReadADIOSInternal() = default;

    // name of adios group, referrable in adios2_config.xml
    const std::string ioName = "read_dump";
    adios2::ADIOS *ad = nullptr; // adios object
    adios2::IO io;     // adios group of variables and attributes in this dump
    adios2::Engine fh; // adios file/stream handle object
    // ADIOS input variables we need to change every step
    adios2::Variable<uint64_t> varNtimestep;
    adios2::Variable<uint64_t> varNatoms;
    adios2::Variable<double> varAtoms;
    // list of column names for the atom table
    // (individual list of 'columns' string)
    std::vector<std::string> columnNames;
    float timeout = 0.0;
};
} // namespace LAMMPS_NS

/* ---------------------------------------------------------------------- */

ReaderADIOS::ReaderADIOS(LAMMPS *lmp) : Reader(lmp)
{
    fieldindex = nullptr;
    nAtoms = 0;
    nAtomsTotal = 0;
    atomOffset = 0;
    nstep = 0;
    nid = 0;
    me = comm->me;

    internal = new ReadADIOSInternal();
    try {
        internal->ad =
            new adios2::ADIOS("adios2_config.xml", world, adios2::DebugON);
    } catch (std::ios_base::failure &e) {
        char str[256];
        snprintf(str, sizeof(str), "ADIOS initialization failed with error: %s",
                 e.what());
        error->one(FLERR, str);
    }

    /* Define the group holding all variables and attributes  */
    internal->io = internal->ad->DeclareIO(internal->ioName);
}

/* ---------------------------------------------------------------------- */

ReaderADIOS::~ReaderADIOS()
{
    if (me == 0) {
        memory->destroy(fieldindex);
    }
    internal->columnNames.clear();
    if (internal->fh) {
        internal->fh.Close();
    }
    delete internal->ad;
    delete internal;
}

/* ----------------------------------------------------------------------
   pass on settings to find and load the proper plugin
   Called by all processors.
------------------------------------------------------------------------- */
void ReaderADIOS::settings(int narg, char **arg)
{
    int idx = 0;
    while (idx < narg) {
        if (!strcmp(arg[idx], "timeout")) {
            if (idx + 1 < narg) {
                internal->timeout = std::stof(arg[idx + 1]);
                internal->io.SetParameter("OpenTimeoutSecs", arg[idx + 1]);
                ++idx;
            } else {
                char str[128];
                snprintf(str, sizeof(str),
                         "Missing value for 'timeout' option for ADIOS "
                         "read_dump command");
                error->one(FLERR, str);
            }
        }
        ++idx;
    }
}

/* ----------------------------------------------------------------------
   try to open given file
   Every process must call this Collective operation
------------------------------------------------------------------------- */

void ReaderADIOS::open_file(const char *file)
{
    int rv;
    char str[1024];

    // close open file, if needed.
    if (internal->fh)
        internal->fh.Close();

    try {
        internal->fh = internal->io.Open(file, adios2::Mode::Read, world);
    } catch (std::ios_base::failure &e) {
        char str[256];
        snprintf(str, sizeof(str), "%s", e.what());
        error->one(FLERR, str);
    }
    if (!internal->fh) {
        snprintf(str, sizeof(str), "Cannot open file %s using ADIOS", file);
        error->one(FLERR, str);
    }
}

/* ----------------------------------------------------------------------
   close current file
   Every process must call this Collective operation
------------------------------------------------------------------------- */

void ReaderADIOS::close_file()
{
    // close open file, if needed.
    if (internal->fh) {
        internal->fh.Close();
    }
}

/* ----------------------------------------------------------------------
   read and return time stamp from dump file
   if first read reaches end-of-file, return 1 so caller can open next file
   Called by all processors.
------------------------------------------------------------------------- */

int ReaderADIOS::read_time(bigint &ntimestep)
{
    char str[1024];

    adios2::StepStatus status =
        internal->fh.BeginStep(adios2::StepMode::Read, internal->timeout);

    switch (status) {
    case adios2::StepStatus::EndOfStream:
    case adios2::StepStatus::NotReady:
    case adios2::StepStatus::OtherError:
        return 1;
    default:
        break;
    }

    internal->varNtimestep =
        internal->io.InquireVariable<uint64_t>("ntimestep");

    if (!internal->varNtimestep) {
        snprintf(str, sizeof(str),
                 "Did not find 'ntimestep' variable in ADIOS file %s",
                 internal->fh.Name().c_str());
        error->one(FLERR, str);
    }

    ntimestep = static_cast<bigint>(internal->varNtimestep.Max());
    // std::cerr << " ****  ReaderADIOS::read_time found step " << ntimestep
    //          << " **** " << std::endl;
    return 0;
}

/* ----------------------------------------------------------------------
   skip snapshot for given timestamp
    Called by all processors.
------------------------------------------------------------------------- */

void ReaderADIOS::skip() { internal->fh.EndStep(); }

/* ----------------------------------------------------------------------
   read remaining header info:
     return natoms
     box bounds, triclinic (inferred), fieldflag (1 if any fields not found),
     xyz flags = from input scaleflag & wrapflag
   if fieldflag set:
     match Nfield fields to per-atom column labels
     allocate and set fieldindex = which column each field maps to
     fieldtype = X,VX,IZ etc
     fieldlabel = user-specified label or nullptr if use fieldtype default
   xyz flag = scaledflag if has fieldlabel name, else set by x,xs,xu,xsu
   only called by proc 0
------------------------------------------------------------------------- */

bigint ReaderADIOS::read_header(double box[3][3], int &boxinfo, int &triclinic,
                                int fieldinfo, int nfield, int *fieldtype,
                                char **fieldlabel, int scaleflag, int wrapflag,
                                int &fieldflag, int &xflag, int &yflag,
                                int &zflag)
{
    char str[1024];
    nid = 0;

    // signal that we have no box info at all so far.

    internal->varNatoms = internal->io.InquireVariable<uint64_t>("natoms");
    if (!internal->varNatoms) {
        snprintf(str, sizeof(str),
                 "Did not find 'natoms' variable in ADIOS file %s",
                 internal->fh.Name().c_str());
        error->one(FLERR, str);
    }

    /* nAtoms */
    nAtomsTotal = internal->varNatoms.Max();
    uint64_t rem = nAtomsTotal % comm->nprocs;
    nAtoms = nAtomsTotal / comm->nprocs;
    atomOffset = comm->me * nAtoms;
    if (comm->me < rem) {
        ++nAtoms;
        atomOffset += comm->me;
    } else {
        atomOffset += rem;
    }

    /* triclinic */
    adios2::Attribute<int32_t> attTriclinic =
        internal->io.InquireAttribute<int32_t>("triclinic");
    if (!attTriclinic) {
        snprintf(str, sizeof(str),
                 "Did not find 'triclinic' attribute in ADIOS file %s",
                 internal->fh.Name().c_str());
        error->one(FLERR, str);
    }

    triclinic = attTriclinic.Data()[0];

    /* read Box */
    adios2::Variable<double> varBoxxlo =
        internal->io.InquireVariable<double>("boxxlo");
    adios2::Variable<double> varBoxxhi =
        internal->io.InquireVariable<double>("boxxhi");
    adios2::Variable<double> varBoxylo =
        internal->io.InquireVariable<double>("boxylo");
    adios2::Variable<double> varBoxyhi =
        internal->io.InquireVariable<double>("boxyhi");
    adios2::Variable<double> varBoxzlo =
        internal->io.InquireVariable<double>("boxzlo");
    adios2::Variable<double> varBoxzhi =
        internal->io.InquireVariable<double>("boxzhi");

    box[0][0] = varBoxxlo.Max();
    box[0][1] = varBoxxhi.Max();
    box[0][2] = 0.0;
    box[1][0] = varBoxylo.Max();
    box[1][1] = varBoxyhi.Max();
    box[1][2] = 0.0;
    box[2][0] = varBoxzlo.Max();
    box[2][1] = varBoxzhi.Max();
    box[2][2] = 0.0;

    if (triclinic) {
        adios2::Variable<double> varBoxxy =
            internal->io.InquireVariable<double>("boxxy");
        adios2::Variable<double> varBoxxz =
            internal->io.InquireVariable<double>("boxxz");
        adios2::Variable<double> varBoxyz =
            internal->io.InquireVariable<double>("boxyz");

        box[0][2] = varBoxxy.Max();
        box[1][2] = varBoxxz.Max();
        box[2][2] = varBoxyz.Max();
    }

    boxinfo = 1;

    // if no field info requested, just return

    if (!fieldinfo)
        return nAtoms;

    memory->create(fieldindex, nfield, "read_dump:fieldindex");

    /* Columns */
    adios2::Attribute<std::string> attColumns =
        internal->io.InquireAttribute<std::string>("columns");

    std::vector<std::string> labelVector = attColumns.Data();
    int nwords = labelVector.size();
    std::map<std::string, int> labels;
    for (int i = 0; i < nwords; ++i) {
        labels.emplace(labelVector[i], i);
    }

    int s_index, u_index, su_index;
    xflag = UNSET;
    yflag = UNSET;
    zflag = UNSET;

    // copy fieldtype list for supported fields

    for (int i = 0; i < nfield; i++) {
        if (fieldlabel[i]) {
            fieldindex[i] = find_label(fieldlabel[i], labels);
            if (fieldtype[i] == X)
                xflag = 2 * scaleflag + wrapflag + 1;
            else if (fieldtype[i] == Y)
                yflag = 2 * scaleflag + wrapflag + 1;
            else if (fieldtype[i] == Z)
                zflag = 2 * scaleflag + wrapflag + 1;
        }

        else if (fieldtype[i] == ID)
            fieldindex[i] = find_label("id", labels);
        else if (fieldtype[i] == TYPE)
            fieldindex[i] = find_label("type", labels);

        else if (fieldtype[i] == X) {
            fieldindex[i] = find_label("x", labels);
            xflag = NOSCALE_WRAP;
            if (fieldindex[i] < 0) {
                fieldindex[i] = nwords;
                s_index = find_label("xs", labels);
                u_index = find_label("xu", labels);
                su_index = find_label("xsu", labels);
                if (s_index >= 0 && s_index < fieldindex[i]) {
                    fieldindex[i] = s_index;
                    xflag = SCALE_WRAP;
                }
                if (u_index >= 0 && u_index < fieldindex[i]) {
                    fieldindex[i] = u_index;
                    xflag = NOSCALE_NOWRAP;
                }
                if (su_index >= 0 && su_index < fieldindex[i]) {
                    fieldindex[i] = su_index;
                    xflag = SCALE_NOWRAP;
                }
            }
            if (fieldindex[i] == nwords)
                fieldindex[i] = -1;

        } else if (fieldtype[i] == Y) {
            fieldindex[i] = find_label("y", labels);
            yflag = NOSCALE_WRAP;
            if (fieldindex[i] < 0) {
                fieldindex[i] = nwords;
                s_index = find_label("ys", labels);
                u_index = find_label("yu", labels);
                su_index = find_label("ysu", labels);
                if (s_index >= 0 && s_index < fieldindex[i]) {
                    fieldindex[i] = s_index;
                    yflag = SCALE_WRAP;
                }
                if (u_index >= 0 && u_index < fieldindex[i]) {
                    fieldindex[i] = u_index;
                    yflag = NOSCALE_NOWRAP;
                }
                if (su_index >= 0 && su_index < fieldindex[i]) {
                    fieldindex[i] = su_index;
                    yflag = SCALE_NOWRAP;
                }
            }
            if (fieldindex[i] == nwords)
                fieldindex[i] = -1;

        } else if (fieldtype[i] == Z) {
            fieldindex[i] = find_label("z", labels);
            zflag = NOSCALE_WRAP;
            if (fieldindex[i] < 0) {
                fieldindex[i] = nwords;
                s_index = find_label("zs", labels);
                u_index = find_label("zu", labels);
                su_index = find_label("zsu", labels);
                if (s_index >= 0 && s_index < fieldindex[i]) {
                    fieldindex[i] = s_index;
                    zflag = SCALE_WRAP;
                }
                if (u_index >= 0 && u_index < fieldindex[i]) {
                    fieldindex[i] = u_index;
                    zflag = NOSCALE_NOWRAP;
                }
                if (su_index >= 0 && su_index < fieldindex[i]) {
                    fieldindex[i] = su_index;
                    zflag = SCALE_NOWRAP;
                }
            }
            if (fieldindex[i] == nwords)
                fieldindex[i] = -1;

        } else if (fieldtype[i] == VX)
            fieldindex[i] = find_label("vx", labels);
        else if (fieldtype[i] == VY)
            fieldindex[i] = find_label("vy", labels);
        else if (fieldtype[i] == VZ)
            fieldindex[i] = find_label("vz", labels);

        else if (fieldtype[i] == FX)
            fieldindex[i] = find_label("fx", labels);
        else if (fieldtype[i] == FY)
            fieldindex[i] = find_label("fy", labels);
        else if (fieldtype[i] == FZ)
            fieldindex[i] = find_label("fz", labels);

        else if (fieldtype[i] == Q)
            fieldindex[i] = find_label("q", labels);

        else if (fieldtype[i] == IX)
            fieldindex[i] = find_label("ix", labels);
        else if (fieldtype[i] == IY)
            fieldindex[i] = find_label("iy", labels);
        else if (fieldtype[i] == IZ)
            fieldindex[i] = find_label("iz", labels);
    }

    // set fieldflag = -1 if any unfound fields

    fieldflag = 0;
    for (int i = 0; i < nfield; i++)
        if (fieldindex[i] < 0)
            fieldflag = -1;

    return nAtoms;
}

/* ----------------------------------------------------------------------
   read N atom lines from dump file
   stores appropriate values in fields array
   return 0 if success, 1 if error
   only called by proc 0
------------------------------------------------------------------------- */

void ReaderADIOS::read_atoms(int n, int nfield, double **fields)
{
    char str[1024];

    /* Read Atoms */
    /* This is the firsts (and last) read of array data, so we
     * call EndStep() here instead of PerformGets()
     */

    adios2::Variable<double> varAtoms =
        internal->io.InquireVariable<double>("atoms");

    if (n != nAtoms) {
        snprintf(
            str, sizeof(str),
            "ReaderADIOS::read_atoms() expects 'n=%d' equal to the number of "
            "atoms (=%" PRIu64 ") for process %d in ADIOS file %s.",
            n, nAtoms, comm->me, internal->fh.Name().c_str());
        error->one(FLERR, str);
    }

    size_t ncols = varAtoms.Count()[1];
    varAtoms.SetSelection({{atomOffset, 0}, {nAtoms, ncols}});

    std::vector<double> table;

    internal->fh.Get<double>(varAtoms, table);
    // EndStep or PerformGets required to make the read happen
    internal->fh.EndStep();

    size_t idx;
    for (int i = 0; i < nAtoms; i++) {
        idx = i * ncols;
        for (int m = 0; m < nfield; m++) {
            fields[i][m] = table[idx + fieldindex[m]];
        }
    }
}

/* ----------------------------------------------------------------------
   match label to any of N labels
   return index of match or -1 if no match
------------------------------------------------------------------------- */

int ReaderADIOS::find_label(const std::string &label,
                            const std::map<std::string, int> &labels)
{
    std::map<std::string, int>::const_iterator it = labels.find(label);
    if (it != labels.end()) {
        return it->second;
    }
    return -1;
}