// Copyright (c) 2025 The Foundry Visionmongers Ltd. All Rights Reserved.

#include "Dust_Knob.h"

#include "DDImage/Row.h"
#include "DDImage/Tile.h"
#include "DDImage/Knobs.h"
#include "DDImage/DDMath.h"
#include "DDImage/Channel.h"

#include <memory>

using namespace DD::Image;

static const char* const CLASS = "DustBust";
static const char* const HELP =
  "Clones multiple areas from a source to a destination, "
#ifdef __APPLE__
  "based on dust points created by alt+command+click in the viewer.\n\n"
#else
  "based on dust points created by alt+ctrl+click in the viewer.\n\n"
#endif
  "Source can be on same frame or on another frame determined "
  "from user-specified frame offset. "
  "Frame offset is global to all dust points on a particular frame "
  "(but frame offset can be animated over time if desired). "
  "Edge softness and size controls act on all dust points on "
  "a particular frame. "
  "The source position can be different for each dust point.";

// Define the C++ class that is the new operator. This may be a subclass
// of Iop, or of some subclass of Iop such as Blur:
class DustBustIop : public Iop
{
  // These are the locations the user interface will store into:
  double ui_enable{1}; // Master adjust to strength of clone.
  bool ui_viewsource{0}; // Allows user to quickly look at source (cleaner) image.
  DustPoints ui_points;
  double ui_frameoffset{1.0}; // Affects source. Used in inputContext().
  double ui_hardness{0.5};

  // Variables calculated in _request():
  std::vector<double> m_pt; // Dust point bounding boxes.
  double m_dest[4]{0.0, 0.0, 0.0, 0.0}; // Area surrounding all dust points.
  double m_hardness{0};
  Channel K_maskchannel{Chan_Black};

public:
  DustBustIop(Node* node) : Iop(node) {}
  ~DustBustIop() = default;
  void _validate(bool) override;
  void _request(int, int, int, int, ChannelMask, int) override;
  void engine(int y, int x, int r, ChannelMask, Row &) override;
  void knobs(Knob_Callback) override;
  int maximum_inputs() const override { return 1; }
  int minimum_inputs() const override { return 1; }
  int split_input(int n) const override { return 2; }
  const OutputContext& inputContext(int, int, OutputContext&) const override;
  const char* Class() const override { return CLASS; }
  const char* node_help() const override { return HELP; }
  static const Op::Description d;
};

// The time for image input n:
const OutputContext& DustBustIop::inputContext(int i, int n, OutputContext& context) const
{
  context = outputContext();
  context.setFrame(context.frame() + n * ui_frameoffset);
  return context;
}

// This is a function that creates an instance of the operator, and is
// needed for the Op::Description to work:
static Op* DustBust_c(Node* node) { return new DustBustIop(node); }

// The Op::Description is how NUKE knows what the name of the operator is
// and how to create one.
const Op::Description DustBustIop::d(CLASS, DustBust_c);

// When the operator runs, "validate" is called first. This function copies
// the "info" data from the input operator(s), modifies it according to
// what this operator does, and saves this information so that operators
// connected to this one can see it:
void DustBustIop::_validate(bool)
{
  copy_info();
  info_.turn_on(K_maskchannel);
}

// After validate is done, "open" is called. This is passed a "viewport"
// which describes which channels and a rectangular "bounding box" that will
// be requested. The operator must translate this to the area that will
// be requested from its inputs and call request() on them:
void DustBustIop::_request(int x, int y, int r, int t, ChannelMask channels, int count)
{
  m_pt.resize(ui_points.size * 4); // Allocate space to store dust bounds

  if (ui_points.size > 0) {
    m_dest[0] = m_dest[2] = ui_points[0].x;
    m_dest[1] = m_dest[3] = ui_points[0].y;
    for (int p = 0; p < ui_points.size; ++p) {
      // Bounding box for dust point p:
      m_pt[p * 4 + 0] = ui_points[p].x - fabs(ui_points[p].w) * 1.01;
      m_pt[p * 4 + 1] = ui_points[p].y - fabs(ui_points[p].h) * 1.01;
      m_pt[p * 4 + 2] = ui_points[p].x + fabs(ui_points[p].w) * 1.01;
      m_pt[p * 4 + 3] = ui_points[p].y + fabs(ui_points[p].h) * 1.01;
      // Adjust total bounding box surrounding all points:
      m_dest[0] = MIN(m_dest[0], m_pt[p * 4 + 0]);
      m_dest[1] = MIN(m_dest[1], m_pt[p * 4 + 1]);
      m_dest[2] = MAX(m_dest[2], m_pt[p * 4 + 2]);
      m_dest[3] = MAX(m_dest[3], m_pt[p * 4 + 3]);
    }
  }

  if (ui_hardness < 1) {
    m_hardness = 1 / (1 - ui_hardness);
  }
  else {
    m_hardness = -1.;
  }

  if (ui_viewsource) {
    input1().request(x, y, r, t, channels, count);
    return;
  }

  input0().request(x, y, r, t, channels, count);

  Box v;
  v.clear();
  for (int p = 0; p < ui_points.size; ++p) {
    v.merge(static_cast<int>(floor(ui_points[p].source_x - ui_points[p].w)) - 1,
      static_cast<int>(floor(ui_points[p].source_y - ui_points[p].h)) - 1,
      static_cast<int>(ceil(ui_points[p].source_x + ui_points[p].w)) + 1,
      static_cast<int>(ceil(ui_points[p].source_y + ui_points[p].h)) + 1);
  }
  input1().request(v.x(), v.y(), v.r(), v.t(), channels, count);
}

// This is the operator that does all the work. For each line in the area
// passed to request(), this will be called. It must calculate the image data
// for a region at vertical position y, and between horizontal positions
// x and r, and write it to the passed row structure. Usually this works
// by asking the input for data, and modifying it:
void DustBustIop::engine(int y, int x, int r, ChannelMask channels, Row& out)
{
  // DustBust 'fixes' areas that are designated by dust boxes.
  // It does this by mixing areas of input0 that are inside dust boxes
  // with 'clean' areas from input1.

  // First, do some trivial-case checks
  // (copy entire source frame, or do nothing).
  // m_dest[] was earlier calculated in _request(),
  // and stores a single bounding box for the union of all the dust boxes:

  if (ui_viewsource) { // User wants to look at source frame
    out.get(input1(), y, x, r, channels); // So just copy from input1 and finish
    return;
  }
  if (y < floor(m_dest[1]) || y > floor(m_dest[3])) { // Outside y bounds of dest
    out.get(input0(), y, x, r, channels); // so just copy from the input and finish
    return;
  }

  // End trivial checks. Perform cloning:
  const float ycen = float(y) + .5; // Y-center of pixels at y (used for clone weighting).

  // First, read from existing (dusty) image:
  out.get(input0(), y, x, r, channels);

  // Modify it where necessary (clone in 'clean areas' from input1):
  for (auto z : channels) {
    const float* inchan = out[z] + x;
    float* outchan = out.writable(z) + x;
    float srcval_xa;
    float srcval_xb;
    float srcval = 0.0f;
    float* MASK = out.writable(K_maskchannel) + x;

    // Get the sources ('cleaners') for each dust from the input:
    // 2 rows to allow fractional y position with interpolation
    // (necessary if srcvec[1] has a fractional part).
    // Need to get these rows for each dust box (ui_points.size of them).
    std::vector<std::unique_ptr<Row>> srcrowa(ui_points.size);
    std::vector<std::unique_ptr<Row>> srcrowb(ui_points.size);

    std::vector<float> fracsy(ui_points.size);
    // Fractional blend between the two rows (for interpolation)
    // (as with the rows, we store this value for each dust box).

    // Arrays of float vectors for inputs.
    std::vector<const float*> a_row(ui_points.size);
    std::vector<const float*> b_row(ui_points.size);

    // For each dust box, get the source (cleaner pixels):
    for (int p = 0; p < ui_points.size; ++p) {
      // Vector to source from dusty pixels:
      const double srcvec[2] = {
        ui_points[p].source_x - ui_points[p].x,
        ui_points[p].source_y - ui_points[p].y
      };

      // X bounds of source pixels:
      const int sx0 = static_cast<int>(floor(m_pt[p * 4 + 0] + srcvec[0]));
      const int sx1 = static_cast<int>(floor(m_pt[p * 4 + 2] + srcvec[0]));

      // Y coordinate of source pixels:
      const double sy = y + srcvec[1];
      // Fractional and integral parts of sy (for interpolation):
      const int isy = static_cast<int>(floor(sy));
      fracsy[p] = sy - isy;

      // Get 2 rows (a and b) from source:
      // (2 rows because y may have fractional part):
      srcrowa[p] = std::make_unique<Row>(sx0, sx1 + 2);
      srcrowa[p]->get(input1(), isy,   sx0, sx1 + 2, channels);
      srcrowb[p] = std::make_unique<Row>(sx0, sx1 + 2);
      srcrowb[p]->get(input1(), isy + 1, sx0, sx1 + 2, channels);
      a_row[p] = (*srcrowa[p])[z];
      b_row[p] = (*srcrowb[p])[z];
    }
    // We now have rows from the source which we can clone into the dest (dusty) pixels

    for (int X = x; X < r; ++X) { // Loop through all pixels on row y.
      if (K_maskchannel != Chan_Black) {
        *MASK = 0.0;
      }
      *outchan = *inchan; // Initially, take an untouched (dusty) copy.
      // This may get modified inside dust box loop.

      if (X >= static_cast<int>(m_dest[0]) && X <= static_cast<int>(m_dest[2])) { // Ff within overall dust bounds (all dust boxes)

        const float xcen = float(X) + .5; // X-center of pixels at X (used for clone weighting).

        // For each dust box, clone clean pixels from source:
        for (int p = 0; p < ui_points.size; ++p) {
          // X coordinate of source pixel:
          const double sx = X + ui_points[p].source_x - ui_points[p].x;
          // Fractional and integral parts of sx (for interpolation):
          const int isx = static_cast<int>(floor(sx));
          const double fracsx = sx - isx;

          float bl = 0.f; // Strength of blend for current dust on current pixel

          // Ellipse-based strength of clone:
          // uses xcen and cyen in relation to dust box to compute a nice soft-edged ellipse:
          if (xcen > m_pt[p * 4 + 0] && xcen <= m_pt[p * 4 + 2]) { // If (xcen,ycen) inside dust box

            // Determine source pixel value by linear interpolation of nearest 4 source pixels:
            srcval_xa = a_row[p][isx];
            srcval_xa += (b_row[p][isx] - srcval_xa) * fracsy[p];
            srcval_xb = a_row[p][isx + 1];
            srcval_xb += (b_row[p][isx + 1] - srcval_xb) * fracsy[p];
            srcval = srcval_xa + (srcval_xb - srcval_xa) * fracsx;

            // Compute (adjusted) distance of (xcen,ycen) from dust center:
            const float fx = (xcen - ui_points[p].x) / ui_points[p].w;
            const float fy = (ycen - ui_points[p].y) / ui_points[p].h;
            const float r5 = static_cast<float>(sqrt(fx * fx + fy * fy)) * .5f + .5f;

            if (r5 < 1) { // If inside ellipse
              // Compute bl for nice soft-edged ellipse:
              if (m_hardness >= 0) {
                bl = 1 - r5 * (1 - r5) * 4;
                bl = 1 - powf(bl, m_hardness);
                bl *= bl;
              }
              else {
                bl = 1.f;
              }
            }
            else {
              bl = 0.f;
            }
          }

          if (bl > 0) { // If blend strength > 0
            // Blend output with source (clean pixels):
            *outchan += (srcval - *outchan) * (bl * ui_enable);
            if (K_maskchannel != Chan_Black) {
              *MASK += bl;
            }
          }

        }
      }

      ++outchan;
      ++inchan;
      ++MASK;
    } // End X loop.

  } // End z/channel loop.

}

// The knobs function describes each user control so the user interface
// can be built. The possible controls are listed in DDImage/Knobs.h:
void DustBustIop::knobs(Knob_Callback f)
{
  Float_knob(f, &ui_enable, "enable");
  Bool_knob(f, &ui_viewsource, "view_source", "view source");
  CustomKnob1(Dust_Knob, f, &ui_points, "points");
  Double_knob(f, &ui_frameoffset, IRange(-1, 1), "frame_offset", "frame offset");
  SetFlags(f, Knob::EARLY_STORE);
  Double_knob(f, &ui_hardness, IRange(0, 1), "edge_hardness", "edge hardness");
  Channel_knob(f, &K_maskchannel, 1, "maskchannel", "Output Mask");
}