#include "OCIOCDLTransform.h"

#include "DDImage/NukeWrapper.h"
#include "DDImage/Op.h"
#include "DDImage/Knob.h"
#include "DDImage/Knobs.h"
#include "DDImage/Enumeration_KnobI.h"

#include <string>
#include <fstream>
#include <stdexcept>

const DD::Image::Op::Description OCIOCDLTransform::description("OCIOCDLTransform", build);


DD::Image::Op* build(Node *node)
{
  DD::Image::NukeWrapper *op = new DD::Image::NukeWrapper(new OCIOCDLTransform(node));
  op->channels(DD::Image::Mask_RGB);
  return op;
}

OCIOCDLTransform::OCIOCDLTransform(Node* node)
  : DD::Image::OCIOWorkingSpaceBase(node)
{
}

const char* OCIOCDLTransform::Class() const
{
  return description.name;
}

const char* OCIOCDLTransform::displayName() const
{
  return description.name;
}

const char* OCIOCDLTransform::node_help() const
{
  return "Use OpenColorIO to apply an ASC CDL grade. Applied using:\n\n"
    "out = (i * s + o)^p\n\n"
    "Where i is the input value, s is slope, o is offset and p is power";
}

void OCIOCDLTransform::knobs(DD::Image::Knob_Callback f)
{
  m_slopeKnob = DD::Image::Color_knob(f, m_slope, DD::Image::IRange(0, 4.0), "slope");
  m_offsetKnob = DD::Image::Color_knob(f, m_offset, DD::Image::IRange(-0.2, 0.2), "offset");
  m_powerKnob = DD::Image::Color_knob(f, m_power, DD::Image::IRange(0.0, 4.0), "power");
  m_saturationKnob = DD::Image::Float_knob(f, &m_saturation, DD::Image::IRange(0, 4.0), "saturation");

  Obsolete_knob(f, "direction", "if {$value==\"inverse\"} { knob invert true }");
  DD::Image::Newline(f);
  m_invertKnob = DD::Image::Bool_knob(f, &m_invert, "invert", "invert direction");
  DD::Image::Tooltip(f, "Specify the transform direction.");

  OCIOWorkingSpaceBase::knobs(f);

  DD::Image::Divider(f);

  DD::Image::Bool_knob(f, &m_readFromFile, "read_from_file", "read from file");
  DD::Image::SetFlags(f, DD::Image::Knob::EARLY_STORE | DD::Image::Knob::KNOB_CHANGED_ALWAYS);
  DD::Image::Tooltip(f, "Load color correction information from a .cc, .ccc or .cdl - ASC file.");

  m_fileKnob = File_knob(f, &m_file, "file", "file");
  DD::Image::Tooltip(f,  "Specify the src ASC CDL file, on disk, to use for this transform. "
    "This can be either a .cc, .ccc or .cdl file. If .ccc is specified, the cccid is required.");
  DD::Image::SetFlags(f, DD::Image::Knob::DISABLED | DD::Image::Knob::KNOB_CHANGED_ALWAYS);

  Button(f, "reload", "reload");
  DD::Image::SetFlags(f, DD::Image::Knob::KNOB_CHANGED_ALWAYS);
  DD::Image::Tooltip(f, "Reloads specified files");
  Int_knob(f, &m_reload_version, "version");
  DD::Image::SetFlags(f, DD::Image::Knob::HIDDEN | DD::Image::Knob::KNOB_CHANGED_ALWAYS | DD::Image::Knob::ENDLINE);

  m_cccidKnob = String_knob(f, &m_cccid, "cccid");
  DD::Image::SetFlags(f, DD::Image::Knob::KNOB_CHANGED_ALWAYS | DD::Image::Knob::KNOB_CHANGED_RECURSIVE);
  DD::Image::Tooltip(f, "If the source file is an ASC CDL CCC (color correction collection), "
    "this specifies the id to lookup. OpenColorIO::Contexts (envvars) are obeyed.");

  DD::Image::PyScript_knob(f, "import ocionuke.cdl; ocionuke.cdl.select_cccid_for_filetransform()", "select_cccid", "select cccid");

  DD::Image::Divider(f);

  DD::Image::PyScript_knob(f, "import ocionuke.cdl; ocionuke.cdl.export_as_cc()", "export_cc", "export grade as .cc");
  DD::Image::Tooltip(f, "Export this grade as a ColorCorrection XML file, which can be loaded with the OCIOFileTransform, or using a FileTransform in an OCIO config");
}

void OCIOCDLTransform::readFromFileChanged(const bool value)
{
  assert(m_slopeKnob);
  m_slopeKnob->enable(!value);
  assert(m_offsetKnob);
  m_offsetKnob->enable(!value);
  assert(m_powerKnob);
  m_powerKnob->enable(!value);
  assert(m_saturationKnob);
  m_saturationKnob->enable(!value);

  assert(m_fileKnob);
  m_fileKnob->enable(value);

  if (value) {
    loadCDLFromFile();
  }
}

DD::Image::Hash OCIOCDLTransform::gpuEngine_shader_hash_at_impl(double time)
{
  DD::Image::Hash hash;
  append(hash);

  if (auto versionKnob = knob("version")) {
    hash.append(versionKnob->get_value_at(time));
  }

  if (DD::Image::Knob* workingSpaceKnob = knob(OCIOWorkingSpaceBase::m_workingSpaceKnobName)) {
    hash.append(workingSpaceKnob->get_value_at(time));
  }

  if (DD::Image::Knob* invertKnob = knob("invert")) {
    hash.append(invertKnob->get_value_at(time));
  }

  bool readFromFile = false;
  if (DD::Image::Knob* readFromFileKnob = knob("read_from_file")) {
    readFromFile = static_cast<bool>(readFromFileKnob->get_value_at(time));
    hash.append(readFromFile);
  }

  if (readFromFile) {

    if (DD::Image::Knob* fileKnob = knob("file")) {
      hash.append(fileKnob->get_value_at(time));
    }

    if (DD::Image::Knob* cccidKnob = knob("cccid")) {
      hash.append(cccidKnob->get_value_at(time));
    }
  }
  else {

    if (DD::Image::Knob* slopeKnob = knob("slope")) {
      hash.append(slopeKnob->get_value_at(time));
    }

    if (DD::Image::Knob* offsetKnob = knob("offset")) {
      hash.append(offsetKnob->get_value_at(time));
    }

    if (DD::Image::Knob* powerKnob = knob("power")) {
      hash.append(powerKnob->get_value_at(time));
    }

    if (DD::Image::Knob* saturationKnob = knob("saturation")) {
      hash.append(saturationKnob->get_value_at(time));
    }
  }

  return hash;
}

int OCIOCDLTransform::knob_changed(DD::Image::Knob* k)
{
  if (k->is("read_from_file")) {
    readFromFileChanged(static_cast<bool>(k->get_value()));
    return true;
  }
  else if (k->is("file") || k->is("cccid") || k->is("reload")) {
    loadCDLFromFile();
    return true;
  }

  return OCIOWorkingSpaceBase::knob_changed(k);
}

std::string OCIOCDLTransform::getCCCID() const
{
  const DD::Image::OutputContext ctx;
  const auto cccid = m_cccidKnob->get_text(&ctx);
  return cccid ? cccid : "";
}

void OCIOCDLTransform::loadCDLFromFile()
{
  const DD::Image::OutputContext ctx;
  const auto filepath = m_fileKnob->get_text(&ctx);

  if (!filepath) {
    return;
  }

  if (std::string(filepath).empty()) {
    return;
  }

  try {
    OCIO::ClearAllCaches();
    auto transform = OCIO::CDLTransform::CreateFromFile(filepath, getCCCID().c_str());
    updateKnobsFromTransform(transform);
  }
  catch (const std::exception& e) {
    error(e.what());
  }
}

void OCIOCDLTransform::updateOCIOProcessors()
{
  OCIO::ConstConfigRcPtr config = getConfigForProject();
  auto transform = OCIO::CDLTransform::Create();
  updateTransformFromKnobs(transform);

  assert(rootOp()->knob("workingSpaceLUT")->enumerationKnob());
  const std::string colorManagement = rootOp()->knob("colorManagement")->enumerationKnob()->getSelectedItemString();
  const std::string refName = (colorManagement == "OCIO") ? rootOp()->knob("workingSpaceLUT")->enumerationKnob()->getSelectedItemString() : OCIO::ROLE_SCENE_LINEAR;
  const std::string inputName = knob("working_space")->enumerationKnob()->getSelectedItemString();

  if (refName != inputName) {
    OCIO::GroupTransformRcPtr group = OCIO::GroupTransform::Create();

    auto create_colorspace_transform = [refName, inputName] (const OCIO::TransformDirection direction) {
      auto colorTransform = OCIO::ColorSpaceTransform::Create();
      colorTransform->setSrc(refName.c_str());
      colorTransform->setDst(inputName.c_str());
      colorTransform->setDirection(direction);
      return colorTransform;
    };

    group->appendTransform(create_colorspace_transform(OCIO::TRANSFORM_DIR_FORWARD));
    group->appendTransform(transform);
    group->appendTransform(create_colorspace_transform(OCIO::TRANSFORM_DIR_INVERSE));

    m_processor = config->getProcessor(group, OCIO::TRANSFORM_DIR_FORWARD);
  }
  else {
    // input and output colorspaces are the same - just do the cc transform
    m_processor = config->getProcessor(transform, OCIO::TRANSFORM_DIR_FORWARD);
  }

  m_cpuProcessor = m_processor->getDefaultCPUProcessor();
}

void OCIOCDLTransform::_validate(bool for_real)
{
  // Call through to the base implementation
  OCIOWorkingSpaceBase::_validate(for_real);
  if (this->inErrorState()) {
    return;
  }

  try {
    updateOCIOProcessors();
    setupGPUShaders(m_processor);
  }
  catch(std::exception &e) {
    error(e.what());
    return;
  }

  if (m_processor->isNoOp()) {
    set_out_channels(DD::Image::Mask_None); // prevents engine() from being called
  }
  else {
    set_out_channels(DD::Image::Mask_All);
  }
}

// Note that this is copied by others (OCIODisplay)
void OCIOCDLTransform::in_channels(int, DD::Image::ChannelSet& mask) const
{
  DD::Image::ChannelSet done;
  foreach (c, mask) {
    if (DD::Image::colourIndex(c) < 3 && !(done & c)) {
      done.addBrothers(c, 3);
    }
  }
  mask += done;
}

// See Saturation::pixel_engine for a well-commented example.
// Note that this is copied by others (OCIODisplay)
void OCIOCDLTransform::pixel_engine(
    const DD::Image::Row& in,
    int /* rowY unused */, int rowX, int rowXBound,
    DD::Image::ChannelMask outputChannels,
    DD::Image::Row& out)
{
    int rowWidth = rowXBound - rowX;

    DD::Image::ChannelSet done;
    foreach (requestedChannel, outputChannels) {
        // Skip channels which had their trios processed already,
        if (done & requestedChannel) {
            continue;
        }

        // Pass through channels which are not selected for processing
        // and non-rgb channels.
        if (colourIndex(requestedChannel) >= 3) {
            out.copy(in, requestedChannel, rowX, rowXBound);
            continue;
        }

        DD::Image::Channel rChannel = DD::Image::brother(requestedChannel, 0);
        DD::Image::Channel gChannel = DD::Image::brother(requestedChannel, 1);
        DD::Image::Channel bChannel = DD::Image::brother(requestedChannel, 2);

        done += rChannel;
        done += gChannel;
        done += bChannel;

        const float *rIn = in[rChannel] + rowX;
        const float *gIn = in[gChannel] + rowX;
        const float *bIn = in[bChannel] + rowX;

        float *rOut = out.writable(rChannel) + rowX;
        float *gOut = out.writable(gChannel) + rowX;
        float *bOut = out.writable(bChannel) + rowX;

        // OCIO modifies in-place
        // Note: xOut can equal xIn in some circumstances, such as when the
        // 'Black' (throwaway) scanline is uses. We thus must guard memcpy,
        // which does not allow for overlapping regions.
        if (rOut != rIn) memcpy(rOut, rIn, sizeof(float) * static_cast<unsigned long>(rowWidth));
        if (gOut != gIn) memcpy(gOut, gIn, sizeof(float) * static_cast<unsigned long>(rowWidth));
        if (bOut != bIn) memcpy(bOut, bIn, sizeof(float) * static_cast<unsigned long>(rowWidth));

        try {
            OCIO::PlanarImageDesc img(rOut, gOut, bOut, nullptr, rowWidth, /*height*/ 1);

            applyImageCPU(m_cpuProcessor, img);
        }
        catch(OCIO::Exception &e) {
            error(e.what());
        }
    }
}


void OCIOCDLTransform::updateKnobsFromTransform(OCIO::CDLTransformRcPtr transform) const
{
  auto UpdateKnob3d = [] (DD::Image::Knob& knob, const std::array<double,3>& f3) {
    knob.clear_animated();
    knob.set_value(double(f3.at(0)), 0);
    knob.set_value(double(f3.at(1)), 1);
    knob.set_value(double(f3.at(2)), 2);
  };

  std::array<double,3> values;

  transform->getSlope(values.data());
  assert(m_slopeKnob);
  UpdateKnob3d(*m_slopeKnob, values);

  transform->getOffset(values.data());
  assert(m_offsetKnob);
  UpdateKnob3d(*m_offsetKnob, values);

  transform->getPower(values.data());
  assert(m_powerKnob);
  UpdateKnob3d(*m_powerKnob, values);

  assert(m_saturationKnob);
  m_saturationKnob->clear_animated();
  m_saturationKnob->set_value(double(transform->getSat()));

  assert(m_cccidKnob);
  m_cccidKnob->clear_animated();
  m_cccidKnob->set_text(transform->getID());

  // direction intentionally left out, let Nuke manage that
}

void OCIOCDLTransform::updateTransformFromKnobs(OCIO::CDLTransformRcPtr& transform) const
{
  transform->setSlope(m_slope);
  transform->setOffset(m_offset);
  transform->setPower(m_power);
  transform->setSat(m_saturation);

  transform->setDirection(m_invert ? OCIO::TransformDirection::TRANSFORM_DIR_INVERSE : OCIO::TransformDirection::TRANSFORM_DIR_FORWARD);
  transform->setID(m_cccid.c_str());
}