#include <pybind11/pybind11.h>
#include <pybind11/pytypes.h>
#include <pybind11/iostream.h>
#include <pybind11/numpy.h>
#include <pybind11/stl_bind.h>

#include <hark-lib/hark-lib.h>
// #include <hark-lib/internal/mutex.h>
#include <hark-lib/datastream.h>
#include <hark-lib/matrix.h>
#include <hark-lib/node.h>

//////////
#include <hark-lib/matrix.h>
#include <hark-lib/tensor.h>
#include <hark-lib/source.h>
#include <hark-lib/timepoint.h>
#include <hark-lib/tf.h>



#include <vector>
#include <map>

#include "offline.h"

namespace py = pybind11;

template<typename T>
hark::data
typed_pyarray_to_data(py::array_t<T> arr, size_t row)
{
    py::buffer_info bi = arr.request();

    hark::data v;
    if(bi.ndim == 2){
        v = hark::data::create<std::vector<T> >(bi.shape[1]);
        auto &vec = v.cast<std::vector<T> >();

        for(ssize_t i = 0; i < bi.shape[1]; i++){
            vec[i] = *arr.data(row, i);
        }

    }
    else if(bi.ndim == 3){
        v = hark::data::create<hark::Matrix<T> >(bi.shape[1], bi.shape[2]);
        auto &mat = v.cast<hark::Matrix<T> >();

        for(ssize_t i = 0; i < bi.shape[1]; i++){
            for(ssize_t j = 0; j < bi.shape[2]; j++){
                mat(i, j) = *arr.data(row, i, j);
            }
        }

    }
    else{
        std::stringstream ss;
        ss << "Unexpected array with ndim=" << bi.ndim << ", expecting 2 or 3";
        throw hark::RuntimeError(ss.str());
    }

    return v;
}

hark::data
pyarray_to_data(py::array arr, size_t row)
{
    py::buffer_info bi = arr.request();

    if(py::dtype::of<float>().is(arr.dtype())){
        return typed_pyarray_to_data<float>(py::cast<py::array_t<float> >(arr), row);

    }
    else if(py::dtype::of<std::complex<float> >().is(arr.dtype())){
        return typed_pyarray_to_data<std::complex<float> >(py::cast<py::array_t<std::complex<float> > >(arr), row);

    }
    else{
        std::stringstream ss;
        std::string dtype(py::str(arr.dtype()));
        ss << "Unexpected array with type=" << dtype << ", expecting float or complex";
        throw hark::RuntimeError(ss.str());
    }
}

// same as object_to_data in hark.cc
/*
hark::data
pyobject_to_data(py::handle obj)
{
    if(false /*obj.is_type<int>()*//*){
        return hark::data::create<int>(obj.cast<int>());

    }
    else{
        return hark::data();
    }
}
*/

template <typename T>
hark::data
py_array_t_to_data(py::handle &object)
{
    auto value = object.cast<py::array_t<T> >();
    const auto &bi = value.request();

    if(bi.ndim == 1){
        std::vector<T> v(value.data(), value.data() + value.size());
        return hark::data(std::move(v));
    }
    else if(bi.ndim == 2){
        hark::Matrix<T> m = hark::Matrix<T>::from_native_image(static_cast<T*>(bi.ptr), bi.shape[0], bi.shape[1]);
        return hark::data(std::move(m));
    }
    else{
        hark::Tensor<T> t = hark::Tensor<T>();
        size_t valid_size = 1;

        t.vsize.resize(bi.shape.size());
        for(size_t i = 0; i < t.vsize.size(); i++){
            t.vsize[i] = bi.shape[i];
            valid_size *= t.vsize[i];
        }

        std::vector<T> v(value.data(), value.data() + value.size());
        std::swap(t.data, v);

        if(valid_size != t.data.size()){
            std::stringstream ss;
            ss << "Unexpected data size: The expected Tensor data size is " << valid_size
               << ", but it was actually " << t.data.size() << ".\n";
            throw hark::RuntimeError(ss.str());
        }
        return hark::data(std::move(t));
    }
    return hark::data();
}

template <typename T>
hark::data
py_list_to_data(py::handle &object)
{
    auto value = object.cast<py::list>();

    std::vector<T> v;
    for(size_t i = 0; i < py::len(value); i++){
        T e = value[i].cast<T>();
        v.emplace_back(e);
    }
    return hark::data(v);
}

static hark::data
object_to_data(py::handle object)
{
    /*
    if(py::isinstance<py::array_t<int16_t> >(object)){
        auto value = object.cast<py::array_t<int16_t> >();
        std::vector<short int> v(value.data(), value.data() + value.size());
        return hark::data(v);
    }
    else if(py::isinstance<py::array_t<float> >(object)){
        auto value = object.cast<py::array_t<float> >();
        const auto &bi = value.request();

        if(bi.ndim == 1){
            std::vector<float> v(value.data(), value.data() + value.size());
            return hark::data(v);
        }else if(bi.ndim == 2){
            hark::Matrix<float> m = hark::Matrix<float>::from_col_major_image(static_cast<float *>(bi.ptr), bi.shape[0], bi.shape[1]);
            return hark::data(m);
        }else{
            throw hark::RuntimeError("Unexpected data type: array<float> with dimensionality more than 2.");
        }
    }
    */
    // array_t -> std::vector<T> , hark::Matrix<T> or hark::Tensor<T>
    if(py::isinstance<py::array_t<int8_t> >(object)){
        return py_array_t_to_data<int8_t>(object);
    }
    else if(py::isinstance<py::array_t<int16_t> >(object)){
        return py_array_t_to_data<int16_t>(object);
    }
    else if(py::isinstance<py::array_t<int32_t> >(object)){
        return py_array_t_to_data<int32_t>(object);
    }
    else if(py::isinstance<py::array_t<int64_t> >(object)){
        return py_array_t_to_data<int64_t>(object);
    }
    else if(py::isinstance<py::array_t<uint8_t> >(object)){
        return py_array_t_to_data<uint8_t>(object);
    }
    else if(py::isinstance<py::array_t<uint16_t> >(object)){
        return py_array_t_to_data<uint16_t>(object);
    }
    else if(py::isinstance<py::array_t<uint32_t> >(object)){
        return py_array_t_to_data<uint32_t>(object);
    }
    else if(py::isinstance<py::array_t<uint64_t> >(object)){
        return py_array_t_to_data<uint64_t>(object);
    }
    else if(py::isinstance<py::array_t<float> >(object)){
        return py_array_t_to_data<float>(object);
    }
    else if(py::isinstance<py::array_t<double> >(object)){
        return py_array_t_to_data<double>(object);
    }
    else if(py::isinstance<py::array_t<std::complex<float> > >(object)){
        return py_array_t_to_data<std::complex<float> >(object);
    }
    else if(py::isinstance<py::array_t<std::complex<double> > >(object)){
        return py_array_t_to_data<std::complex<double> >(object);
    }
    // list -> std::vector<hark::data>
    else if(py::isinstance<py::list>(object)){
        if(py::len(object.cast<py::list>()) == 0){
            std::vector<hark::data> v;
            return hark::data(v);
        }
        else if(py::isinstance<hark::Source>(object.cast<py::list>()[0])){
            auto value = object.cast<py::list>();

            std::vector<hark::data> v;
            for(size_t i = 0; i < py::len(value); i++){
                hark::Source src = value[i].cast<hark::Source>();
                v.emplace_back(hark::data(src));
            }
            return hark::data(v);
        }
        else if(py::isinstance<py::bool_>(object.cast<py::list>()[0])){
            return py_list_to_data<bool>(object);
        }
        else if(py::isinstance<py::int_>(object.cast<py::list>()[0])){
            return py_list_to_data<int64_t>(object);
        }
        else if(py::isinstance<py::float_>(object.cast<py::list>()[0])){
            return py_list_to_data<double>(object);
        }
        /*
        else if(py::isinstance<bool>(object.cast<py::list>()[0])){
            return py_list_to_data<bool>(object);
        }
        else if(py::isinstance<int8_t>(object.cast<py::list>()[0])){
            return py_list_to_data<int8_t>(object);
        }
        else if(py::isinstance<int16_t>(object.cast<py::list>()[0])){
            return py_list_to_data<int16_t>(object);
        }
        else if(py::isinstance<int32_t>(object.cast<py::list>()[0])){
            return py_list_to_data<int32_t>(object);
        }
        else if(py::isinstance<int64_t>(object.cast<py::list>()[0])){
            return py_list_to_data<int64_t>(object);
        }
        else if(py::isinstance<uint8_t>(object.cast<py::list>()[0])){
            return py_list_to_data<uint8_t>(object);
        }
        else if(py::isinstance<uint16_t>(object.cast<py::list>()[0])){
            return py_list_to_data<uint16_t>(object);
        }
        else if(py::isinstance<uint32_t>(object.cast<py::list>()[0])){
            return py_list_to_data<uint32_t>(object);
        }
        else if(py::isinstance<uint64_t>(object.cast<py::list>()[0])){
            return py_list_to_data<uint64_t>(object);
        }
        else if(py::isinstance<float>(object.cast<py::list>()[0])){
            return py_list_to_data<float>(object);
        }
        else if(py::isinstance<double>(object.cast<py::list>()[0])){
            return py_list_to_data<double>(object);
        }
        else if(py::isinstance<std::complex<float> >(object.cast<py::list>()[0])){
            return py_list_to_data<std::complex<float> >(object);
        }
        else if(py::isinstance<std::complex<double> >(object.cast<py::list>()[0])){
            return py_list_to_data<std::complex<double> >(object);
        }
        */
        else{
            throw hark::RuntimeError("Unexpected data type.  To support, implement method harklib::object_to_data(). py::list");
        }
    }
    else if(py::isinstance<py::dict>(object)){
        if(py::len(object.cast<py::dict>()) == 0){
            std::map<int, hark::data> m;
            return hark::data(std::move(m));
        }
        std::map<int, hark::data> m;
        for(std::pair<py::handle, py::handle> item: object.cast<py::dict>()){
            auto key = item.first.cast<int>();
            m[key] = object_to_data(item.second);
        }
        return hark::data(std::move(m));
    }
    else if(py::isinstance<py::str>(object)){
        if(py::len(object.cast<py::str>()) == 0){
            std::string s;
            return hark::data(std::move(s));
        }
        std::string s(object.cast<std::string>());
        return hark::data(std::move(s));
    }
    else if(py::isinstance<py::bool_>(object)){
        return hark::data(std::move(object.cast<bool>()));
    }
    else if(py::isinstance<py::int_>(object)){
        return hark::data(std::move(object.cast<int64_t>()));
    }
    else if(py::isinstance<py::float_>(object)){
        return hark::data(std::move(object.cast<double>()));
    }
    else if(py::isinstance<hark::Source>(object)){
        hark::Source src = object.cast<hark::Source>();
        return hark::data(std::move(src));
    }
    else if(py::isinstance<hark::TransferFunction>(object)){
        hark::TransferFunction tf = object.cast<hark::TransferFunction>();
        return hark::data(std::move(tf));
    }
    else if(py::isinstance<hark::Time>(object)){
        hark::Time t = object.cast<hark::Time>();
        return hark::data(std::move(t));
    }
    else if(py::isinstance<hark::TimeStamp>(object)){
        hark::TimeStamp ts = object.cast<hark::TimeStamp>();
        return hark::data(std::move(ts));
    }
    else if(py::isinstance<hark::data>(object)){
        return object.cast<hark::data>();
    }
    else{
        throw hark::RuntimeError("Unexpected data type.  To support, implement method harklib::object_to_data().");
    }
    return hark::data();
}

hark::data
pyobject_to_data(py::handle obj)
{
    return object_to_data(obj);
}




// same as data_to_object in hark.cc
/*
py::object
data_to_pyobject(hark::data data)
{
    if(data.is_type<int>()){

    }
    else if(data.is_type<std::vector<float> >()){
        const auto &v = data.cast<std::vector<float> >();

        py::array_t<float> w(v.size(), v.data());
        return w;

    }
    else if(data.is_type<hark::Matrix<float> >()){
        const auto &v = data.cast<hark::Matrix<float> >();

        py::array_t<float> w(v.size(), v.data());
        w.resize({v.rows(), v.cols()});
        return w;

    }
    else if(data.is_type<std::vector<std::complex<float> > >()){
        const auto &v = data.cast<std::vector<std::complex<float> > >();

        py::array_t<std::complex<float> > w(v.size(), v.data());
        return w;

    }
    else if(data.is_type<hark::Matrix<std::complex<float> > >()){
        const auto &v = data.cast<hark::Matrix<std::complex<float> > >();

        py::array_t<std::complex<float> > w(v.size(), v.data());
        w.resize({v.rows(), v.cols()});
        return w;

    }

    return py::none();
}
*/
template <typename T>
py::object
vector_to_list_object(hark::data &value)
{
    std::vector<T> &v = value.cast<std::vector<T> >();
    py::list l;
    for(auto it: v){ l.append(py::cast(it)); }

    return l;
}

template <typename T>
py::object
vector_to_array_t_object(hark::data &value)
{
    std::vector<T> &v = value.cast<std::vector<T> >();
    py::array_t<T> a(v.size(), v.data());

    return a;
}

template <typename T>
py::object
matrix_to_array_t_object(hark::data &value)
{
    hark::Matrix<T> &m = value.cast<hark::Matrix<T> >();
    bool row_major = m.native_is_row_major();
    py::buffer_info bi = {
        m.data(),
        sizeof(T),
        py::format_descriptor<T>::format(),
        2,
        { m.rows(), m.cols() },
        { sizeof(T) * (row_major ? m.cols() : 1),
          sizeof(T) * (row_major ? 1 : m.rows())}
    };
    return py::array_t<T>(bi);
}

template <typename T>
py::object
tensor_to_array_t_object(hark::data &value)
{
    hark::Tensor<T> &t = value.cast<hark::Tensor<T> >();
    py::array_t<T> a(t.vsize, t.data.data());

    return a;
}

py::object
data_to_object(hark::data value)
{
    // container data elements are hark::data (recursive)
    if(value.is_type<std::map<int, hark::data> >()){
        auto &s = value.cast<std::map<int, hark::data> >();
        py::dict d;
        for(auto it: s){
            d[py::int_(it.first)] = data_to_object(it.second);
        }
        return d;
    }
    else if(value.is_type<std::vector<hark::data> >()){
        auto &s = value.cast<std::vector<hark::data> >();
        py::list l = py::list(s.size());
        for(size_t i = 0; i < s.size(); i++){
            l[i] = data_to_object(s[i]);
        }
        return l;
    }
    // std::vector<T> -> list or array_t
    else if(value.is_type<std::vector<bool> >()){
        return vector_to_list_object<bool>(value);
    }
    else if(value.is_type<std::vector<int8_t> >()){
        return vector_to_array_t_object<int8_t>(value);
    }
    else if(value.is_type<std::vector<int16_t> >()){
        return vector_to_array_t_object<int16_t>(value);
    }
    else if(value.is_type<std::vector<int32_t> >()){
        return vector_to_array_t_object<int32_t>(value);
    }
    else if(value.is_type<std::vector<int64_t> >()){
        return vector_to_array_t_object<int64_t>(value);
    }
    else if(value.is_type<std::vector<float> >()){
        return vector_to_array_t_object<float>(value);
    }
    else if(value.is_type<std::vector<double> >()){
        return vector_to_array_t_object<double>(value);
    }
    else if(value.is_type<std::vector<std::complex<float> > >()){
        return vector_to_array_t_object<std::complex<float> >(value);
    }
    else if(value.is_type<std::vector<std::complex<double> > >()){
        return vector_to_array_t_object<std::complex<double> >(value);
    }
    // kinoshita-san version
    /*
    else if(value.is_type<hark::Matrix<float> >()){
        hark::Matrix<float> &m = value.cast<hark::Matrix<float>>();

        bool row_major = hark::Matrix<float>::Flags & Eigen::RowMajorBit;

        py::buffer_info bi = {
            m.data(),
            sizeof(float),
            py::format_descriptor<float>::format(),
            2,
            { m.rows(), m.cols() },
            { sizeof(float) * (row_major ? m.cols() : 1),
              sizeof(float) * (row_major ? 1 : m.rows())}
        };

        return py::array_t<float>(bi);
    }
    else if(value.is_type<hark::Matrix<std::complex<float> > >()){
        hark::Matrix<std::complex<float> > &m = value.cast<hark::Matrix<std::complex<float> > >();

        bool row_major = hark::Matrix<std::complex<float> >::Flags & Eigen::RowMajorBit;

        py::buffer_info bi = {
            m.data(),
            sizeof(float),
            py::format_descriptor<float>::format(),
            2,
            { m.rows(), m.cols() },
            { sizeof(float) * (row_major ? m.cols() : 1),
              sizeof(float) * (row_major ? 1 : m.rows())}
        };

        return py::array_t<float>(bi);
    }
    */
    // hark::Matrix<T> -> array_t
    else if(value.is_type<hark::Matrix<float> >()){
        return matrix_to_array_t_object<float>(value);
    }
    else if(value.is_type<hark::Matrix<double> >()){
        return matrix_to_array_t_object<double>(value);
    }
    else if(value.is_type<hark::Matrix<std::complex<float> > >()){
        return matrix_to_array_t_object<std::complex<float> >(value);
    }
    else if(value.is_type<hark::Matrix<std::complex<double> > >()){
        return matrix_to_array_t_object<std::complex<double> >(value);
    }
    // hark::Tensor<T> -> array_t
    else if(value.is_type<hark::Tensor<float> >()){
        return tensor_to_array_t_object<float>(value);
    }
    else if(value.is_type<hark::Tensor<double> >()){
        return tensor_to_array_t_object<double>(value);
    }
    else if(value.is_type<hark::Tensor<std::complex<float> > >()){
        return tensor_to_array_t_object<std::complex<float> >(value);
    }
    else if(value.is_type<hark::Tensor<std::complex<double> > >()){
        return tensor_to_array_t_object<std::complex<double> >(value);
    }
    // hark:: namespace classes
    else if(value.is_type<hark::Source>()){
        return py::cast(value.cast<hark::Source>());
    }
    else if(value.is_type<hark::TransferFunction>()){
        return py::cast(value.cast<hark::TransferFunction>());
    }
    else if(value.is_type<hark::Time>()){
        return py::cast(value.cast<hark::Time>());
    }
    else if(value.is_type<hark::TimeStamp>()){
        return py::cast(value.cast<hark::TimeStamp>());
    }
    else{
        return py::str("Unexpected data");
    }

    return py::none();
}

py::object
data_to_pyobject(hark::data value)
{
    return data_to_object(value);
}




struct array_handle
{
    py::array arr;
    ssize_t   count;
    ssize_t   size;
};

class Receiver
{
public:
    Receiver():
        _is_equally_shaped_array(false)
    {
    }

    void
    push(hark::data v)
    {
        auto o = data_to_pyobject(v);

        if(py::isinstance<py::array>(o)){
            auto arr = o.cast<py::array>();
            auto bi  = arr.request();

            bool    is_dense = true;
            ssize_t st = bi.itemsize;
            for(size_t i = 0; is_dense && i < bi.ndim; i++){
                if(bi.strides[bi.ndim - i - 1] != st){
                    is_dense = false;
                }
                st *= bi.shape[bi.ndim - i - 1];
            }

            if(is_dense){
                if(this->_buffer.empty()){
                    this->_is_equally_shaped_array = true;
                    this->_array_itemsize = bi.itemsize;
                    this->_array_shape = bi.shape;
                    this->_array_strides = bi.strides;
                    this->_array_dtype = arr.dtype();

                }else if(this->_is_equally_shaped_array && 
                            this->_array_itemsize != bi.itemsize ||
                            this->_array_shape != bi.shape ||
                            this->_array_strides != bi.strides ||
                            !this->_array_dtype.is(arr.dtype())){
                    this->_is_equally_shaped_array = false;
                }
            }else{
                this->_is_equally_shaped_array = false;
            }

        }else{
            this->_is_equally_shaped_array = false;
        }

        this->_buffer.push_back(o);
    }

    py::object
    as_pyobject()
    {
        if(this->_is_equally_shaped_array){
            size_t chunk_size = this->_array_itemsize;
            for(auto s: this->_array_shape){
                chunk_size *= s;
            }

            size_t total_size = chunk_size * this->_buffer.size();

            unsigned char *output_buffer = new unsigned char[total_size];
            for(size_t i = 0; i < this->_buffer.size(); i++){
                auto arr = this->_buffer[i].cast<py::array>();
                memcpy(output_buffer + i * chunk_size, arr.data(), chunk_size);
            }

            py::capsule deleter(output_buffer, [](void *f) {
                delete[] reinterpret_cast<unsigned char *>(f);
            });

            std::vector<size_t> shape = {this->_buffer.size()};
            std::vector<size_t> strides = {chunk_size};
            shape.insert(shape.end(), this->_array_shape.begin(), this->_array_shape.end());
            strides.insert(strides.end(), this->_array_strides.begin(), this->_array_strides.end());

            return  py::array(this->_array_dtype, shape, strides, output_buffer, deleter);

        }else{
            py::list v;
            for(auto &o: this->_buffer){
                v.append(o);
            }
            return v;
        }
    }

private:
    std::vector<py::object> _buffer;
    bool _is_equally_shaped_array;

    std::vector<ssize_t> _array_shape, _array_strides;
    ssize_t _array_itemsize;
    py::dtype _array_dtype;

};

py::object
offline_process(hark::NodeDef *nodedef, py::kwargs kwargs)
{
    hark::Node node(nodedef, nullptr);

    std::map<std::string, py::iterator> iterable_input_map;
    std::map<std::string, array_handle> array_input_map;
    std::vector<std::string> stream_proxy_names;
    for (const std::pair<py::handle, py::handle>& kv : kwargs)
    {
        std::string ks = py::str(kv.first);
        std::string vs = py::str(kv.second);
        if(py::isinstance<py::list>(kv.second)){
            iterable_input_map.emplace(py::str(kv.first), py::iterator(py::cast<py::list>(kv.second)));
            stream_proxy_names.push_back(py::str(kv.first));
        }else if (py::isinstance<py::array>(kv.second)){
            py::array arr = py::cast<py::array>(kv.second);
            py::buffer_info bi = arr.request();
            array_handle ah = {arr, 0, bi.shape[0]};
            array_input_map.emplace(py::str(kv.first), ah);
            stream_proxy_names.push_back(py::str(kv.first));
        }else{
            //std::cout << "pararm=" << py::str(kv.first).cast<std::string>() << ", value=" << py::str(kv.second).cast<std::string>() << " @offline_process" << std::endl;
            node.add_input(hark::Input(py::str(kv.first), pyobject_to_data(kv.second)));
        }
    }

    std::vector<std::string> outputs = node.get_output_names();
    std::map<std::string, std::shared_ptr<hark::FunctionalDataSink> > output_sinks;
    std::map<std::string, std::shared_ptr<Receiver> > output_receivers;

    for(auto &name: outputs){
        std::shared_ptr<Receiver> r(new Receiver());
        std::shared_ptr<hark::FunctionalDataSink> sink(new hark::FunctionalDataSink([r](const hark::data &v){
            r->push(v);
        }));
        output_sinks.emplace(name, sink);
        output_receivers.emplace(name, r);

        node.get_output(name).connect_to(sink.get());
    }

    node.initialize();
    while(true){
        if(PyErr_CheckSignals() != 0){
            throw py::error_already_set();
        }

        for(auto &it: iterable_input_map){
            const std::string &name = it.first;
            py::iterator &iter = it.second;

            auto &sink = node.get_input(name);
            while(iter != py::iterator::sentinel() && sink.is_sink_ready()){
                sink.push(pyobject_to_data(*iter));
            }
        }
        for(auto &it: array_input_map){
            const std::string &name = it.first;
            array_handle &ah = it.second;

            auto &sink = node.get_input(name);
            while(ah.count < ah.size && sink.is_sink_ready()){
                sink.push(pyarray_to_data(ah.arr, ah.count));
                ah.count++;
            }
        }

        bool processed = node.try_process();
        if(processed){
            for(auto &name: stream_proxy_names){
                auto &proxy = node.get_input_proxy(name);
                if(proxy.is_synchronous()){
                    proxy.pop();
                }
            }

        }else{
            break;
        }
    }
    node.reset();

    py::object SimpleNamespace = py::module_::import("types").attr("SimpleNamespace");

    auto r = SimpleNamespace();
    for(auto &name: outputs){
        py::object v = output_receivers[name]->as_pyobject();
        py::setattr(r, py::str(name), v);
    }

    return r;
}

