//
// SPDX-License-Identifier: BSD-3-Clause
// Copyright (c) Weta Digital, Ltd and Contributors to the OpenEXR Project.
//

#include "ImfCompositeDeepScanLine.h"
#include "IlmThreadPool.h"
#include "ImfChannelList.h"
#include "ImfDeepCompositing.h"
#include "ImfDeepFrameBuffer.h"
#include "ImfDeepScanLineInputFile.h"
#include "ImfDeepScanLineInputPart.h"
#include "ImfFrameBuffer.h"
#include "ImfPixelType.h"

#include <Iex.h>
#include <stddef.h>
#include <vector>
OPENEXR_IMF_INTERNAL_NAMESPACE_SOURCE_ENTER

using ILMTHREAD_NAMESPACE::Task;
using ILMTHREAD_NAMESPACE::TaskGroup;
using ILMTHREAD_NAMESPACE::ThreadPool;
using IMATH_NAMESPACE::Box2i;
using std::string;
using std::vector;

struct CompositeDeepScanLine::Data
{
public:
    vector<DeepScanLineInputFile*> _file; // array of files
    vector<DeepScanLineInputPart*> _part; // array of parts
    FrameBuffer _outputFrameBuffer;       // output frame buffer provided
    bool _zback; // true if we are using zback (otherwise channel 1 = channel 0)
    vector<vector<float>>
        _channeldata; // pixel values, read from the input, one array per channel
    vector<int>      _sampleCounts; // total per-pixel sample counts,
    Box2i            _dataWindow;   // data window of combined inputs
    DeepCompositing* _comp;         // user-provided compositor
    vector<string>   _channels;     // names of channels that will be composited
    vector<int>
        _bufferMap; // entry _outputFrameBuffer[n].name() == _channels[ _bufferMap[n] ].name()

    void check_valid (
        const Header&
            header); // check newly added part/file is OK; on first good call, set _zback/_dataWindow

    //
    // set up the given deep frame buffer to contain the required channels
    // resize counts and pointers to the width of _dataWindow
    // zero-out all counts, since the datawindow may be smaller than/not include this part
    //

    void handleDeepFrameBuffer (
        DeepFrameBuffer&      buf,
        vector<unsigned int>& counts, //per-pixel counts
        vector<vector<float*>>&
                      pointers, //per-channel-per-pixel pointers to data
        const Header& header,
        int           start,
        int           end);

    Data ();
};

CompositeDeepScanLine::Data::Data () : _zback (false), _comp (NULL)
{}

CompositeDeepScanLine::CompositeDeepScanLine () : _Data (new Data)
{}

CompositeDeepScanLine::~CompositeDeepScanLine ()
{
    delete _Data;
}

void
CompositeDeepScanLine::addSource (DeepScanLineInputPart* part)
{
    _Data->check_valid (part->header ());
    _Data->_part.push_back (part);
}

void
CompositeDeepScanLine::addSource (DeepScanLineInputFile* file)
{
    _Data->check_valid (file->header ());
    _Data->_file.push_back (file);
}

int
CompositeDeepScanLine::sources () const
{
    return int (_Data->_part.size ()) + int (_Data->_file.size ());
}

void
CompositeDeepScanLine::Data::check_valid (const Header& header)
{

    bool has_z     = false;
    bool has_alpha = false;
    // check good channel names
    for (ChannelList::ConstIterator i = header.channels ().begin ();
         i != header.channels ().end ();
         ++i)
    {
        std::string n (i.name ());
        if (n == "ZBack") { _zback = true; }
        else if (n == "Z") { has_z = true; }
        else if (n == "A") { has_alpha = true; }
    }

    if (!has_z)
    {
        throw IEX_NAMESPACE::ArgExc (
            "Deep data provided to CompositeDeepScanLine is missing a Z channel");
    }

    if (!has_alpha)
    {
        throw IEX_NAMESPACE::ArgExc (
            "Deep data provided to CompositeDeepScanLine is missing an alpha channel");
    }

    if (_part.size () == 0 && _file.size () == 0)
    {
        // first in - update and return

        _dataWindow = header.dataWindow ();

        return;
    }

    const Header* const match_header = _part.size () > 0 ? &_part[0]->header ()
                                                         : &_file[0]->header ();

    // check the sizes match
    if (match_header->displayWindow () != header.displayWindow ())
    {
        throw IEX_NAMESPACE::ArgExc (
            "Deep data provided to CompositeDeepScanLine has a different displayWindow to previously provided data");
    }

    _dataWindow.extendBy (header.dataWindow ());
}
void
CompositeDeepScanLine::Data::handleDeepFrameBuffer (
    DeepFrameBuffer&             buf,
    std::vector<unsigned int>&   counts,
    vector<std::vector<float*>>& pointers,
    const Header&                header,
    int                          start,
    int                          end)
{
    ptrdiff_t width      = _dataWindow.size ().x + 1;
    size_t    pixelcount = width * (end - start + 1);
    pointers.resize (_channels.size ());
    counts.resize (pixelcount);
    buf.insertSampleCountSlice (Slice (
        OPENEXR_IMF_INTERNAL_NAMESPACE::UINT,
        (char*) (&counts[0] - _dataWindow.min.x - start * width),
        sizeof (unsigned int),
        sizeof (unsigned int) * width));

    pointers[0].resize (pixelcount);
    buf.insert (
        "Z",
        DeepSlice (
            OPENEXR_IMF_INTERNAL_NAMESPACE::FLOAT,
            (char*) (&pointers[0][0] - _dataWindow.min.x - start * width),
            sizeof (float*),
            sizeof (float*) * width,
            sizeof (float)));

    if (_zback)
    {
        pointers[1].resize (pixelcount);
        buf.insert (
            "ZBack",
            DeepSlice (
                OPENEXR_IMF_INTERNAL_NAMESPACE::FLOAT,
                (char*) (&pointers[1][0] - _dataWindow.min.x - start * width),
                sizeof (float*),
                sizeof (float*) * width,
                sizeof (float)));
    }

    pointers[2].resize (pixelcount);
    buf.insert (
        "A",
        DeepSlice (
            OPENEXR_IMF_INTERNAL_NAMESPACE::FLOAT,
            (char*) (&pointers[2][0] - _dataWindow.min.x - start * width),
            sizeof (float*),
            sizeof (float*) * width,
            sizeof (float)));

    size_t i = 0;
    for (FrameBuffer::ConstIterator qt = _outputFrameBuffer.begin ();
         qt != _outputFrameBuffer.end ();
         qt++)
    {
        int channel_in_source = _bufferMap[i];
        if (channel_in_source > 2)
        {
            // not dealt with yet (0,1,2 previously inserted)
            pointers[channel_in_source].resize (pixelcount);
            buf.insert (
                qt.name (),
                DeepSlice (
                    OPENEXR_IMF_INTERNAL_NAMESPACE::FLOAT,
                    (char*) (&pointers[channel_in_source][0] -
                             _dataWindow.min.x - start * width),
                    sizeof (float*),
                    sizeof (float*) * width,
                    sizeof (float)));
        }

        i++;
    }
}

void
CompositeDeepScanLine::setCompositing (DeepCompositing* c)
{
    _Data->_comp = c;
}

const IMATH_NAMESPACE::Box2i&
CompositeDeepScanLine::dataWindow () const
{
    return _Data->_dataWindow;
}

void
CompositeDeepScanLine::setFrameBuffer (const FrameBuffer& fr)
{

    //
    // count channels; build map between channels in frame buffer
    // and channels in internal buffers
    //

    _Data->_channels.resize (3);
    _Data->_channels[0] = "Z";
    _Data->_channels[1] = _Data->_zback ? "ZBack" : "Z";
    _Data->_channels[2] = "A";
    _Data->_bufferMap.resize (0);

    for (FrameBuffer::ConstIterator q = fr.begin (); q != fr.end (); q++)
    {

        //
        // Frame buffer must have xSampling and ySampling set to 1
        // (Sampling in FrameBuffers must match sampling in file,
        //  and Header::sanityCheck enforces sampling in deep files is 1)
        //

        if (q.slice ().xSampling != 1 || q.slice ().ySampling != 1)
        {
            THROW (
                IEX_NAMESPACE::ArgExc,
                "X and/or y subsampling factors "
                "of \""
                    << q.name ()
                    << "\" channel in framebuffer "
                       "are not 1");
        }

        string name (q.name ());
        if (name == "ZBack") { _Data->_bufferMap.push_back (1); }
        else if (name == "Z") { _Data->_bufferMap.push_back (0); }
        else if (name == "A") { _Data->_bufferMap.push_back (2); }
        else
        {
            _Data->_bufferMap.push_back (
                static_cast<int> (_Data->_channels.size ()));
            _Data->_channels.push_back (name);
        }
    }

    _Data->_outputFrameBuffer = fr;
}

namespace
{

class LineCompositeTask : public Task
{
public:
    LineCompositeTask (
        TaskGroup*                      group,
        CompositeDeepScanLine::Data*    data,
        int                             y,
        int                             start,
        vector<const char*>*            names,
        vector<vector<vector<float*>>>* pointers,
        vector<unsigned int>*           total_sizes,
        vector<unsigned int>*           num_sources)
        : Task (group)
        , _Data (data)
        , _y (y)
        , _start (start)
        , _names (names)
        , _pointers (pointers)
        , _total_sizes (total_sizes)
        , _num_sources (num_sources)
    {}

    virtual ~LineCompositeTask () {}

    virtual void                    execute ();
    CompositeDeepScanLine::Data*    _Data;
    int                             _y;
    int                             _start;
    vector<const char*>*            _names;
    vector<vector<vector<float*>>>* _pointers;
    vector<unsigned int>*           _total_sizes;
    vector<unsigned int>*           _num_sources;
};

void
composite_line (
    int                                   y,
    int                                   start,
    CompositeDeepScanLine::Data*          _Data,
    vector<const char*>&                  names,
    const vector<vector<vector<float*>>>& pointers,
    const vector<unsigned int>&           total_sizes,
    const vector<unsigned int>&           num_sources)
{
    vector<float> output_pixel (names.size ()); //the pixel we'll output to
    vector<const float*> inputs (names.size ());
    DeepCompositing      d; // fallback compositing engine
    DeepCompositing*     comp = _Data->_comp ? _Data->_comp : &d;

    int pixel =
        (y - start) * (_Data->_dataWindow.max.x + 1 - _Data->_dataWindow.min.x);

    for (int x = _Data->_dataWindow.min.x; x <= _Data->_dataWindow.max.x; x++)
    {
        // set inputs[] to point to the first sample of the first part of each channel
        // if there's a zback, set all channel independently...

        if (_Data->_zback)
        {

            for (size_t channel = 0; channel < names.size (); channel++)
            {
                inputs[channel] = pointers[0][channel][pixel];
            }
        }
        else
        {

            // otherwise, set 0 and 1 to point to Z

            inputs[0] = pointers[0][0][pixel];
            inputs[1] = pointers[0][0][pixel];
            for (size_t channel = 2; channel < names.size (); channel++)
            {
                inputs[channel] = pointers[0][channel][pixel];
            }
        }
        comp->composite_pixel (
            &output_pixel[0],
            &inputs[0],
            &names[0],
            static_cast<int> (names.size ()),
            total_sizes[pixel],
            num_sources[pixel]);

        size_t channel_number = 0;

        //
        // write out composited value into internal frame buffer
        //
        for (FrameBuffer::Iterator it = _Data->_outputFrameBuffer.begin ();
             it != _Data->_outputFrameBuffer.end ();
             it++)
        {

            float value = output_pixel
                [_Data->_bufferMap[channel_number]]; // value to write
            intptr_t base = reinterpret_cast<intptr_t> (it.slice ().base);

            // cast to half float if necessary
            if (it.slice ().type == OPENEXR_IMF_INTERNAL_NAMESPACE::FLOAT)
            {
                float* ptr = reinterpret_cast<float*> (
                    base + y * it.slice ().yStride + x * it.slice ().xStride);
                *ptr = value;
            }
            else if (it.slice ().type == HALF)
            {
                half* ptr = reinterpret_cast<half*> (
                    base + y * it.slice ().yStride + x * it.slice ().xStride);
                *ptr = half (value);
            }

            channel_number++;
        }

        pixel++;

    } // next pixel on row
}

void
LineCompositeTask::execute ()
{
    composite_line (
        _y, _start, _Data, *_names, *_pointers, *_total_sizes, *_num_sources);
}

} // namespace

namespace
{
int64_t maximumSampleCount = 0;
}

void
CompositeDeepScanLine::setMaximumSampleCount (int64_t c)
{
    maximumSampleCount = c;
}

int64_t
CompositeDeepScanLine::getMaximumSampleCount ()
{
    return maximumSampleCount;
}

void
CompositeDeepScanLine::readPixels (int start, int end)
{
    size_t parts =
        _Data->_file.size () + _Data->_part.size (); // total of files+parts

    vector<DeepFrameBuffer>      framebuffers (parts);
    vector<vector<unsigned int>> counts (parts);

    //
    // for each part, a pointer to an array of channels
    //
    vector<vector<vector<float*>>> pointers (parts);
    vector<const Header*>          headers (parts);

    {
        size_t i;
        for (i = 0; i < _Data->_file.size (); i++)
        {
            headers[i] = &_Data->_file[i]->header ();
        }

        for (size_t j = 0; j < _Data->_part.size (); j++)
        {
            headers[i + j] = &_Data->_part[j]->header ();
        }
    }

    for (size_t i = 0; i < parts; i++)
    {
        _Data->handleDeepFrameBuffer (
            framebuffers[i], counts[i], pointers[i], *headers[i], start, end);
    }

    //
    // set frame buffers and read scanlines from all parts
    // TODO what happens if SCANLINE not in data window?
    //

    {
        size_t i = 0;
        for (i = 0; i < _Data->_file.size (); i++)
        {
            _Data->_file[i]->setFrameBuffer (framebuffers[i]);
            _Data->_file[i]->readPixelSampleCounts (start, end);
        }
        for (size_t j = 0; j < _Data->_part.size (); j++)
        {
            _Data->_part[j]->setFrameBuffer (framebuffers[i + j]);
            _Data->_part[j]->readPixelSampleCounts (start, end);
        }
    }

    //
    //  total width
    //

    size_t               total_width  = _Data->_dataWindow.size ().x + 1;
    size_t               total_pixels = total_width * (end - start + 1);
    vector<unsigned int> total_sizes (total_pixels);
    vector<unsigned int> num_sources (
        total_pixels); //number of parts with non-zero sample count

    int64_t overall_sample_count =
        0; // sum of all samples in all images between start and end

    //
    // accumulate pixel counts
    //
    for (size_t ptr = 0; ptr < total_pixels; ptr++)
    {
        total_sizes[ptr] = 0;
        num_sources[ptr] = 0;
        for (size_t j = 0; j < parts; j++)
        {
            if (total_sizes[ptr] > std::numeric_limits<unsigned int>::max() - counts[j][ptr])
                throw IEX_NAMESPACE::ArgExc (
                    "Cannot composite scanline: pixel cannot have more than UINT_MAX samples");
                
            total_sizes[ptr] += counts[j][ptr];
            if (counts[j][ptr] > 0) num_sources[ptr]++;
        }
        overall_sample_count += total_sizes[ptr];
    }

    if (maximumSampleCount > 0 && overall_sample_count > maximumSampleCount)
    {
        throw IEX_NAMESPACE::ArgExc (
            "Cannot composite scanline: total sample count on scanline exceeds "
            "limit set by CompositeDeepScanLine::setMaximumSampleCount()");
    }

    //
    // allocate arrays for pixel data
    // samples array accessed as in pixels[channel][sample]
    //

    vector<vector<float>> samples (_Data->_channels.size ());

    for (size_t channel = 0; channel < _Data->_channels.size (); channel++)
    {
        if (channel != 1 || _Data->_zback)
        {
            samples[channel].resize (overall_sample_count);
        }
    }

    for (size_t channel = 0; channel < samples.size (); channel++)
    {

        if (channel != 1 || _Data->_zback)
        {

            samples[channel].resize (overall_sample_count);

            //
            // allocate pointers for channel data
            //

            int64_t offset = 0;

            for (size_t pixel = 0; pixel < total_pixels; pixel++)
            {
                for (size_t part = 0;
                     part < parts && offset < overall_sample_count;
                     part++)
                {
                    pointers[part][channel][pixel] = &samples[channel][offset];
                    offset += counts[part][pixel];
                }
            }
        }
    }

    //
    // read data
    //

    for (size_t i = 0; i < _Data->_file.size (); i++)
    {
        _Data->_file[i]->readPixels (start, end);
    }
    for (size_t j = 0; j < _Data->_part.size (); j++)
    {
        _Data->_part[j]->readPixels (start, end);
    }

    //
    // composite pixels and write back to framebuffer
    //

    // turn vector of strings into array of char *
    // and make sure 'ZBack' channel is correct
    vector<const char*> names (_Data->_channels.size ());
    for (size_t i = 0; i < names.size (); i++)
    {
        names[i] = _Data->_channels[i].c_str ();
    }

    if (!_Data->_zback)
        names[1] = names[0]; // no zback channel, so make it point to z

    TaskGroup g;
    for (int y = start; y <= end; y++)
    {
        ThreadPool::addGlobalTask (new LineCompositeTask (
            &g,
            _Data,
            y,
            start,
            &names,
            &pointers,
            &total_sizes,
            &num_sources));
    } //next row
}

const FrameBuffer&
CompositeDeepScanLine::frameBuffer () const
{
    return _Data->_outputFrameBuffer;
}

OPENEXR_IMF_INTERNAL_NAMESPACE_SOURCE_EXIT
