//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Device/IO/ReadWriteINT.cpp
//! @brief     Implements functions read|writeBAInt.
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2018
//! @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
//
//  ************************************************************************************************

#include "Device/IO/ReadWriteINT.h"
#include "Base/Axis/Frame.h"
#include "Base/Axis/MakeScale.h"
#include "Base/Axis/Scale.h"
#include "Base/Math/Numeric.h"
#include "Base/Util/Assert.h"
#include "Base/Util/StringUtil.h"
#include "Device/Data/DataUtil.h"
#include "Device/Data/Datafield.h"
#include <iomanip>

namespace {

void writeBlock(const std::vector<double>& values, std::ostream& output_stream, size_t nrows,
                size_t ncols)
{
    output_stream.imbue(std::locale::classic());
    output_stream << std::scientific << std::setprecision(15);

    ASSERT(values.size() == nrows * ncols);
    for (size_t col = 0; col < ncols; ++col) {
        for (size_t row = 0; row < nrows; ++row) {
            if (row)
                output_stream << " ";
            output_stream << Numeric::ignoreDenormalized(values[row * ncols + col]);
        }
        output_stream << std::endl;
    }
}

//! Reads ListScan and GenericScale lines, already parsed as 'type(text)'.
std::vector<double> parse_x_list(std::string text, const std::string& type)
{
    if (text.substr(0, 1) == ",")
        throw std::runtime_error("file format from BornAgain <= 20 is obsolete");
    if (text.substr(0, 1) == "[" && text.substr(text.size() - 1, 1) == "]")
        text = text.substr(1, text.size() - 2);
    std::vector<std::string> arr = Base::String::split(text, ",");
    std::vector<double> result;
    for (std::string e : arr) {
        e = Base::String::trim(e);
        if (e.empty())
            continue;
        double x;
        if (!(Base::String::to_double(e, &x)))
            throw std::runtime_error("Reading " + type + ": cannot read entry '" + e + "'");
        result.push_back(x);
    }
    if (result.empty())
        throw std::runtime_error("Reading " + type + ": found empty list");
    return result;
}

//! Creates axis of certain type from input stream.
Scale* parseScale(std::istream& input_stream)
{
    std::string line;
    std::getline(input_stream, line);

    size_t j = line.find_first_of('(');
    if (j == std::string::npos)
        throw std::runtime_error("Scale constructor has no '('");
    std::string type = line.substr(0, j);

    line = line.substr(j + 1);
    if (line.back() != ')')
        throw std::runtime_error("Scale constructor call not ending with ')'");
    line = line.substr(0, line.size() - 1);

    if (line[0] != '"')
        throw std::runtime_error("Scale constructor arg 1 does not start with \"");
    j = 1 + line.substr(1).find_first_of('"');
    if (j == std::string::npos)
        throw std::runtime_error("Scale constructor arg 1 has no closing \"");
    std::string name = line.substr(1, j - 1);

    if (line[j + 1] != ',')
        throw std::runtime_error("Scale constructor arg 1 not followed by comma");
    std::string body = line.substr(j + 2);
    body = Base::String::trimFront(body, " ");

    if (type == "EquiDivision") {
        std::vector<std::string> arr = Base::String::split(body, ",");
        int nbins;
        double xmi, xma;
        if (!(Base::String::to_int(arr[0], &nbins) && Base::String::to_double(arr[1], &xmi)
              && Base::String::to_double(arr[2], &xma)))
            throw std::runtime_error("Reading EquiDivision: cannot read parameters");
        return EquiDivision(name, nbins, xmi, xma);
    }
    if (type == "ListScan" || type == "DiscreteAxis")
        return ListScan(name, parse_x_list(body, type));

    if (type == "GenericScale")
        return GenericScale(name, parse_x_list(body, type));

    throw std::runtime_error("Unknown axis type '" + type + "'");
}

std::vector<double> readBlock(std::istream& input_stream, size_t nrows, size_t ncols)
{
    std::string line;
    std::vector<double> invec;
    while (std::getline(input_stream, line)) {
        line = Base::String::trim(line);
        if (line.empty())
            break;
        if (line[0] == '#')
            continue;
        for (double x : Base::String::parse_doubles(line))
            invec.emplace_back(x);
    }
    if (invec.empty())
        return {};
    if (invec.size() != nrows * ncols) {
        std::stringstream msg;
        msg << "data block has size " << invec.size() << " instead of expected " << nrows << " x "
            << ncols;
        throw std::runtime_error(msg.str());
    }

    std::vector<double> result(nrows * ncols);
    for (size_t row = 0; row < nrows; ++row)
        for (size_t col = 0; col < ncols; ++col)
            result[col * nrows + row] = invec[row * ncols + col];
    return result;
}

} // namespace


Datafield RW::readBAInt(std::istream& input_stream)
{
    std::string line;

    std::vector<const Scale*> axes;
    while (std::getline(input_stream, line)) {
        line = Base::String::trim(line);

        if (line.find("axis") != std::string::npos)
            axes.emplace_back(::parseScale(input_stream));

        if (line.find("data") != std::string::npos)
            break;
    }
    size_t nrows;
    size_t ncols;
    if (axes.size() == 1) {
        nrows = 1;
        ncols = axes[0]->size();
    } else if (axes.size() == 2) {
        nrows = axes[0]->size();
        ncols = axes[1]->size();
    } else
        ASSERT_NEVER;

    std::vector<double> datvec = ::readBlock(input_stream, nrows, ncols);

    // skip lines before errorbars
    while (std::getline(input_stream, line))
        if (line[0] == '#')
            break;

    std::vector<double> errvec = ::readBlock(input_stream, nrows, ncols);

    return {Frame(axes), datvec, errvec};
}

void RW::writeBAInt(const Datafield& data, std::ostream& output_stream)
{
    output_stream << "# BornAgain Intensity Data\n";

    for (size_t i = 0; i < data.rank(); ++i) {
        const Scale& axis = data.axis(i);
        output_stream << std::endl;
        output_stream << "# axis-" << i << "\n";
        output_stream << axis << "\n";
    }
    size_t ncols = data.axis(0).size();
    size_t nrows = data.rank() == 1 ? 1 : data.axis(1).size();

    output_stream << "\n# data\n";
    ::writeBlock(data.flatVector(), output_stream, nrows, ncols);

    if (data.hasErrorSigmas()) {
        output_stream << "\n# errorbars\n";
        ::writeBlock(data.errorSigmas(), output_stream, nrows, ncols);
    }
    output_stream << std::endl;
}
