///////////////////////////////////////////////////////////////////////////
//
// Copyright (c) 2012, Weta Digital Ltd
// 
// All rights reserved.
// 
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
// *       Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
// *       Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
// *       Neither the name of Weta Digital nor the names of
// its contributors may be used to endorse or promote products derived
// from this software without specific prior written permission. 
// 
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
///////////////////////////////////////////////////////////////////////////


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

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

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



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)
{
    int 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++)
    {
        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(_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],
                                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


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

               channel_number++;

           }

           pixel++;

       }// next pixel on row
}

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


}

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
   
   size_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++)
       {
          total_sizes[ptr]+=counts[j][ptr];
          if(counts[j][ptr]>0) num_sources[ptr]++;
       }
       overall_sample_count+=total_sizes[ptr];
       
       
       
   }
   
  
  
   
   //
   // 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
          //
          
          size_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