//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Tests/Unit/PyBinding/EmbeddedTest.cpp
//!
//! @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 "BABuild.h"
#include "Base/Const/Units.h"
#include "PyCore/Embed/PyInterpreter.h"
#include "PyCore/Embed/PyObjectPtr.h"
#include "Tests/GTestWrapper/google_test.h"
#include <cstdint>
#include <string>
#include <vector>

//! Importing numpy and accessing its version string.

TEST(Embedded, PyInterpreterTest)
{
    // initialize Python interpreter
    PyInterpreter::initialize();
    EXPECT_TRUE(PyInterpreter::isInitialized());

    // add Python path
    PyInterpreter::addPythonPath("/Some/Extra/Python/Path/");
    EXPECT_FALSE(PyInterpreter::checkError());

    // get runtime info
    std::string runtime_info = PyInterpreter::runtimeInfo();
    std::cout << "Python runtime info:\n" << runtime_info << std::endl;

    // set Python path
    PyInterpreter::setPythonPath("/Some/Extra/Python/Path/");
}


TEST(Embedded, PyInterpreterNumpyTest)
{
    // TODO: re-implement the Numpy Test

    // initialize Python interpreter
    PyInterpreter::initialize();
    EXPECT_TRUE(PyInterpreter::isInitialized());

    // import Numpy
    PyObjectPtr numpy_module = PyInterpreter::import("numpy");
    EXPECT_TRUE(numpy_module.valid());
}


TEST(Embedded, BornAgainPyFunctionTest)
{
    // Test importing a Python script which uses BornAgain Python package

    // initialize Python interpreter
    PyInterpreter::initialize();
    EXPECT_TRUE(PyInterpreter::isInitialized());

    const std::string script{"import bornagain as ba; from bornagain import deg;"
                             "d0 = ba.deg; d1 = deg; "
                             "get_sample = lambda: (d0 == d1 and 'BornAgain.deg = %.3f' % d1)"};
    const std::string functionName{"get_sample"};

    // locate the `get_sample` function (it is an attribute of the module)
    PyObjectPtr ret{PyInterpreter::BornAgain::callScriptFunction(functionName, script, "")};

    if (!ret.valid())
        throw std::runtime_error("Failed executing Python function '" + functionName + "'");

    const std::string return_str = PyInterpreter::pyStrtoString(ret.get());
    std::stringstream _stream;
    _stream << std::fixed << std::setprecision(3) << Units::deg;
    const std::string expected_str = "BornAgain.deg = " + _stream.str();

    // verify that the returned string starts with the expected string
    if (!(return_str == expected_str))
        throw std::runtime_error("Result '" + return_str + "' does not match the expected '"
                                 + expected_str + "'");
}


TEST(Embedded, BornAgainPyFabioTest)
{
    // Test importing a Python script which uses BornAgain Python package

    // initialize Python interpreter
    PyInterpreter::initialize();
    EXPECT_TRUE(PyInterpreter::isInitialized());

    // import Fabio
    PyObjectPtr fabio_module = PyInterpreter::Fabio::import();
    EXPECT_TRUE(fabio_module.valid());

    // mar format ([2300, 2300], dtype = uint32)
    const std::string datafile_mar{BABuild::testdataDir() + "/fabio/200mMmgso4_001.mar2300"};
    std::cout << "BornAgainPyFabioTest: datafile '" << datafile_mar << std::endl;

    NumpyArrayWrapper np_arr1 = PyInterpreter::Fabio::load(datafile_mar);
    const ArrayMetadata& metadata1 = np_arr1.metadata();
    const std::size_t num_pixels = std::stoi(std::get<std::string>(metadata1.at("NumPixels")));
    const std::size_t n_rows1 = np_arr1.dimensions()[0], n_columns1 = np_arr1.dimensions()[1];
    std::cout << "Numpy array size = " << np_arr1.size() << ", "
              << "n_dimensions = " << np_arr1.rank() << ", "
              << "dtype = " << np_arr1.dtype() << ", "
              << "nr of pixels = " << num_pixels << ", "
              << "rows = " << n_rows1 << ", "
              << "columns = " << n_columns1 << std::endl;

    // check the nominal and actual number of pixels
    EXPECT_TRUE(num_pixels == np_arr1.size());

    auto arr1 = reinterpret_cast<const std::uint32_t*>(np_arr1.arrayPtr());
    const std::uint32_t arr1_test[] = {38, 21, 24}; // expected result
    const std::size_t i_r = 249, offset_r = i_r * n_columns1;
    for (std::size_t i_c = 1095, idx = 0; i_c < 1098; ++i_c) {
        const std::size_t offset = offset_r + i_c;
        const std::uint32_t val = arr1[offset], val_expected = arr1_test[idx];
        ++idx;

        std::printf("data[%zu, %zu] = %u =?= %u\n", i_r, i_c, val, val_expected);
        EXPECT_TRUE(val == val_expected);
    }

    // tif format ([480, 640], dtype = uint8)
    const std::string datafile_tif{BABuild::testdataDir() + "/fabio/at3_1m4_01.tif"};
    std::cout << "BornAgainPyFabioTest: datafile '" << datafile_tif << std::endl;

    NumpyArrayWrapper np_arr2 = PyInterpreter::Fabio::load(datafile_tif);
    const ArrayMetadata& metadata2 = np_arr2.metadata();
    const unsigned long n_rows2 = std::get<long int>(metadata2.at("nRows")),
                        n_columns2 = std::get<long int>(metadata2.at("nColumns")),
                        n_pixels2 = n_rows2 * n_columns2;

    std::cout << "Numpy array size = " << np_arr2.size() << ", "
              << "n_dimensions = " << np_arr2.rank() << ", "
              << "dtype = " << np_arr2.dtype() << ", "
              << "nr of rows = " << n_rows2 << ", "
              << "nr of columns = " << n_columns2 << std::endl;

    EXPECT_TRUE(n_pixels2 == np_arr2.size());

    auto arr2 = reinterpret_cast<const std::uint8_t*>(np_arr2.arrayPtr());
    const std::uint8_t arr2_test[] = {70, 72, 69, 67, 66, 68, 65, 67, 68}; // expected result
    for (std::size_t i_r = 200, idx = 0; i_r < 203; ++i_r) {
        const std::size_t offset_r = i_r * n_columns2;
        for (std::size_t i_c = 300; i_c < 303; ++i_c) {
            const std::size_t offset = offset_r + i_c;
            const std::uint8_t val = arr2[offset], val_expected = arr2_test[idx];
            ++idx;

            std::printf("data[%zu, %zu] = %u =?= %u\n", i_r, i_c, val, val_expected);
            EXPECT_TRUE(val == val_expected);
        }
    }
}
