// exrReaderDeep.cpp

///////////////////////////////////////////////////////////////////////////
//
// Copyright (c) 2011 The Foundry Visionmongers Ltd.  All Rights Reserved.
// Portions contributed and copyright held by others as indicated.  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 The Foundry Visionmongers nor any other contributors 
//      to this software 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 <OpenEXR/ImfMultiPartInputFile.h>
#include <OpenEXR/ImfDeepScanLineInputPart.h>
#include <OpenEXR/ImfDeepFrameBuffer.h>

#include <OpenEXR/ImfPartType.h>
#include <OpenEXR/ImfChannelList.h>

#include <OpenEXR/ImfStringAttribute.h>
#include <OpenEXR/ImfIntAttribute.h>
#include <OpenEXR/ImfFloatAttribute.h>
#include <OpenEXR/ImfVecAttribute.h>
#include <OpenEXR/ImfBoxAttribute.h>
#include <OpenEXR/ImfStringVectorAttribute.h>
#include <OpenEXR/ImfTimeCodeAttribute.h>
#include <OpenEXR/ImfStandardAttributes.h>
#include <OpenEXR/ImfMatrixAttribute.h>

#include <OpenEXR/ImfFramesPerSecond.h>

#include <stdio.h>

#include "DDImage/DeepReader.h"
#include "DDImage/DeepOp.h"
#include "DDImage/Thread.h"

#include "exrGeneral.h"

using namespace DD::Image;

class exrDeepReaderFormat : public DeepReaderFormat
{

  bool _doNotAttachPrefix;

public:

  exrDeepReaderFormat()
  {
    _doNotAttachPrefix = false;
  }

  bool doNotAttachPrefix() const { return _doNotAttachPrefix; }

  void knobs(Knob_Callback f)
  {
    Bool_knob(f, &_doNotAttachPrefix, "noprefix",  "do not attach prefix");
    Tooltip(f, "By default the 'exr' prefix is attached to metadata keys to make it distinct from other metadata in the tree.  Enable this option to read the metadata 'as is'' without attaching the exr prefix.");
  }

  void append(Hash& hash)
  {
    hash.append(_doNotAttachPrefix);
  }
};

/**
 * OpenEXR2 reader for deep data.
 *
 * Missing features, to be adapted from the existing EXR reader:
 *  stereo
 */
class exrReaderDeep : public DeepReader
{
  std::string _filename;
  ChannelSet _decodeChannels;
  ChannelMap _decodeChannelMap;
  Imf::MultiPartInputFile *_file;
  Imf::DeepScanLineInputPart *_part;
  int _partNumber;
  std::map<Channel, std::string> _chans;
  std::map<Channel, Imf::PixelType> _chanTypes;
  MetaData::Bundle _meta;

  Lock _lock;

  void createChannelMap(const Imf::Header& header);

  const char* chanName(Channel chan)
  {
    return _chans[chan].c_str();
  }

  bool decodeLine(Imf::DeepScanLineInputPart& part, int exrY, const Box& box, const ChannelSet& reqChannels, DeepOutputPlane& plane);
  void setMetaData( const Imf::Header& header, DD::Image::MetaData::Bundle& meta, bool doNotAttachPrefix );
  
public:

  exrReaderDeep(DeepReaderOwner* iop, const std::string& filename);
  virtual ~exrReaderDeep();
  
  virtual bool doDeepEngine(Box box, const ChannelSet& reqChannels, DeepOutputPlane& plane);
  const MetaData::Bundle& fetchMetaData(const char* key);
};

exrReaderDeep::exrReaderDeep(DeepReaderOwner* iop, const std::string& filename) : DeepReader(iop), _filename(filename)
{
  _file = 0;
  _part = 0;
  exrDeepReaderFormat *exrOptions = dynamic_cast<exrDeepReaderFormat*>( iop->handler() );
  bool doNotAttachPrefix = exrOptions ? exrOptions->doNotAttachPrefix() : false;

  try {
    _file = new Imf::MultiPartInputFile( filename.c_str() );

    std::string view( iop->readerOutputContext().viewname() );

    int partNumber = -1;

    for(int i=0; i < _file->parts(); i++) {
      const Imf::Header& header = _file->header(i);
      if ( i == 0 )
        setMetaData( header, _metaData, doNotAttachPrefix ); // read 'global' metadata from first part

      // look for a part with a type set to DEEPSCANLINE
      if ( ! header.hasType() || header.type() != Imf::DEEPSCANLINE )
        continue;

      if( header.hasView() ) {
        if( view == header.view() ) {
          partNumber = i;
          break;
        } else {
          // wrong view, but it'll do if all else fails
          partNumber = i;
        }
      } else {
        // no view mentioned - if we wanted main, that's exactly right
        if( view == "main" ) {
          partNumber = i;
          break;
        } else {
          // no view in file - if we aren't using an earlier part, use this one
          if ( partNumber == -1 ) 
            partNumber = i;
        }
      }
    }

    if( partNumber == -1) {
      throw Iex::InputExc("no deep data found in file");
    }
    
    _partNumber = partNumber;
    _part = new Imf::DeepScanLineInputPart( *_file, _partNumber);

    const Imf::Header& header = _part->header();

    createChannelMap(header);

    // read part metadata - overrides global metadata

    setMetaData( header, _metaData, doNotAttachPrefix );

    Box displayBox = boxToBox(header.displayWindow(), header.displayWindow());
    Box dataBox = boxToBox(header.dataWindow(), header.displayWindow());
    
    setInfo(displayBox.r(),
            displayBox.h(),
            OutputContext(),
            _decodeChannels );

    _deepInfo.box().set(dataBox);
  }
  
  catch ( Iex::BaseExc& e ) {
    _owner->readerInternalError( e.what() );
  }
}

exrReaderDeep::~exrReaderDeep()
{
  delete _file;
  delete _part;
}

/**
 * decode the ofx channel names to Nuke channel numbers and 
 * names
 */
void exrReaderDeep::createChannelMap(const Imf::Header& header)
{
  for (Imf::ChannelList::ConstIterator it = header.channels().begin();
       it != header.channels().end();
       it++) {
    
    const char* chanName = it.name();
    Channel channel = getExrChannel(chanName);
    _chans[channel] = chanName;
    _decodeChannels += channel;
  }
  
  _decodeChannelMap = ChannelMap(_decodeChannels);
}

bool exrReaderDeep::decodeLine(Imf::DeepScanLineInputPart& part, int exrY, const Box& box, const ChannelSet& reqChannels, DeepOutputPlane& plane)
{
  const Imf::Header& header = part.header();
  
  const int dataWid = header.dataWindow().size().x + 1;
  const int dataX = header.dataWindow().min.x;
  
  if (exrY < header.dataWindow().min.y || exrY > header.dataWindow().max.y)
    return false;
  
  std::vector<unsigned> sampleCounts(dataWid, 0);

  /// set the slice for the sample counts
  Imf::DeepFrameBuffer frameBuffer;
  frameBuffer.insertSampleCountSlice(Imf::Slice(Imf::UINT, 
                                                (char*)(&sampleCounts[0] - dataX),
                                                sizeof(unsigned), 0));

  unsigned chanCount = _decodeChannelMap.size();
  
  std::vector<std::vector<const float*> > samples;
  samples.reserve(chanCount);

  // set up the slices for the actual data.  Note these pointers
  // don't actually point anywhere yet, they can't be filled in
  // until we have done readPixelSampleCounts (but they need to be
  // present for setFrameBuffer)
  //
  // Due to an apparent bug in the EXR library, we need to decode
  // ALL the channels in the file, even if we aren't interested
  // in them.
  foreach(z, _decodeChannels) {
    samples.push_back(std::vector<const float*>(dataWid));
    frameBuffer.insert(chanName(z),
                       Imf::DeepSlice(Imf::FLOAT,
                                      (char *)(&samples.back()[0] - dataX),
                                      sizeof(const float*),        // xstride
                                      0,                           // ystride
                                      sizeof(float) * chanCount)); // samplestride
  }

  std::vector<float> data;

  {
    Guard g(_lock);
    part.setFrameBuffer(frameBuffer);
    part.readPixelSampleCounts(exrY, exrY);

    // get the total number of samples

    unsigned sampleCountSum = 0;
    for (unsigned i = 0; i < sampleCounts.size(); i++)
      sampleCountSum += sampleCounts[i];

    // if there were no samples, skip out

    if (sampleCountSum == 0) {
      return false;
    }

    // allocate the actual data for the samples, and then go back and
    // set the samples vectors that the framebuffer is pointing at
    // to point to spans within this.

    data.resize(sampleCountSum * chanCount, 0);

    foreach(z, _decodeChannels) {
      const int chanNo = _decodeChannelMap.chanNo(z);
      const float* ptr = &data[0] + chanNo;

      for (int x = 0; x < dataWid; x++) {
        samples[chanNo][x] = ptr;
        ptr += chanCount * sampleCounts[x];
      }
    }

    // read the actual data
    part.readPixels(exrY);
  }
  
  for (int x = box.x(); x < box.r(); x++) {
    
    const unsigned lineX = x - dataX;

    if (lineX >= sampleCounts.size() || sampleCounts[lineX] == 0)
      plane.addHole(); // if we're out of range, add a hole (no data)
    else {
      DeepOutPixel outPixel;

      const unsigned sampleCount = sampleCounts[lineX];

      // copy data from the unpacked buffer into the output plane.
      // in some future version it might be nice to unpack directly
      // into the output plane's buffer.  Possibly need some API changes
      // for this.

      for (unsigned sampleNo = 0; sampleNo < sampleCount; sampleNo++)
        foreach(z, reqChannels) {
          const int chanNo = _decodeChannelMap.chanNo(z);
          if (_decodeChannelMap.contains(z)) {
            outPixel.push_back(samples[chanNo][lineX][sampleNo * chanCount]);
          }
          else
            outPixel.push_back(0);
        }
      
      plane.addPixel(outPixel);
    }
  }
  return true;
}

void exrReaderDeep::setMetaData( const Imf::Header& header, DD::Image::MetaData::Bundle& meta, bool doNotAttachPrefix)
{
  std::map<Imf::PixelType, int> pixelTypes;

  for (Imf::ChannelList::ConstIterator it = header.channels().begin();
       it != header.channels().end();
       it++) {

    Channel channel = getExrChannel( it.name() );
    if ( channel != Chan_DeepBack && channel != Chan_DeepFront )
      pixelTypes[ it.channel().type ]++;
  }

  if (pixelTypes[Imf::FLOAT] > 0) {
    meta.setData(MetaData::DEPTH, MetaData::DEPTH_FLOAT);
  }
  else if (pixelTypes[Imf::HALF] > 0) {
    meta.setData(MetaData::DEPTH, MetaData::DEPTH_HALF);
  }

  exrHeaderToMetadata( header, meta, doNotAttachPrefix );
}

bool exrReaderDeep::doDeepEngine(Box box, const ChannelSet& reqChannels, DeepOutputPlane& plane)
{
  const Imf::Header& header = _part->header();
  plane = DeepOutputPlane(reqChannels, box);
  
  try {
    for (int y = box.y(); y < box.t(); y++) {
      int exrY = lineToLine(y, header.displayWindow());
      bool decoded = decodeLine(*_part, exrY, box, reqChannels, plane);

      if (!decoded)
        for (int x = box.x(); x < box.r(); x++)
          plane.addHole();
    }
  }
  catch ( Iex::BaseExc& e ) {
    _op->error( e.what() );
    return false;
  }
  return true;
}

static DeepReader* build(DeepReaderOwner* iop, const std::string& fn)
{
  return new exrReaderDeep(iop, fn);
}

static DeepReaderFormat* buildFormat(DeepReaderOwner* op)
{
  return new exrDeepReaderFormat();
}

static const DeepReader::Description d("exr\0", build, buildFormat);