#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 <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());
    }
}

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();
    }
}

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();
}

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{
            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;
}

