//
// SPDX-License-Identifier: BSD-3-Clause
// Copyright (c) Contributors to the OpenEXR Project.
//

//-----------------------------------------------------------------------------
//
//    Code examples that show how class MultiPartInputFile and
//    class MultiPartOutputFile can be used to read and write
//    OpenEXR image files containing multiple images.
//
//-----------------------------------------------------------------------------

#include <ImfHeader.h>
#include <ImfPartType.h>
#include <ImfChannelList.h>
#include <ImfArray.h>

#include <ImfFrameBuffer.h>
#include <ImfDeepFrameBuffer.h>

#include <ImfMultiPartInputFile.h>
#include <ImfMultiPartOutputFile.h>

#include <ImfInputPart.h>
#include <ImfTiledInputPart.h>
#include <ImfDeepScanLineInputPart.h>
#include <ImfDeepTiledInputPart.h>

#include <ImfOutputPart.h>
#include <ImfTiledOutputPart.h>
#include <ImfDeepScanLineOutputPart.h>
#include <ImfDeepTiledOutputPart.h>

#include <vector>
#include <list>

#include "namespaceAlias.h"
using namespace IMF;
using namespace std;
using namespace IMATH_NAMESPACE;


template <class InputPartType, class OutputPartType>
void copyPixels (
    MultiPartInputFile & inputFile,
    MultiPartOutputFile & outputFile,
    int inputPartNumber,
    int outputPartNumber)
{
    InputPartType  inPart  (inputFile, inputPartNumber);
    OutputPartType outPart (outputFile, outputPartNumber);
    outPart.copyPixels (inPart);
}

void copyPixels(
    const std::string & type,
    MultiPartInputFile & inputFile,
    MultiPartOutputFile & outputFile,
    int inputPartNumber,
    int outputPartNumber)
{
    //
    // Copy pixels from a given part of a multipart input file
    // to a given part of a multipart output file.
    //
    
    if (type == SCANLINEIMAGE)
    {
        copyPixels<InputPart, OutputPart> (inputFile, outputFile, inputPartNumber, outputPartNumber);
    }
    else if (type == TILEDIMAGE)
    {
        copyPixels<TiledInputPart, TiledOutputPart> (inputFile, outputFile, inputPartNumber, outputPartNumber);
    }
    else if (type == DEEPSCANLINE)
    {
        copyPixels<DeepScanLineInputPart, DeepScanLineOutputPart> (inputFile, outputFile, inputPartNumber, outputPartNumber);
    }
    else if (type == DEEPTILE)
    {
        copyPixels<DeepTiledInputPart, DeepTiledOutputPart> (inputFile, outputFile, inputPartNumber, outputPartNumber);
    }
}

void combineFiles ()
{
    //
    // Read multiple single-part input files and write them as a multi-part file.
    // If an input file is multi-part, only one part is copied.
    // All input files dimentions must be the same.
    //
    
    std::vector<MultiPartInputFile *> inputFiles;
    std::vector<Header> headers;
    
    const char * filenames[] = {"gz1.exr", "tiledgz1.exr", "test.deep.exr", "testTiled.deep.exr"};
    
    for (size_t i = 0; i < sizeof (filenames) / sizeof (filenames[0]); i++)
    {
        MultiPartInputFile * in_file = new MultiPartInputFile (filenames[i]);
        Header header = in_file->header (0);
        
        if (!header.hasName ())
        {
            header.setName (filenames[i]);
        }
        
        inputFiles.push_back (in_file);
        headers.push_back (header);
    }
    
    MultiPartOutputFile outputFile ("multipart.exr", headers.data (), (int)headers.size ());
    
    for (size_t i = 0; i < sizeof (filenames) / sizeof (filenames[0]); i++)
    {
        Header & header = headers[i];
        const string& type = header.type ();
        copyPixels (type, *inputFiles[i], outputFile, 0, i);
        delete inputFiles[i];
    }
}

void splitFiles ()
{
    //
    // Read a multi-part input file and write all parts as separate files.
    //
    
    MultiPartInputFile inputFile ("modified.exr");
    
    for (int i = 0; i < inputFile.parts (); i++)
    {
        Header header = inputFile.header (i);
        string out_path = string ("split_part_") + to_string (i) + string (".exr");
        const string& type = header.type ();
        MultiPartOutputFile outputFile (out_path.c_str (), &header, 1);
        copyPixels (type, inputFile, outputFile, i, 0);
    }
}

template <typename ValueType>
void insertSlice(
    Box2i & dataWindow,
    FrameBuffer & frameBuffer,
    const char name[],
    PixelType pixelType,
    list<Array2D<ValueType>> & channels)
{
    //
    // Allocate a pixel buffer and describe the layout of
    // the buffer in the FrameBuffer object.
    //
    
    size_t width  = dataWindow.max.x - dataWindow.min.x + 1;
    size_t height = dataWindow.max.y - dataWindow.min.y + 1;
    
    channels.emplace_back ();
    Array2D<ValueType> & buffer = channels.back ();
    buffer.resizeErase (height, width);
    
    char * base = (char *)(&buffer[0][0] - dataWindow.min.x - dataWindow.min.y * width);
    
    frameBuffer.insert (
        name,                           // name
        Slice (
            pixelType,                  // type
            base,                       // base
            sizeof (ValueType) * 1,     // xStride
            sizeof (ValueType) * width  // yStride
        )
    );
}

template <typename ValueType>
void insertDeepSlice(
    Box2i & dataWindow,
    DeepFrameBuffer & frameBuffer,
    const char name[],
    PixelType pixelType,
    list<Array2D<ValueType *>> & channels)
{
    //
    // Allocate a pixel buffer and describe the layout of
    // the buffer in the DeepFrameBuffer object.
    //
    
    size_t width  = dataWindow.max.x - dataWindow.min.x + 1;
    size_t height = dataWindow.max.y - dataWindow.min.y + 1;
    
    channels.emplace_back ();
    Array2D<ValueType *> & buffer = channels.back ();
    buffer.resizeErase (height, width);
    
    char * base = (char *)(&buffer[0][0] - dataWindow.min.x - dataWindow.min.y * width);
    
    frameBuffer.insert (
        name,                               // name
        DeepSlice (
            pixelType,                      // type
            base,                           // base
            sizeof (ValueType *) * 1,       // xStride
            sizeof (ValueType *) * width,   // yStride
            sizeof (ValueType) * width      // sample stride
        )
    );
}

FrameBuffer setupFramebuffer (
    const Header & header,
    list<Array2D<uint32_t>> & intChannels,
    list<Array2D<half>>     & halfChannels,
    list<Array2D<float>>    & floatChannels
)
{
    //
    // Allocate pixel buffers for all channels specified in the header, describe the layout of
    // the buffers in the FrameBuffer object.
    //
    
    FrameBuffer frameBuffer;
    Box2i dataWindow = header.dataWindow ();
    
    for (auto i = header.channels ().begin (); i != header.channels ().end (); i++)
    {
        if (i.channel ().type == UINT)
        {
            insertSlice<uint32_t> (dataWindow, frameBuffer, i.name (), i.channel ().type, intChannels);
        }
        else if (i.channel ().type == HALF)
        {
            insertSlice<half> (dataWindow, frameBuffer, i.name (), i.channel ().type, halfChannels);
        }
        else if (i.channel ().type == FLOAT)
        {
            insertSlice<float> (dataWindow, frameBuffer, i.name (), i.channel ().type, floatChannels);
        }
    }

    return frameBuffer;
}

DeepFrameBuffer setupDeepFramebuffer (
    const Header & header,
    Array2D<uint32_t> & sampleCount,
    list<Array2D<uint32_t *>> & intChannels,
    list<Array2D<half *>>     & halfChannels,
    list<Array2D<float *>>    & floatChannels
)
{
    //
    // Allocate pixel buffers for all channels specified in the header, describe the layout of
    // the buffers in the DeepFrameBuffer object.
    //

    DeepFrameBuffer frameBuffer;

    Box2i dataWindow = header.dataWindow ();

    size_t width  = dataWindow.max.x - dataWindow.min.x + 1;
    size_t height = dataWindow.max.y - dataWindow.min.y + 1;
    
    sampleCount.resizeErase (height, width);
    
    frameBuffer.insertSampleCountSlice (
        Slice (
            UINT,
            (char*) (&sampleCount[0][0] - dataWindow.min.x - dataWindow.min.y * width),
            sizeof (unsigned int) * 1,      // xStride
            sizeof (unsigned int) * width   // yStride
        )
    );
    
    for (auto i = header.channels ().begin (); i != header.channels ().end (); i++)
    {
        if (i.channel ().type == UINT)
        {
            insertDeepSlice<uint32_t> (dataWindow, frameBuffer, i.name (), i.channel ().type, intChannels);
        }
        else if (i.channel ().type == HALF)
        {
            insertDeepSlice<half> (dataWindow, frameBuffer, i.name (), i.channel ().type, halfChannels);
        }
        else if (i.channel ().type == FLOAT)
        {
            insertDeepSlice<float> (dataWindow, frameBuffer, i.name (), i.channel ().type, floatChannels);
        }
    }

    return frameBuffer;
}

template <typename T>
void resizeDeepBuffers (
    Array2D<uint32_t> & sampleCount,
    list<Array2D<T *>> & channels
)
{
    //
    // Allocate memory for samples in all pixel buffers according to the data in sampleCount buffer.
    //
    
    for (auto i = channels.begin (); i != channels.end (); i++)
    {
        Array2D<T *> & channel = *i;
        for (int y = 0; y < channel.height (); y++)
        {
            for (int x = 0; x < channel.width (); x++)
            {
                uint32_t count = sampleCount[y][x];
                channel[y][x] = new T [count];
            }
        }
    }
}

template <typename T>
void modifyChannels(list<Array2D<T>> & channels, T delta)
{
    //
    // Dummy code modifying each pixel by incrementing every channel by a given delta.
    //
    
    for (auto i = channels.begin (); i != channels.end (); i++)
    {
        Array2D<T> & channel = *i;
        
        for (int y = 0; y < channel.height (); y++)
        {
            for (int x = 0; x < channel.width (); x++)
            {
                channel[y][x] += delta;
            }
        }
    }
}

template <typename T>
void modifyDeepChannels(Array2D<uint32_t> & sampleCount, list<Array2D<T *>> & channels, T delta)
{
    //
    // Dummy code modifying each deep pixel by incrementing every sample of each channel by a given delta.
    //
    
    for (auto i = channels.begin (); i != channels.end (); i++)
    {
        Array2D<T *> & channel = *i;
        
        for (int y = 0; y < channel.height (); y++)
        {
            for (int x = 0; x < channel.width (); x++)
            {
                uint32_t count = sampleCount[y][x];
                for (uint32_t j = 0; j < count; j++)
                    channel[y][x][j] += delta;
            }
        }
    }
}

void modifyMultipart ()
{
    //
    // Read all channels from a multi-part file, modify each pixel value, write the modified data as a multi-part file.
    // The parts in the file can be scanline- or tile-based, either flat or deep.
    // Every channel of the input file gets modified.
    //
    
    MultiPartInputFile inputFile ("multipart.exr");
    
    std::vector<Header> headers (inputFile.parts ());
    
    for (int i = 0; i < inputFile.parts (); i++)
    {
        headers[i] = inputFile.header (i);
    }
    
    MultiPartOutputFile outputFile("modified.exr", headers.data (), (int)headers.size ());
    
    for (int i = 0; i < inputFile.parts (); i++)
    {
        Header & header = headers[i];
        
        const string& type = header.type ();

        if (type == SCANLINEIMAGE ||
            type == TILEDIMAGE)
        {
            list<Array2D<uint32_t>> intChannels;
            list<Array2D<half>>     halfChannels;
            list<Array2D<float>>    floatChannels;
            
            FrameBuffer frameBuffer = setupFramebuffer (header, intChannels, halfChannels, floatChannels);
            
            if (type == SCANLINEIMAGE)
            {
                InputPart inputPart (inputFile, i);
                inputPart.setFrameBuffer (frameBuffer);
                inputPart.readPixels (header.dataWindow ().min.y, header.dataWindow ().max.y);
            }
            else
            {
                TiledInputPart inputPart (inputFile, i);
                inputPart.setFrameBuffer (frameBuffer);
                inputPart.readTiles (0, inputPart.numXTiles () - 1, 0, inputPart.numYTiles () - 1);
            }
            
            modifyChannels<uint32_t> (intChannels, 1);
            modifyChannels<half>     (halfChannels, 0.3);
            modifyChannels<float>    (floatChannels, 0.5);
            
            if (type == SCANLINEIMAGE)
            {
                Box2i dataWindow = header.dataWindow ();
                OutputPart outputPart (outputFile, i);
                outputPart.setFrameBuffer (frameBuffer);
                outputPart.writePixels (dataWindow.max.y - dataWindow.min.y + 1);
            }
            else
            {
                TiledOutputPart outputPart (outputFile, i);
                outputPart.setFrameBuffer (frameBuffer);
                outputPart.writeTiles (0, outputPart.numXTiles () - 1, 0, outputPart.numYTiles () - 1);
            }
        }
        else if (type == DEEPSCANLINE || type == DEEPTILE)
        {
            Array2D<uint32_t>         sampleCount;
            list<Array2D<uint32_t *>> intChannels;
            list<Array2D<half *>>     halfChannels;
            list<Array2D<float *>>    floatChannels;
            
            DeepFrameBuffer frameBuffer = setupDeepFramebuffer (header, sampleCount, intChannels, halfChannels, floatChannels);
            
            if (type == DEEPSCANLINE)
            {
                DeepScanLineInputPart inputPart (inputFile, i);
                inputPart.setFrameBuffer (frameBuffer);
                inputPart.readPixelSampleCounts (header.dataWindow ().min.y, header.dataWindow ().max.y);
                
                resizeDeepBuffers<uint32_t> (sampleCount, intChannels);
                resizeDeepBuffers<half>     (sampleCount, halfChannels);
                resizeDeepBuffers<float>    (sampleCount, floatChannels);
                
                inputPart.readPixels (header.dataWindow ().min.y, header.dataWindow ().max.y);
            }
            else
            {
                DeepTiledInputPart inputPart (inputFile, i);
                inputPart.setFrameBuffer (frameBuffer);
                inputPart.readPixelSampleCounts (0, inputPart.numXTiles () - 1, 0, inputPart.numYTiles () - 1);
                
                resizeDeepBuffers<uint32_t> (sampleCount, intChannels);
                resizeDeepBuffers<half>     (sampleCount, halfChannels);
                resizeDeepBuffers<float>    (sampleCount, floatChannels);
                
                inputPart.readTiles (0, inputPart.numXTiles () - 1, 0, inputPart.numYTiles () - 1);
            }
            
            modifyDeepChannels<uint32_t> (sampleCount, intChannels, 1);
            modifyDeepChannels<half>     (sampleCount, halfChannels, 0.3);
            modifyDeepChannels<float>    (sampleCount, floatChannels, 0.5);
            
            if (type == DEEPSCANLINE)
            {
                Box2i dataWindow = header.dataWindow ();
                DeepScanLineOutputPart outputPart (outputFile, i);
                outputPart.setFrameBuffer (frameBuffer);
                outputPart.writePixels (dataWindow.max.y - dataWindow.min.y + 1);
            }
            else
            {
                DeepTiledOutputPart outputPart (outputFile, i);
                outputPart.setFrameBuffer (frameBuffer);
                outputPart.writeTiles (0, outputPart.numXTiles () - 1, 0, outputPart.numYTiles () - 1);
            }
        }
    }
}

void multipartExamples()
{
    // Read multiple single-part files and write them out as a single multi-part file.
    combineFiles ();
    
    // Read all parts from a multi-part file, modify each channel of every pixel by incrementing its value, write out as a multi-part file.
    modifyMultipart ();
    
    // Read a multi-part file and write out as multiple single-part files.
    splitFiles ();
}