// exrWriter.C
// Copyright (c) 2009 The Foundry Visionmongers Ltd.  All Rights Reserved.

/* Reads exr files using libexr.
   This is an example of a file reader that is not a subclass of
   FileWriter. Instead this uses the library's reader functions and
   a single lock so that multiple threads do not crash the library.

   04/14/03     Initial Release                Charles Henrich
   12/04/03     User selectable compression,   Charles Henrich
        float precision, and autocrop
   10/04    Defaulted autocrop to off    spitzak
   5/06        black-outside and reformatting    spitzak
 */
#include "DDImage/DDWindows.h"
#include "DDImage/Writer.h"
#include "DDImage/Row.h"
#include "DDImage/Knobs.h"
#include "DDImage/Tile.h"
#include "DDImage/DDString.h"
#include "DDImage/MetaData.h"
#include "DDImage/LUT.h"
#include <errno.h>
#include <stdio.h>

#include <OpenEXR/ImfOutputFile.h>
#include <OpenEXR/ImfChannelList.h>
#include <OpenEXR/ImfArray.h>
#include <OpenEXR/ImfCompression.h>
#include <OpenEXR/ImfStringAttribute.h>
#include <OpenEXR/ImfIntAttribute.h>
#include <OpenEXR/ImfBoxAttribute.h>
#include <OpenEXR/ImfDoubleAttribute.h>
#include <OpenEXR/ImfStringVectorAttribute.h>
#include <OpenEXR/ImfTimeCodeAttribute.h>
#include <OpenEXR/ImfStandardAttributes.h>
#include <OpenEXR/ImfFramesPerSecond.h>
#include <OpenEXR/ImfMatrixAttribute.h>
#include <OpenEXR/ImathBox.h>
#include <OpenEXR/half.h>

#include "exrGeneral.h"

using namespace DD::Image;

class exrWriter : public Writer
{
  void autocrop_tile(Tile& img, ChannelMask channels, int* bx, int* by, int* br, int* bt);
  int datatype;
  int compression;
  bool autocrop;
  bool writeHash;
  int hero;
  int _metadataMode;
  bool _doNotWriteNukePrefix;
  bool _followStandard;

public:

  exrWriter(Write* iop);
  ~exrWriter();

  Iop* firstInput(const std::set<int>& wantViews);

  void execute();
  void knobs(Knob_Callback f);
  int knob_changed(Knob *k);
  static const Writer::Description d;

  // Make it default to linear colorspace:
  LUT* defaultLUT() const { return LUT::getLut(LUT::FLOAT); }

  int split_input(int i) const
  {
    return int(executingViews().size() ? executingViews().size() : 1);
  }

  //! This writer is capable of writing out the overscan, so passthrough should not
  //! clip to the format.
  virtual bool clipToFormat() const
  {
    return false;
  }

  /**
   * return the view which we are expecting on input N
   */
  int view(int n) const
  {
    std::set<int> views = executingViews();

    std::set<int>::const_iterator i = views.begin();
    while (i != views.end()) {
      if (!n) {
        return *i;
      }
      n--;
      i++;
    }

    return 0;
  }

  const OutputContext& inputContext(int i, OutputContext& o) const
  {
    o = iop->outputContext();
    o.view(view(i));
    return o;
  }

  const char* help() { return "OpenEXR high dynamic range format from ILM"; }

};

static Writer* build(Write* iop)
{
  return new exrWriter(iop);
}

const Writer::Description exrWriter::d("exr\0sxr\0", build);

exrWriter::exrWriter(Write* iop) : Writer(iop)
{
  setFlags(DONT_CHECK_INPUT0_CHANNELS);
  datatype = 0;
  compression = 1;
  autocrop = false;
  writeHash = true;
  hero = 1;
  _metadataMode = eDefaultMetaData;
  _doNotWriteNukePrefix = false;
  _followStandard = 0;
}

exrWriter::~exrWriter()
{
}

class RowGroup
{
  std::vector<Row*> row;

public:
  RowGroup(size_t n, int x, int r)
  {
    row.resize(n);
    for (size_t i = 0; i < n; i++) {
      row[i] = new Row(x, r);
    }
  }

  ~RowGroup()
  {
    for (size_t i = 0; i < row.size(); i++) {
      delete row[i];
    }
  }

  Row& operator[](int i)
  {
    if (i >= int(row.size()))
      abort();
    return *row[i];
  }

  const Row& operator[](int i) const
  {
    if (i >= int(row.size()))
      abort();
    return *row[i];
  }
};

Iop* exrWriter::firstInput(const std::set<int>& wantViews)
{
  for (int i = 0; i < iop->inputs(); i++) {
    if (wantViews.find(view(i)) == wantViews.end())
      continue;
    
    return iop->input(i);
  }
  return &input0();
}


void exrWriter::execute()
{
  const std::set<int> execViews = executingViews();
  std::set<int> wantViews = iop->executable()->viewsToExecute();

  if (wantViews.size() == 0) {
    wantViews = execViews;
  }

  int floatdepth = datatype ? 32 : 16;

  Imf::Compression compression = ctypes[this->compression];

  ChannelSet channels(firstInput(wantViews)->channels());
  channels &= (iop->channels());
  if (!channels) {
    iop->critical("exrWriter: No channels selected (or available) for write\n");
    return;
  }

  if (premult() && !lut()->linear() &&
      (channels & Mask_RGB) && (input0().channels() & Mask_Alpha))
    channels += (Mask_Alpha);

  std::vector<int> views;
  std::vector<std::string> viewstr;

  viewstr.push_back(OutputContext::viewname(hero));

  if (wantViews.size() == 1) {
    hero = *wantViews.begin();
  }

  for (std::set<int>::const_iterator i = execViews.begin(); i != execViews.end(); i++) {
    views.push_back(*i);
    if (*i != hero) {
      viewstr.push_back(OutputContext::viewname(*i));
    }
  }

  DD::Image::Box bound;

  bool sizewarn = false;

  bool firstInputBbox = true;

  for (int i = 0; i < iop->inputs(); i++) {

    if (wantViews.find(view(i)) == wantViews.end())
      continue;

    Iop* input = iop->input(i);
    input->validate(true);

  }

  for (int i = 0; i < iop->inputs(); i++) {

    if (wantViews.find(view(i)) == wantViews.end())
      continue;

    Iop* input = iop->input(i);
    int bx = input->x();
    int by = input->y();
    int br = input->r();
    int bt = input->t();
    if (input->black_outside()) {
      if (bx + 2 < br) {
        bx++;
        br--;
      }
      if (by + 2 < bt) {
        by++;
        bt--;
      }
    }

    input->request(bx, by, br, bt, channels, 1);

    if (br - bx > input0().format().width() * 1.5 ||
        bt - by > input0().format().height() * 1.5) {
      // print this warning before it possibly crashed due to requesting a
      // huge buffer!
      if (sizewarn) {
        fprintf(stderr, "!WARNING! Bounding Box Area is > 1.5 times larger "
                        "than format. You may want crop your image before writing it.\n");
        sizewarn = true;
      }
    }

    if (autocrop) {

      Tile img(*input, input->x(), input->y(),
               input->r(), input->t(), channels, true);

      if (iop->aborted()) {
        //iop->critical("exrWriter: Write failed [Unable to get input tile]\n");
        return;
      }

      autocrop_tile(img, channels, &bx, &by, &br, &bt);
      bt++; /* We (aka nuke) want r & t to be beyond the last pixel */
      br++;
    }

    if (firstInputBbox) {
      bound.y(by);
      bound.x(bx);
      bound.r(br);
      bound.t(bt);
    }
    else {
      bound.y(std::min(bound.y(), by));
      bound.x(std::min(bound.x(), bx));
      bound.r(std::max(bound.r(), br));
      bound.t(std::max(bound.t(), bt));
    }

    firstInputBbox = false;
  }

  const Format& inputFormat = firstInput(wantViews)->format();

  Imath::Box2i C_datawin;
  C_datawin.min.x = bound.x();
  C_datawin.min.y = inputFormat.height() - bound.t();
  C_datawin.max.x = bound.r() - 1;
  C_datawin.max.y = inputFormat.height() - bound.y() - 1;

  Imath::Box2i C_dispwin;
  C_dispwin.min.x = 0;
  C_dispwin.min.y = 0;
  C_dispwin.max.x = inputFormat.width() - 1;
  C_dispwin.max.y = inputFormat.height() - 1;

  try {
    int numchannels = channels.size();
    Imf::OutputFile* outfile;
    Imf::FrameBuffer fbuf;

    RowGroup renderrow(executingViews().size(), bound.x(), bound.r());
    RowGroup writerow(executingViews().size(), bound.x(), bound.r());

    Imf::Header exrheader(C_dispwin, C_datawin, iop->format().pixel_aspect(),
                          Imath::V2f(0, 0), 1, Imf::INCREASING_Y, compression);

    if (wantViews.size() > 1) {
      // only write multi view string if a stereo file
      Imf::StringVectorAttribute multiViewAttr;

      multiViewAttr.value() = viewstr;
      exrheader.insert("multiView", multiViewAttr);
    }

    Iop* metaInput = NULL;
    for (size_t viewIdx = 0; viewIdx < views.size(); viewIdx ++) {
      if (wantViews.find(views[viewIdx]) == wantViews.end()) {
        continue;
      }
      if (metaInput == NULL || views[viewIdx] == hero) {
        metaInput = iop->input(viewIdx);
      }
    }
    if (metaInput == NULL) {
      metaInput = iop->input(0);
    }

    const MetaData::Bundle& metadata = metaInput->fetchMetaData(NULL);
    Hash nodeHash = iop->getHashOfInputs();
    metadataToExrHeader( (enum ExrMetaDataMode)  _metadataMode, metadata, exrheader, iop, writeHash ? &nodeHash : 0, _doNotWriteNukePrefix );

    Imf::Array2D<half> halfwriterow(numchannels * views.size(), bound.r() - bound.x());

    std::map<int, ChannelSet> channelsperview;

    for (int v = 0; v < int(views.size()); v++) {

      if (wantViews.find(views[v]) == wantViews.end()) {
        continue;
      }

      int curchan = 0;
      foreach(z, channels) {
        std::string channame;

        switch (z) {
          case Chan_Red: channame = "R";
            break;
          case Chan_Green: channame = "G";
            break;
          case Chan_Blue: channame = "B";
            break;
          case Chan_Alpha: channame = "A";
            break;
          default: channame = iop->channel_name(z);
            break;
        }

        if (executingViews().size() > 1 && views[v] != hero) {
          if (_followStandard){
            size_t i = channame.find('.');
            if (i != channame.npos){
              std::string layerName = channame.substr(0, i);
              std::string channelName = channame.substr(i + 1);
              channame = layerName + "." + OutputContext::viewname(views[v]) + "." + channelName;
            }
            else
              channame = OutputContext::viewname(views[v]) + "." + channame;
          }
          else
            channame = OutputContext::viewname(views[v]) + "." + channame;

          if (z == Chan_Stereo_Disp_Left_X ||
              z == Chan_Stereo_Disp_Left_Y ||
              z == Chan_Stereo_Disp_Right_X ||
              z == Chan_Stereo_Disp_Right_Y) {
            continue;
          }
        }

        channelsperview[v].insert(z);

        if (floatdepth == 32) {
          exrheader.channels().insert(channame.c_str(), Imf::Channel(Imf::FLOAT));
        }
        else {
          exrheader.channels().insert(channame.c_str(), Imf::Channel(Imf::HALF));
        }

        writerow[v].writable(z);

        if (floatdepth == 32) {
          fbuf.insert(channame.c_str(),
                      Imf::Slice(Imf::FLOAT, (char*)(float*)writerow[v][z],
                                 sizeof(float), 0));
        }
        else {
          fbuf.insert(channame.c_str(),
                      Imf::Slice(Imf::HALF,
                                 (char*)(&halfwriterow[v * numchannels + curchan][0] - C_datawin.min.x),
                                 sizeof(halfwriterow[v * numchannels * curchan][0]), 0));
          curchan++;
        }
      }
    }

    std::string temp_name;
    temp_name = createFileHash().c_str();
    temp_name += ".tmp";
    outfile = new Imf::OutputFile(temp_name.c_str(), exrheader);
    outfile->setFrameBuffer(fbuf);

    for (int scanline = bound.t() - 1;
         scanline >= bound.y();
         scanline--) {

      for (int v = 0; v < int(views.size()); v++) {

        channels = channelsperview[v];

        if (wantViews.find(views[v]) == wantViews.end()) {
          continue;
        }

        writerow[v].pre_copy(renderrow[v], channels);

        iop->inputnget(v, scanline, bound.x(), bound.r(), channels, renderrow[v], 1.0/wantViews.size());
        if (iop->aborted())
          break;

	if (bound.is_constant()) {
	  foreach(z, channels)
	    renderrow[v].erase(z);
	  continue;
	}

        int curchan = 0;

        foreach(z, channels) {

          const float* from = renderrow[v][z];
          const float* alpha = renderrow[v][Chan_Alpha];
          float* to = writerow[v].writable(z);

          if (!lut()->linear() && z <= Chan_Blue) {
            to_float(z - 1, to + C_datawin.min.x,
                     from + C_datawin.min.x,
                     alpha + C_datawin.min.x,
                     C_datawin.max.x - C_datawin.min.x + 1);
            from = to;
          }

          if (bound.r() > iop->input(v)->r()) {
            float* end = renderrow[v].writable(z)   + bound.r();
            float* start = renderrow[v].writable(z) + iop->input(v)->r();
            while (start < end) {
              *start = 0;
              start++;
            }
          }

          if (bound.x() < iop->input(v)->x()) {
            float* end = renderrow[v].writable(z)   + bound.x();
            float* start = renderrow[v].writable(z) + iop->input(v)->x();
            while (start > end) {
              *start = 0;
              start--;
            }
          }

          if (floatdepth == 32) {
            if (to != from)
              for (int count = C_datawin.min.x; count < C_datawin.max.x + 1; count++)
                to[count] = from[count];
          }
          else {
            for (int count = C_datawin.min.x; count < C_datawin.max.x + 1; count++)
              halfwriterow[v * numchannels + curchan][count - C_datawin.min.x] = from[count];
            curchan++;
          }
        }

        progressFraction(double(bound.t() - scanline) / (bound.t() - bound.y()));
      }
      outfile->writePixels(1);
    }

    delete outfile;

    if( !FileIop::renameFile( temp_name.c_str(), filename() ) ) 
      iop->critical("Can't rename .tmp to final, %s", strerror(errno));
  }
  catch (const std::exception& exc) {
    iop->critical("EXR: Write failed [%s]\n", exc.what());
    return;
  }
}

void exrWriter::knobs(Knob_Callback f)
{
  Bool_knob(f, &autocrop, "autocrop");
  Tooltip(f, "Reduce the bounding box to the non-zero area. This is normally "
             "not needed as the zeros will compress very small, and it is slow "
             "as the whole image must be calculated before any can be written. "
             "However this may speed up some programs reading the files.");

  Bool_knob(f, &writeHash, "write_hash", "write hash");
  SetFlags(f, Knob::INVISIBLE);
  Tooltip(f, "Write the hash of the node graph into the exr file.  Useful to see if your image is up to date when doing a precomp.");

  Enumeration_knob(f, &datatype, dnames, "datatype");
  Enumeration_knob(f, &compression, cnames, "compression");

  Obsolete_knob(f, "stereo", 0);

  OneView_knob(f, &hero, "heroview");
  Tooltip(f, "If stereo is on, this is the view that is written as the \"main\" image");

  Enumeration_knob(f, &_metadataMode, metadata_modes, "metadata");
  Tooltip(f, "Which metadata to write out to the EXR file."
             "<p>'no metadata' means that no custom attributes will be created and only metadata that fills required header fields will be written.<p>'default metadata' means that the optional timecode, edgecode, frame rate and exposure header fields will also be filled using metadata values.");
  Bool_knob(f, &_doNotWriteNukePrefix, "noprefix", "do not attach prefix" );
  Tooltip(f, "By default unknown metadata keys have the prefix 'nuke' attached to them before writing them into the file.  Enable this option to write the metadata 'as is' without the nuke prefix.");

  Newline(f);
  Bool_knob(f, &_followStandard, "Standard layer name format");
  Tooltip(f, "Older versions of Nuke write out channel names in the format: view.layer.channel."
              "Check this option to follow the EXR standard format: layer.view.channel");
}

int exrWriter::knob_changed(Knob *k)
{
  if ( k == &Knob::showPanel || k->is("metadata") ) {
    iop->knob( "noprefix")->enable( ExrMetaDataMode(_metadataMode) >=  eAllMetadataExceptInput );
    return 1;
  }
  return 0;
}

void exrWriter::autocrop_tile(Tile& img, ChannelMask channels,
                              int* bx, int* by, int* br, int* bt)
{
  int xcount, ycount;

  *bx = img.r();
  *by = img.t();
  *br = img.x();
  *bt = img.y();

  foreach (z, channels) {
    for (ycount = img.y(); ycount < img.t(); ycount++) {
      for (xcount = img.x(); xcount < img.r(); xcount++) {
        if (img[z][ycount][xcount] != 0) {
          if (xcount < *bx)
            *bx = xcount;
          if (ycount < *by)
            *by = ycount;
          break;
        }
      }
    }

    for (ycount = img.t() - 1; ycount >= img.y(); ycount--) {
      for (xcount = img.r() - 1; xcount >= img.x(); xcount--) {
        if (img[z][ycount][xcount] != 0) {
          if (xcount > *br)
            *br = xcount;
          if (ycount > *bt)
            *bt = ycount;
          break;
        }
      }
    }
  }

  if (*bx > *br || *by > *bt)
    *bx = *by = *br = *bt = 0;
}