//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Tests/Unit/PyBinding/Embedded.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/Py/PyCore.h"
#include "Base/Py/PyFmt.h"
#include "Base/Py/PyUtil.h"
#include "Sample/Multilayer/MultiLayer.h"
#include "Sample/Multilayer/PyImport.h"
#include "Sample/StandardSamples/ExemplarySamples.h"
#include "Sim/Export/ExportToPython.h"
#include "Tests/GTestWrapper/google_test.h"
#include <iostream>
#include <sstream>

//! Importing numpy and accessing its version string.

TEST(Embedded, ImportNumpy)
{
    Py_Initialize();

    PyObject* pmod = PyImport_ImportModule("numpy");
    EXPECT_TRUE(pmod);

    PyObject* pvar = PyObject_GetAttrString(pmod, "__version__");
    Py_DecRef(pmod);
    EXPECT_TRUE(pvar);

    auto version_string = Base::Python::toString(pvar);
    Py_DecRef(pvar);
    std::cout << "numpy_version_string=" << version_string << std::endl;

    Py_Finalize();

    EXPECT_TRUE(!version_string.empty());
}

//! Creating instance of Cylinder and calling its method in embedded Python.

TEST(Embedded, MethodCall)
{
    const double radius(5.0), height(6.0);
    Py_Initialize();

    PyObject* sysPath = PySys_GetObject((char*)"path");
    PyList_Append(sysPath, PyString_FromString("."));
    PyList_Append(sysPath, PyString_FromString(BABuild::buildLibDir().c_str()));

    PyObject* pmod = PyImport_ImportModule("bornagain");
    EXPECT_TRUE(pmod);

    PyObject* pclass = PyObject_GetAttrString(pmod, "Cylinder");
    Py_DecRef(pmod);
    EXPECT_TRUE(pclass);

    PyObject* pargs = Py_BuildValue("(dd)", radius, height);
    EXPECT_TRUE(pargs);

    PyObject* pinst = PyObject_Call(pclass, pargs, nullptr);
    Py_DecRef(pclass);
    Py_DecRef(pargs);
    EXPECT_TRUE(pinst);

    // result of Cylinder
    PyObject* pmeth = PyObject_GetAttrString(pinst, "height");
    Py_DecRef(pinst);
    EXPECT_TRUE(pmeth);

    PyObject* pargs2 = Py_BuildValue("()");
    EXPECT_TRUE(pargs2);

    PyObject* pres = PyObject_Call(pmeth, pargs2, nullptr);
    Py_DecRef(pmeth);
    Py_DecRef(pargs);

    EXPECT_TRUE(pres);

    double value(0);
    EXPECT_TRUE(PyArg_Parse(pres, "d", &value));

    Py_DecRef(pres);

    Py_Finalize();

    EXPECT_TRUE(value == height);
}

//! From https://www.awasu.com/weblog/embedding-python/calling-python-code-from-your-program/

TEST(Embedded, CompiledFunction)
{
    Py_Initialize();

    // compile our function
    std::stringstream buf;
    buf << "def add( n1 , n2 ) :" << std::endl << "    return n1+n2" << std::endl;

    PyObject* pCompiledFn = Py_CompileString(buf.str().c_str(), "", Py_file_input);
    EXPECT_TRUE(pCompiledFn);

    // create a module
    PyObject* pModule = PyImport_ExecCodeModule((char*)"test", pCompiledFn);
    EXPECT_TRUE(pModule);

    // locate the "add" function (it's an attribute of the module)
    PyObject* pAddFn = PyObject_GetAttrString(pModule, "add");
    EXPECT_TRUE(pAddFn);

    // clean up
    Py_DecRef(pAddFn);
    Py_DecRef(pModule);
    Py_DecRef(pCompiledFn);

    // ------------------------
    // using compiled function
    // ------------------------

    // create a new tuple with 2 elements
    PyObject* pPosArgs = PyTuple_New(2);

    PyObject* pVal1 = PyInt_FromLong(10);
    EXPECT_TRUE(pVal1);

    int rc = PyTuple_SetItem(pPosArgs, 0, pVal1); // nb: tuple position 0
    EXPECT_TRUE(rc == 0);

    PyObject* pVal2 = PyInt_FromLong(20);
    EXPECT_TRUE(pVal2);

    rc = PyTuple_SetItem(pPosArgs, 1, pVal2); // nb: tuple position 0
    EXPECT_TRUE(rc == 0);

    // create a new dictionary
    PyObject* pKywdArgs = PyDict_New();
    EXPECT_TRUE(pKywdArgs);

    // call our function
    PyObject* pResult = PyObject_Call(pAddFn, pPosArgs, pKywdArgs);
    EXPECT_TRUE(pResult);

    // convert the result to a string
    PyObject* pResultRepr = PyObject_Repr(pResult);
    std::string result = Base::Python::toString(pResultRepr);
    Py_DecRef(pResultRepr);

    Py_Finalize();

    EXPECT_TRUE(result == "30");
}

//! Creating MultiLayer in Python and extracting object to C++.
//! https://stackoverflow.com/questions/9040669/how-can-i-implement-a-c-class-in-python-to-be-called-by-c/

TEST(Embedded, ObjectExtract)
{
    PyObject* pmod = Base::Python::import_bornagain({BABuild::buildLibDir()});
    EXPECT_TRUE(pmod);

    PyObject* ml = PyObject_GetAttrString(pmod, "MultiLayer");
    Py_DecRef(pmod);
    EXPECT_TRUE(ml);

    PyObject* instance = PyObject_CallFunctionObjArgs(ml, NULL);

    void* argp1 = nullptr;
    swig_type_info* pTypeInfo = SWIG_TypeQuery("MultiLayer *");

    const int res = SWIG_ConvertPtr(instance, &argp1, pTypeInfo, 0);
    EXPECT_TRUE(SWIG_IsOK(res));

    auto* sample = reinterpret_cast<MultiLayer*>(argp1);
    std::string name = sample->className();

    Py_DecRef(instance);
    Py_DecRef(ml);

    Py_Finalize();

    EXPECT_TRUE(name == "MultiLayer");
}

//! Running Python snippet which creates a sample in embedded way.
//! Casting resulting PyObject to C++ MultiLayer.

TEST(Embedded, EmbeddedMultiLayer)
{
    PyObject* pmod = Base::Python::import_bornagain({BABuild::buildLibDir()});
    EXPECT_TRUE(pmod);

    // compile our function
    std::stringstream buf;
    buf << "import bornagain as ba\n";
    buf << "\n";
    buf << "def get_simulation():\n";
    buf << "    m_vacuum = ba.RefractiveMaterial(\"Vacuum\", 0.0, 0.0)\n";
    buf << "    vacuum_layer = ba.Layer(m_vacuum)\n";
    buf << "    sample = ba.MultiLayer()\n";
    buf << "    sample.addLayer(vacuum_layer)\n";
    buf << "    return sample\n";

    PyObject* pCompiledFn = Py_CompileString(buf.str().c_str(), "", Py_file_input);
    EXPECT_TRUE(pCompiledFn);

    // create a module
    PyObject* pModule = PyImport_ExecCodeModule((char*)"test", pCompiledFn);
    EXPECT_TRUE(pModule);

    // locate the "get_simulation" function (it's an attribute of the module)
    PyObject* pAddFn = PyObject_GetAttrString(pModule, "get_simulation");
    EXPECT_TRUE(pAddFn);

    PyObject* instance = PyObject_CallFunctionObjArgs(pAddFn, NULL);
    EXPECT_TRUE(instance);

    // clean up
    Py_DecRef(pAddFn);
    Py_DecRef(pModule);
    Py_DecRef(pCompiledFn);

    void* argp1 = nullptr;
    swig_type_info* pTypeInfo = SWIG_TypeQuery("MultiLayer *");

    const int res = SWIG_ConvertPtr(instance, &argp1, pTypeInfo, 0);
    EXPECT_TRUE(SWIG_IsOK(res));

    auto* sample = reinterpret_cast<MultiLayer*>(argp1);
    size_t n_layers = sample->numberOfLayers();

    Py_DecRef(instance);

    Py_Finalize();

    EXPECT_TRUE(n_layers == 1);
}

//! We use one of our standard sample builders to build a sample, then generate Python snippet
//! using our standard ExportToPython machinery. Ditto for a cloned sample.
//! Two exported code snippets must be identical.

TEST(Embedded, CloneAndExportToPython)
{
    std::unique_ptr<MultiLayer> sample1(ExemplarySamples::createMultiLayerWithNCRoughness());

    const std::string code1 = Py::Export::sampleCode(*sample1);
    std::cout << "Test ExportToPythonAndBack: code1:\n" << code1 << std::endl;

    std::unique_ptr<MultiLayer> sample2(sample1->clone());
    const std::string code2 = Py::Export::sampleCode(*sample2);
    if (code2 != code1)
        std::cout << "Test ExportToPythonAndBack: code2:\n" << code2 << std::endl;
    else
        std::cout << "Test ExportToPythonAndBack: code2 = code1" << std::endl;
    EXPECT_TRUE(code2 == code1);
}

//! We use one of our standard sample builders to build a sample, then generate Python snippet
//! using our standard ExportToPython machinery.
//! Given snippet is compiled and executed in embedded interpretor. Resulting multi layer
//! is casted back to C++ object and used again, to generate code snippet.
//! Two code snippets must be identical.

TEST(Embedded, ExportToPythonAndBack)
{
    std::unique_ptr<MultiLayer> sample1(ExemplarySamples::createCylindersAndPrisms());

    const std::string code1 = Py::Export::sampleCode(*sample1);
    std::cout << "Test ExportToPythonAndBack: code1:\n" << code1 << std::endl;

    const std::string snippet =
        "import bornagain as ba\n" + Py::Fmt::printImportedSymbols(code1) + "\n\n" + code1;
    const auto sample2 =
        Py::Import::createFromPython(snippet, "get_sample", BABuild::buildLibDir());
    const std::string code2 = Py::Export::sampleCode(*sample2);
    if (code2 != code1)
        std::cout << "Test ExportToPythonAndBack: code2:\n" << code2 << std::endl;
    else
        std::cout << "Test ExportToPythonAndBack: code2 = code1" << std::endl;
    EXPECT_TRUE(code2 == code1);
}

//! Retrieves list of functions from the imported script and checks, that there is
//! one function in a dictionary with name "get_simulation".

TEST(Embedded, ModuleFunctionsList)
{
    // compile our function
    std::stringstream buf;
    buf << "import bornagain as ba\n";
    buf << "\n";
    buf << "def get_simulation():\n";
    buf << "    m_vacuum = ba.RefractiveMaterial(\"Vacuum\", 0.0, 0.0)\n";
    buf << "    vacuum_layer = ba.Layer(m_vacuum)\n";
    buf << "    sample = ba.MultiLayer()\n";
    buf << "    sample.addLayer(vacuum_layer)\n";
    buf << "    return sample\n";

    auto listOfFunc = Py::Import::listOfFunctions(buf.str(), BABuild::buildLibDir());
    for (auto s : listOfFunc)
        std::cout << "AAA" << s << std::endl;
    EXPECT_TRUE(listOfFunc.size() == 1 && listOfFunc.at(0) == "get_simulation");
}
