/**
 * OpenColorIO Display Iop.
 */

#include "OCIODisplay.h"

namespace OCIO = OCIO_NAMESPACE;

#include <algorithm>
#include <string>
#include <sstream>
#include <fstream>
#include <stdexcept>

#include <DDImage/Channel.h>
#include <DDImage/PixelIop.h>
#include <DDImage/NukeWrapper.h>
#include <DDImage/Row.h>
#include <DDImage/Knobs.h>
#include <DDImage/Enumeration_KnobI.h>
#include <DDImage/Colorspace_KnobI.h>
#include <DDImage/plugins.h>

#ifndef FN_EXAMPLE_PLUGIN
#include <OCIOUtility/OCIOUtility.h>
#else
#include "OCIOUtility.h"
#endif

enum eDisplayedLayer
{
  eDisplayedLayerRGB = 0,
  eDisplayedLayerR,
  eDisplayedLayerG,
  eDisplayedLayerB,
  eDisplayedLayerA,
  eDisplayedLayerLum,
  eDisplayedLayerMat
};

OCIODisplay::OCIODisplay(Node *n) : DD::Image::OCIOPluginBase(n)
{
    m_layersToProcess = DD::Image::Mask_All;
    m_hasLists = false;
    m_colorSpaceIndex = m_displayIndex = m_viewIndex = 0;
    m_displayKnob = m_viewKnob = nullptr;
    m_gain = 1.0;
    m_gamma = 1.0;
    m_channel = eDisplayedLayerRGB;
    m_allowGPUAcceleration = false;
    m_transform = OCIO::DisplayViewTransform::Create();

    setExtendedOpProperties(std::make_shared<DD::Image::ExtendedOpProperties>(kNumOCIOContextVariables));

    configChanged();
}

OCIODisplay::~OCIODisplay()
{

}

void OCIODisplay::configChanged()
{
  int defaultDispIndex = 0;

  try {
    const auto config = getConfigForProject();
    m_displayNames.clear();
    if (!config) {
      throw OCIO::Exception("OCIO config is invalid");
    }

    const std::string& defaultDisplay = config->getDefaultDisplay();

    for (int i = 0; i < config->getNumDisplays(); ++i) {
      std::string display = config->getDisplay(i);
      m_displayNames.push_back(display);
      if (display == defaultDisplay) {
        defaultDispIndex = i;
      }
    }
  }
  catch(OCIO::Exception& e) {
    std::cerr << "OCIODisplay: " << e.what() << std::endl;
    return;
  }
  catch (...) {
    std::cerr << "OCIODisplay: Unknown exception during OCIO setup." << std::endl;
    return;
  }

  m_displayCstrNames.clear();
  for(unsigned int i=0; i<m_displayNames.size(); ++i) {
      m_displayCstrNames.push_back(m_displayNames[i].c_str());
  }
  m_displayCstrNames.push_back(nullptr);

  m_displayIndex = defaultDispIndex;
  if (m_displayKnob)  {
    DD::Image::Enumeration_KnobI *enumKnob = m_displayKnob->enumerationKnob();
    if (enumKnob) {
      enumKnob->menu(m_displayNames);
      m_displayKnob->set_value(m_displayIndex);
    }
  }

  refreshDisplayTransforms();

  m_hasLists = !(m_displayNames.empty() || m_viewNames.empty());

  if (!m_hasLists) {
      std::cerr << "OCIODisplay: Missing one or more of display devices, or display transforms." << std::endl;
  }
}

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

  // Using append() takes into account the current config.
  // TODO: No doubt there's some way to animate the config selection but this doesn't
  // support that. In practice, I can't believe anyone would do this.
  append(hash);

  if (m_colorspaceKnob) {
    hash.append(m_colorspaceKnob->get_value_at(time));
  }

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

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

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

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

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

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

  return hash;
}

void OCIODisplay::knobs(DD::Image::Knob_Callback f)
{
    updateConfigBeforeCreatingKnobs(f);

    const int oldDisplayIndex = m_displayIndex;

    m_colorspaceKnob = DD::Image::Colorspace_knob(f, &m_colorSpaceIndex, OCIO::ROLE_SCENE_LINEAR, "colorspace", "input");
    DD::Image::Tooltip(f, "\nInput data is taken to be in this color transform. \n\n"
                          "The first menu level is the color space family name. \n\n"
                          "The sub menu items are the names of the color spaces themselves."
                          "\n\n"
                          "<b>Warning</b>"
                          "\n"
                          "The label for this knob has been changed - however the knob name is the same as in 13.0"
                          " - in a future major release the knob name will also be changed.");

    m_displayKnob = DD::Image::Enumeration_knob(f,
        &m_displayIndex, &m_displayCstrNames[0], "display", "display device");
    DD::Image::Tooltip(f, "Display device for output.");
    DD::Image::SetFlags(f, DD::Image::Knob::ALWAYS_SAVE);

    m_viewKnob = DD::Image::Enumeration_knob(f,
        &m_viewIndex, &m_viewCstrNames[0], "view", "view transform");
    DD::Image::Tooltip(f, "Display transform for output.");
    DD::Image::SetFlags(f, DD::Image::Knob::ALWAYS_SAVE);

    DD::Image::Newline(f);
    DD::Image::Bool_knob(f, &m_invert, "invert", "invert direction");
    DD::Image::Tooltip(f, "Specify the transform direction.");

    DD::Image::Float_knob(f, &m_gain, DD::Image::IRange(1.0 / 64.0f, 64.0f), "gain");
    DD::Image::SetFlags(f, DD::Image::Knob::NO_ANIMATION | DD::Image::Knob::LOG_SLIDER);
    DD::Image::Tooltip(f, "Exposure adjustment, in scene-linear, prior to the display transform.");

    DD::Image::Float_knob(f, &m_gamma, DD::Image::IRange(0.01, 4.0), "gamma");
    DD::Image::SetFlags(f, DD::Image::Knob::NO_ANIMATION | DD::Image::Knob::LOG_SLIDER);
    DD::Image::Tooltip(f, "Gamma correction applied after the display transform.");

    static const char* const channelvalues[] = {
      "RGB",
      "R",
      "G",
      "B",
      "A",
      "Luminance",
      "Matte overlay",
      nullptr
    };
    DD::Image::Enumeration_knob(f, &m_channel, channelvalues, "channel_selector", "channel view");
    DD::Image::SetFlags(f, DD::Image::Knob::NO_ANIMATION | DD::Image::Knob::NO_UNDO | DD::Image::Knob::NODEGRAPH_ONLY);
    DD::Image::Tooltip(f, "Specify which channels to view (prior to the display transform).");

    if (nodeContext() != eTimeline) {
      DD::Image::Divider(f);
    }

    DD::Image::Input_ChannelSet_knob(f, &m_layersToProcess, 0, "layer", "layer");
    DD::Image::SetFlags(f, DD::Image::Knob::NODEGRAPH_ONLY);
    DD::Image::Tooltip(f, "Set which layer to process. This should be a layer with rgb data.");

    DD::Image::Newline(f);
    DD::Image::Bool_knob(f, &m_allowGPUAcceleration, "allowGPUAcceleration");
    DD::Image::SetFlags(f, DD::Image::Knob::INVISIBLE);

    getExtendedOpProperties()->extendedPropertiesKnobs(f, "Context", nullptr, DD::Image::OCIOPluginBase::kContextTooltip);

    if (!f.makeKnobs()) {
      validateConfig();

      if (oldDisplayIndex != m_displayIndex) {
        refreshDisplayTransforms();
      }
    }
}

void OCIODisplay::append(DD::Image::Hash& localhash)
{
    // TODO: Hang onto the context, what if getting it
    // (and querying getCacheID) is expensive?
    try {
        OCIO::ConstConfigRcPtr config = getConfigForProject();
        if (!config) {
          throw OCIO::Exception("OCIO config is invalid");
        }

        OCIO::ConstContextRcPtr context = getLocalContext();
        std::string configCacheID = config->getCacheID(context);
        localhash.append(configCacheID);

        // This is required due to our custom channel overlay mode post-processing
        localhash.append(m_channel);
    }
    catch(OCIO::Exception &e) {
        error(e.what());
        return;
    }
}

void OCIODisplay::_validate(bool for_real)
{
    input0().validate(for_real);

    validateConfig();

    if(!m_hasLists) {
        error("Missing one or more of display devices, or display transforms.");
        return;
    }

    try {
        OCIO::ConstConfigRcPtr config = getConfigForProject();
        if (!config) {
          throw OCIO::Exception("OCIO config is invalid");
        }

        const std::string colorspaceName = m_colorspaceKnob->enumerationKnob()->getSelectedItemString();

        m_transform->setSrc(colorspaceName.c_str());
        m_transform->setDisplay(m_displayCstrNames[m_displayIndex]);
        m_transform->setView(m_viewCstrNames[m_viewIndex]);
        m_transform->setDirection(m_invert ? OCIO::TransformDirection::TRANSFORM_DIR_INVERSE : OCIO::TransformDirection::TRANSFORM_DIR_FORWARD);
        m_viewingPipeline = OCIO::LegacyViewingPipeline::Create();
        m_viewingPipeline->setDisplayViewTransform(m_transform);

        // Specify an (optional) linear color correction
        {
            double m44[16];
            double offset4[4];

            const double slope4[] = { m_gain, m_gain, m_gain, m_gain };
            OCIO::MatrixTransform::Scale(m44, offset4, slope4);

            OCIO::MatrixTransformRcPtr mtx =  OCIO::MatrixTransform::Create();
            mtx->setMatrix(m44);
            mtx->setOffset(offset4);

            m_viewingPipeline->setLinearCC(mtx);
        }

        // Specify an (optional) post-display transform.
        {
            float exponent = 1.0/std::max(1e-2, m_gamma);
            const double exponent4[] = { exponent, exponent, exponent, exponent };
            OCIO::ExponentTransformRcPtr cc =  OCIO::ExponentTransform::Create();
            cc->setValue(exponent4);
            m_viewingPipeline->setDisplayCC(cc);
        }

        // Add Channel swizzling
        {
            int channelHot[4] = { 0, 0, 0, 0};

            switch(m_channel)
            {
            case eDisplayedLayerRGB: // RGB
                channelHot[0] = 1;
                channelHot[1] = 1;
                channelHot[2] = 1;
                channelHot[3] = 1;
                break;
            case eDisplayedLayerR: // R
                channelHot[0] = 1;
                break;
            case eDisplayedLayerG: // G
                channelHot[1] = 1;
                break;
            case eDisplayedLayerB: // B
                channelHot[2] = 1;
                break;
            case eDisplayedLayerA: // A
                channelHot[3] = 1;
                break;
            default:
                break;
            case eDisplayedLayerLum: // Luma
                channelHot[0] = 1;
                channelHot[1] = 1;
                channelHot[2] = 1;
                break;
            case eDisplayedLayerMat: //  Channel overlay mode. Do rgb, and then swizzle later
                channelHot[0] = 1;
                channelHot[1] = 1;
                channelHot[2] = 1;
                channelHot[3] = 1;
                break;
            }

            double lumacoef[3];
            config->getDefaultLumaCoefs(lumacoef);
            double m44[16];
            double offset[4];
            OCIO::MatrixTransform::View(m44, offset, channelHot, lumacoef);
            OCIO::MatrixTransformRcPtr swizzle = OCIO::MatrixTransform::Create();
            swizzle->setMatrix(m44);
            swizzle->setOffset(offset);
            m_viewingPipeline->setChannelView(swizzle);
        }

        OCIO::ConstContextRcPtr context = getLocalContext();
        m_processor = m_viewingPipeline->getProcessor(config, context);
        m_cpuProcessor = m_processor->getDefaultCPUProcessor();
        setupGPUShaders(m_processor);
    }
    catch (OCIO::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);
    }

    DD::Image::PixelIop::_validate(for_real);
}

/*static*/ DD::Image::Channel OCIODisplay::FindAlphaChannel(DD::Image::ChannelMask outputChannels, DD::Image::Channel rChannel, DD::Image::Channel gChannel, DD::Image::Channel bChannel)
{
    // Bug 45992 - outputChannels.next will only work if the channel indicies are incrementing, otherwise it'll return
    // the gChannel again (bChannel + 1) and the output will be tinted strangely. Hence when the channel indicies are
    // decrementing we need to use outputChannels.previous to get the correct channel.
    const bool incrementingChannels = bChannel > gChannel;
    DD::Image::Channel aChannel = incrementingChannels ? outputChannels.next(bChannel) : outputChannels.previous(bChannel);

    // This Op always assumes that any channel following an RGB channel is an alpha channel.
    // However, in some cases when switching between RGB and a custom visible layer in the viewer,
    // this isn't the case.
    // This is because the cache code combines the channels that were stored in the cache with channels
    // requested by the viewer.
    // In some cases, this can result in channel list such as R, G, B, customR, customG, customB, customA,
    // which means that the aChannel returned from outputChannels.next(bChannel) is actually customR.
    // This can cause tinting of the output image.
    // To check that aChannel is actually asociated with the r,g,b channels, we check if they are brothers
    // If not, we set they aChannel to the black (throwaway) channel.
    // Note that the brother relationship sometimes only works one way for the alpha channel - i.e r, g, b are brothers of a,
    // but a is not a brother of r, g or b.
    if ( rChannel != DD::Image::brother(aChannel, 0)
      || gChannel != DD::Image::brother(aChannel, 1)
      || bChannel != DD::Image::brother(aChannel, 2)  ) {

        aChannel = DD::Image::Chan_Black;
    }

    return aChannel;
}

// Note: Same as OCIO ColorSpace::in_channels.
void OCIODisplay::in_channels(int /* n unused */, DD::Image::ChannelSet& mask) const
{
    DD::Image::ChannelSet done;
    foreach (c, mask) {
        if ((m_layersToProcess & c) && DD::Image::colourIndex(c) < 4 && !(done & c)) {
            done.addBrothers(c, 4);
        }
    }
    mask += done;
}

// Note: Same as OCIO ColorSpace::pixel_engine.
void OCIODisplay::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 (!(m_layersToProcess & requestedChannel)) {
            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);
        DD::Image::Channel aChannel = FindAlphaChannel(outputChannels, rChannel, gChannel, bChannel);

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

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

        float *rOut = out.writable(rChannel) + rowX;
        float *gOut = out.writable(gChannel) + rowX;
        float *bOut = out.writable(bChannel) + rowX;
        float *aOut = out.writable(aChannel) + 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)*rowWidth);
        if (gOut != gIn) memcpy(gOut, gIn, sizeof(float)*rowWidth);
        if (bOut != bIn) memcpy(bOut, bIn, sizeof(float)*rowWidth);
        if (aOut != aIn) memcpy(aOut, aIn, sizeof(float)*rowWidth);

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

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

        // Hack to emulate Channel overlay mode
        if (m_channel == eDisplayedLayerMat) {
            for (int i=0; i<rowWidth; ++i) {
                rOut[i] = rOut[i] + (1.0f - rOut[i]) * (0.5f * aOut[i]);
                gOut[i] = gOut[i] - gOut[i] * (0.5f * aOut[i]);
                bOut[i] = bOut[i] - bOut[i] * (0.5f * aOut[i]);
            }
        }
    }
}

void OCIODisplay::refreshDisplayTransforms()
{
  try {
    OCIO::ConstConfigRcPtr config = getConfigForProject();
    if (!config) {
      throw OCIO::Exception("OCIO config is invalid");
    }

    _config = config;

    const char * display = m_displayCstrNames[m_displayIndex];
    std::string defaultViewName = config->getDefaultView(display);

    // Try to maintain an old transform name, or use the default.
    bool hasOldViewName = false;
    std::string oldViewName = "";

    auto python_appended_view_from_different_display = [this]() {
        return m_viewNames.size() > 0 && m_viewIndex >= static_cast<int>(m_viewNames.size());
      };

    auto has_old_view = [this]() {
        return m_viewIndex >= 0 && m_viewIndex < static_cast<int>(m_viewNames.size());
      };

    if (python_appended_view_from_different_display()) {
      hasOldViewName = true;
      oldViewName = m_viewKnob->enumerationKnob()->getSelectedItemString();
    }
    else if (has_old_view()) {
        hasOldViewName = true;
        oldViewName = m_viewNames[m_viewIndex];
    }
    int defaultViewIndex = 0, newViewIndex = -1;

    m_viewCstrNames.clear();

    m_viewNames = Foundry::OCIOUtility::getDisplayActiveViews(config, display);

    for (unsigned int i = 0; i < m_viewNames.size(); i++) {
      const std::string& view = m_viewNames[i];
      if (hasOldViewName && view == oldViewName) {
        newViewIndex = i;
      }
      if (view == defaultViewName) {
        defaultViewIndex = i;
      }
    }

    for (unsigned int i = 0; i < m_viewNames.size(); ++i) {
        m_viewCstrNames.push_back(m_viewNames[i].c_str());
    }
    m_viewCstrNames.push_back(nullptr);

    if (newViewIndex == -1) {
        newViewIndex = defaultViewIndex;
    }

    m_viewIndex = newViewIndex;
    if (m_viewKnob) {
        DD::Image::Enumeration_KnobI *enumKnob = m_viewKnob->enumerationKnob();
        enumKnob->menu(m_viewNames);
        m_viewKnob->set_value(newViewIndex);
    }
  }
  catch (OCIO::Exception &e) {
      error(e.what());
      return;
  }
}

int OCIODisplay::knob_changed(DD::Image::Knob *k)
{
  if (k == m_displayKnob && m_displayKnob != nullptr) {
      refreshDisplayTransforms();
      return 1;
  }
  else {
      return 0;
  }
}

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

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

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

const char* OCIODisplay::node_help() const
{
    // TODO more detailed help text
    return "Use OpenColorIO to convert for a display device.";
}

DD::Image::Op* build(Node *node)
{
    DD::Image::NukeWrapper *op = new DD::Image::NukeWrapper(new OCIODisplay(node));
    op->noMix();
    op->noMask();
    op->noChannels(); // prefer our own channels control without checkboxes / alpha
    op->noUnpremult();
    return op;
}