/*------------------------------------------------------------------------------
 src/plugin/djImageIoExr/save.cxx
 
 Copyright (c) 2004-2006 Darby Johnston.
 All rights reserved.
 Use is subject to license terms.
------------------------------------------------------------------------------*/

//#define DJ_DBG
#include "io.h"
#include <half.h>
#include <ImfChannelList.h>
#include <ImfCompressionAttribute.h>
#include <ImfStandardAttributes.h>
#include <ImfHeader.h>
#include <ImfOutputFile.h>

extern "C" {
const dj::List<dj::ImageSave *> djImageSave()
{
  dj::List<dj::ImageSave *> out;
  out += new djImageIoExr::Save;
  return out;
}
}

namespace djImageIoExr {

Save::Save() :
  _compress(COMPRESS_NONE),
  _bitDepth(F16),
  _f(0)
{}

Save::~Save()
{
  if (_f) delete _f;
}

const StrList Save::attrNameList() const { return saveAttrStr(); }

bool Save::attr(const Str & in, const StrList & data)
{
  StrList tmp = data;
  if (compareCase(saveAttrStr()[ATTR_COMPRESS], in)) {
    if (! inflateStr(tmp, &_compress, compressStr())) return false;
  } else if (compareCase(saveAttrStr()[ATTR_BIT_DEPTH], in)) {
    if (! inflateStr(tmp, &_bitDepth, bitDepthStr())) return false;
  } else
    return false;
  return true;
}

const StrList Save::attr(const Str & in) const
{
  StrList out;
  if (compareCase(saveAttrStr()[ATTR_COMPRESS], in))
    out += deflateStr(_compress, compressStr());
  else if (compareCase(saveAttrStr()[ATTR_BIT_DEPTH], in))
    out += deflateStr(_bitDepth, bitDepthStr());
  return out;
}

bool Save::open(const Str & fileName, const ImageIoInfo & info)
{
  DJ_DBG_FNC("djImageIoExr::Save::open");
  DJ_DBG_PRN("file name = " << fileName);
  
  _fileInfo.set(fileName);
  _info.size(info.size());
  _info.format(info.format());
  _info.type(PIXEL_F32);
  
  return true;
}

bool Save::imageBegin(const Image & in, int frame)
{
  DJ_DBG_FNC("djImageIoExr::Save::imageBegin");
  DJ_DBG_PRN("in = " << in);
  DJ_DBG_PRN("frame = " << frame);
  
  if (! imageOpen(_fileInfo.get(frame == -1 ? _fileInfo.start() : frame), in))
    return false;

  return true;
}

bool Save::imageOpen(const Str & fileName, const Image & in)
{
  DJ_DBG_FNC("djImageIoExr::Save::imageOpen");
  DJ_DBG_PRN("file name = " << fileName);
  
  bool out = false;
try {
  
  // Setup header.
  //
  Imf::Header header(_info.width(), _info.height());
  
  switch (_bitDepth) {
    case BIT_DEPTH_AUTO:
      switch (_info.type()) {
        case PIXEL_U8:
        case PIXEL_U16: _pixelType = Imf::HALF; break;
        case PIXEL_F32: _pixelType = Imf::FLOAT; break;
      }
      break;
    case F16: _pixelType = Imf::HALF; break;
    case F32: _pixelType = Imf::FLOAT; break;
  }
  DJ_DBG_PRN("pixel type = " << _pixelType);
  
  switch (_info.format()) {
    case PIXEL_L:
      header.channels().insert("Y", Imf::Channel(_pixelType));
      break;
    case PIXEL_LA:
      header.channels().insert("Y", Imf::Channel(_pixelType));
      header.channels().insert("A", Imf::Channel(_pixelType));
      break;
    case PIXEL_RGB:
      header.channels().insert("R", Imf::Channel(_pixelType));
      header.channels().insert("G", Imf::Channel(_pixelType));
      header.channels().insert("B", Imf::Channel(_pixelType));
      break;
    case PIXEL_RGBA:
      header.channels().insert("R", Imf::Channel(_pixelType));
      header.channels().insert("G", Imf::Channel(_pixelType));
      header.channels().insert("B", Imf::Channel(_pixelType));
      header.channels().insert("A", Imf::Channel(_pixelType));
      break;
  }
  
  Imf::CompressionAttribute compress;
  switch (_compress) {
    case COMPRESS_NONE: compress = Imf::NO_COMPRESSION; break;
    case COMPRESS_PIZ: compress = Imf::PIZ_COMPRESSION; break;
    case COMPRESS_ZIP: compress = Imf::ZIP_COMPRESSION; break;
    case COMPRESS_RLE: compress = Imf::RLE_COMPRESSION; break;
  }
  header.insert(Imf::CompressionAttribute::staticTypeName(), compress);
  
  // Tags.
  //
  const ImageTagList tagList = in.tagList();
  Str str;
  StrList strList;
  double floatValue;
  StrList data;
  
  data = tagList[imageTagStr()[IMAGE_TAG_CREATOR]];
  if (inflateStr(data, &str)) addOwner(header, std::string(str.toAscii()));
  data = tagList[imageTagStr()[IMAGE_TAG_DESCRIPTION]];
  if (inflateStr(data, &str)) addComments(header, std::string(str.toAscii()));
  data = tagList[imageTagStr()[IMAGE_TAG_TIME_STAMP]];
  if (inflateStr(data, &str)) addCapDate(header, std::string(str.toAscii()));
  data = tagList[imageTagStr()[IMAGE_TAG_UTC_OFFSET]];
  if (inflateStr(data, &floatValue)) addutcOffset(header, floatValue);

  data = tagList[tagStr()[TAG_LONGITUDE]];
  if (inflateStr(data, &floatValue)) addLongitude(header, floatValue);
  data = tagList[tagStr()[TAG_LATITUDE]];
  if (inflateStr(data, &floatValue)) addLatitude(header, floatValue);
  data = tagList[tagStr()[TAG_ALTITUDE]];
  if (inflateStr(data, &floatValue)) addAltitude(header, floatValue);

  data = tagList[tagStr()[TAG_FOCUS]];
  if (inflateStr(data, &floatValue)) addFocus(header, floatValue);
  data = tagList[tagStr()[TAG_EXPOSURE]];
  if (inflateStr(data, &floatValue)) addExpTime(header, floatValue);
  data = tagList[tagStr()[TAG_APERTURE]];
  if (inflateStr(data, &floatValue)) addAperture(header, floatValue);
  data = tagList[tagStr()[TAG_ISO_SPEED]];
  if (inflateStr(data, &floatValue)) addIsoSpeed(header, floatValue);

  data = tagList[tagStr()[TAG_CHROMATICITIES]];
  if (inflateStr(data, &strList))
    if (strList.size() == 8)
      addChromaticities(
        header,
        Imf::Chromaticities(
          Imath::V2f(toFloat(strList[0]), toFloat(strList[1])),
          Imath::V2f(toFloat(strList[2]), toFloat(strList[3])),
          Imath::V2f(toFloat(strList[4]), toFloat(strList[5])),
          Imath::V2f(toFloat(strList[6]), toFloat(strList[7]))
        )
      );
  data = tagList[tagStr()[TAG_WHITE_LUMINANCE]];
  if (inflateStr(data, &floatValue)) addWhiteLuminance(header, floatValue);
  data = tagList[tagStr()[TAG_X_DENSITY]];
  if (inflateStr(data, &floatValue)) addXDensity(header, floatValue);
  
  data = tagList[imageTagStr()[IMAGE_TAG_KEY_CODE]];
  if (inflateStr(data, &str)) {
    int id = 0, type = 0, prefix = 0, count = 0, offset = 0;
    strToKeyCode(str, &id, &type, &prefix, &count, &offset);
    addKeyCode(header, Imf::KeyCode(id, type, prefix, count, offset));
  }
  data = tagList[imageTagStr()[IMAGE_TAG_TIME_CODE]];
  if (inflateStr(data, &str))
    addTimeCode(header, Imf::TimeCode(strToTimeCode(str)));
  
  // Open file.
  //
  if (_f) delete _f;
  _f = new Imf::OutputFile(fileName.toAscii(), header);

  out = true;
} catch (const std::exception &) {}
  return out;
}

bool Save::imageSet(const PixelData & in)
{
  DJ_DBG_FNC("djImageIoExr::Save::imageSet");
  DJ_DBG_PRN("in = " << in);

  bool r = false;
try {

  const int channels = pixelSize(_info.format());
  int bytes = 0;
  if (Imf::HALF == _pixelType) bytes = 2;
  else if (Imf::FLOAT == _pixelType) bytes = 4;
  const int strideX = channels * bytes;
  const int strideY = _info.width() * strideX;
  Str channelListStr;
  switch (_info.format()) {
    case PIXEL_L: channelListStr = Str("Y"); break;
    case PIXEL_LA: channelListStr = Str("YA"); break;
    case PIXEL_RGB: channelListStr = Str("RGB"); break;
    case PIXEL_RGBA: channelListStr = Str("RGBA"); break;
  }

  // Convert.
  //
  PixelData framebuf(_info);
  framebuf.zero();
  pixelCopy(in, &framebuf);
  DJ_DBG_PRN("framebuf = " << framebuf);
  
  // Write.
  //
  int y = 0;
  for (; y < _info.height(); ++y) {
    //DJ_DBG_PRN("y = " << y);
    uint8_t * const imageP =
      reinterpret_cast<uint8_t *>(framebuf.data(0, _info.height() - 1 - y));
    
    // Convert.
    //
    if (Imf::HALF == _pixelType) {
      const PixelF32 * inP = reinterpret_cast<PixelF32 *>(imageP);
      const PixelF32 * const end = inP + _info.width() * channels;
      half * outP = reinterpret_cast<half *>(imageP);
      while (inP < end) *outP++ = *inP++;
    }
    
    Imf::FrameBuffer frameBuffer;
    for (StrIt it = channelListStr; it.isValid(); ++it)
      frameBuffer.insert(
        Str(it()).toAscii(),
        Imf::Slice(
          _pixelType,
          reinterpret_cast<char *>(imageP) + -y * strideY + bytes * it.count(),
          strideX, strideY
        )
      );
    _f->setFrameBuffer(frameBuffer);
    _f->writePixels(1);
  }
  
  r = y == _info.height();
} catch (const std::exception & in) {}
  delete _f;
  _f = 0;
  return r;
}

}
