#include <render/imageblock.h>
#include <iostream>
#include <fmt/core.h>

ImageBlock::ImageBlock(const Array2i &offset, const Array2i &blockSize, const ArrayXd &data)
    : m_offset(offset), m_blockSize(blockSize)
{
    m_pixelGenerated = 0;
    m_curPixel = Array2i(0, 0);
    m_pixelTotal = blockSize.prod();
    m_data = decltype(m_data)::Zero(blockSize.x(), blockSize.y());
    std::vector<Spectrum> d = from_tensor_to_spectrum_list(data, m_pixelTotal);
    for (int y = 0; y < m_blockSize.y(); ++y)
    {
        for (int x = 0; x < m_blockSize.x(); ++x)
        {
            m_data(x, y) = d[y * m_blockSize.x() + x];
        }
    }
}

Spectrum ImageBlock::get(const Array2i &pos) const
{
    Array2i offset = pos - m_offset;
    if (offset.x() < 0 || offset.y() < 0 ||
        offset.x() >= m_blockSize.x() || offset.y() >= m_blockSize.y())
    {
        std::cout << "ImageBlock::get: " << offset << " out of bounds" << std::endl;
        return Spectrum(0.0);
    }
    return m_data(offset(0), offset(1));
}

void ImageBlock::put(const ImageBlock &block)
{
    Array2i target_min = m_offset;
    Array2i target_max = m_offset + m_blockSize;

    Array2i source_min = block.m_offset;
    Array2i source_max = block.m_offset + block.m_blockSize;

    Array2i its_min = target_min.cwiseMax(source_min);
    Array2i its_max = (target_max).cwiseMin(source_max);

    Array2i target_offset = (its_min - target_min).cwiseMax(Array2i{0, 0});
    Array2i source_offset = (its_min - target_max).cwiseMax(Array2i{0, 0});
    Array2i block_size = its_max - its_min;

    m_data.block(target_offset(0), target_offset(1),
                 block_size(0), block_size(1)) +=
        block.m_data.block(source_offset(0), source_offset(1),
                           block_size(0), block_size(1));
}

void ImageBlock::put(const Vector2i &_pos, const Spectrum &value)
{
    Array2i pos = _pos.array() - m_offset;
    if (pos.x() < 0 || pos.y() < 0 ||
        pos.x() >= m_blockSize.x() || pos.y() >= m_blockSize.y())
    {
        std::cout << "ImageBlock::put: " << pos << " out of bounds" << std::endl;
        return;
    }
    m_data(pos(0), pos(1)) += value;
}

ArrayXd ImageBlock::flattened() const
{
    std::vector<Spectrum> data(m_pixelTotal);
    for (int y = 0; y < m_blockSize.y(); ++y)
    {
        for (int x = 0; x < m_blockSize.x(); ++x)
        {
            data[y * m_blockSize.x() + x] = m_data(x, y);
        }
    }
    return from_spectrum_list_to_tensor(data, m_pixelTotal);
}

ImageBlock &operator+=(ImageBlock &a, const ImageBlock &b)
{
    a.put(b);
    return a;
}

// py::array_t<Float> ImageBlock::getData()
// {

//     std::vector<ssize_t> shape = {m_data.rows(), m_data.cols(), 3};
//     return py::array_t<Float>(shape,                                                                             // shape
//                               {shape[0] * shape[1] * sizeof(double), shape[1] * sizeof(double), sizeof(double)}, // strides
//                               (Float *)m_data.data());
// }

// void ImageBlock::setData(py::array_t<Float> data)
// {
//     py::buffer_info buffer_info = data.request();

//     // extract data an shape of input array
//     Float *buffer = static_cast<Float *>(buffer_info.ptr);
//     std::vector<ssize_t> shape = buffer_info.shape;
//     if (shape.size() != 3)
//         throw std::runtime_error("Expected 3D array");
//     if (shape[2] != 3)
//         throw std::runtime_error("Expected 3 channels");
//     if (shape[0] != m_data.rows() || shape[1] != m_data.cols())
//         throw std::runtime_error("Expected shape of " + std::to_string(m_data.rows()) + "x" + std::to_string(m_data.cols()));

//     m_data = Eigen::Array<Spectrum, -1, -1>(shape[0], shape[1]);
//     for (int i = 0; i < m_data.rows(); i++)
//     {
//         for (int j = 0; j < m_data.cols(); j++)
//         {
//             buffer[i * m_data.cols() + j]
//             m_data(i, j) = Spectrum();
//         }
//     }
// }