// tcg includes
#include "tcg/tcg_misc.h"

#include "trop.h"

/*! \file terodilate.cpp

This file contains an implementation of a greyscale (ie per-channel)
erode/dilate
morphological operator, following the van Herk/Gil-Werman O(row*cols) algorithm.

An extension with circular structuring element is attempted - unfortunately I
could
not retrieve a copy of Miyataka's paper about that, which seemingly claimed
O(rows * cols) too. The implemented algorithm is a sub-optimal
O(rows*cols*radius).
*/

//********************************************************
//    Auxiliary  functions
//********************************************************

namespace {

template <typename Pix>
void copyMatte(const TRasterPT<Pix> &src,
               const TRasterPT<typename Pix::Channel> &matte) {
  typedef typename Pix::Channel Chan;

  int y, lx = src->getLx(), ly = src->getLy();
  for (y = 0; y != ly; ++y) {
    Pix *s, *sBegin = src->pixels(y), *sEnd = sBegin + lx;
    Chan *m, *mBegin = matte->pixels(y);

    for (s = sBegin, m = mBegin; s != sEnd; ++s, ++m) *m = s->m;
  }
}

//--------------------------------------------------------------

template <typename Pix>
void copyChannels_erode(const TRasterPT<Pix> &src,
                        const TRasterPT<typename Pix::Channel> &matte,
                        const TRasterPT<Pix> &dst) {
  typedef typename Pix::Channel Chan;

  // Just assemble src and matte, remembering to depremultiply src pixels before
  // applying the new matte

  double fac;

  int y, lx = src->getLx(), ly = src->getLy();
  for (y = 0; y != ly; ++y) {
    const Pix *s, *sBegin = src->pixels(y), *sEnd = sBegin + lx;
    Pix *d, *dBegin = dst->pixels(y);

    Chan *m, *mBegin = matte->pixels(y);

    for (s = sBegin, d = dBegin, m = mBegin; s != sEnd; ++s, ++d, ++m) {
      fac  = double(*m) / double(s->m);
      d->r = fac * s->r, d->g = fac * s->g, d->b = fac * s->b, d->m = *m;
    }
  }
}

//--------------------------------------------------------------

template <typename Pix>
void copyChannels_dilate(const TRasterPT<Pix> &src,
                         const TRasterPT<typename Pix::Channel> &matte,
                         const TRasterPT<Pix> &dst) {
  typedef typename Pix::Channel Chan;

  // Trickier - since src is presumably premultiplied, increasing its pixels'
  // alpha by direct
  // substitution would expose the excessive RGB discretization of pixels with a
  // low matte value.
  // So, let's just put the pixels on a black background. It should do fine.

  double max = Pix::maxChannelValue;

  int y, lx = src->getLx(), ly = src->getLy();
  for (y = 0; y != ly; ++y) {
    const Pix *s, *sBegin = src->pixels(y), *sEnd = sBegin + lx;
    Pix *d, *dBegin = dst->pixels(y);

    const Chan *m, *mBegin = matte->pixels(y);

    for (s = sBegin, d = dBegin, m = mBegin; s != sEnd; ++s, ++d, ++m) {
      *d   = *s;
      d->m = s->m + (1.0 - s->m / max) * *m;
    }
  }
}

}  // namespace

//********************************************************
//    EroDilate  algorithms
//********************************************************

namespace {

template <typename Chan>
struct MaxFunc {
  inline Chan operator()(const Chan &a, const Chan &b) {
    return std::max(a, b);
  }
};

template <typename Chan>
struct MinFunc {
  inline Chan operator()(const Chan &a, const Chan &b) {
    return std::min(a, b);
  }
};

//--------------------------------------------------------------

// NOTE: src and dst must be NOT OVERLAPPING (eg src != dst)

template <typename Chan, typename Func>
void erodilate_row(int len, const Chan *src, int sIncr, Chan *dst, int dIncr,
                   int rad, double radR, Func func) {
  assert(rad >= 0);

  // Segment the row of specified length into wCount windows of max wSize
  // elements
  int w, wSize = 2 * rad + 1, wCount = len / wSize + 1;
  int swIncr = wSize * sIncr, srIncr = rad * sIncr;
  int dwIncr = wSize * dIncr, drIncr = rad * dIncr;

  const Chan *s, *sEnd = src + len * sIncr;
  Chan *d, *dEnd       = dst + len * dIncr;

  double one_radR = (1.0 - radR);

  for (w = 0; w != wCount; ++w) {
    Chan *dwBegin = dst + w * dwIncr, *dwEnd = std::min(dwBegin + dwIncr, dEnd);

    // Compute prefixes
    const Chan *swBegin = src + std::max(w * swIncr - srIncr - sIncr, 0),
               *swEnd =
                   src + std::min(w * swIncr + srIncr + sIncr, len * sIncr);

    s = swEnd - sIncr, d = dst + ((s - src) / sIncr) * dIncr +
                           drIncr;  // d already decremented by dIncr

    Chan val = *s, oldVal;

    for (s -= sIncr; (d >= dEnd) && (s >= swBegin);
         s -= sIncr, d -= dIncr)  // s decremented here
    {
      assert(s >= src);
      assert(s < sEnd);
      assert((s - src) % sIncr == 0);
      assert(d >= dst);
      assert((d - dst) % dIncr == 0);

      val = func(oldVal = val, *s);
    }

    for (; s >= swBegin; s -= sIncr, d -= dIncr) {
      assert(s >= src);
      assert(s < sEnd);
      assert((s - src) % sIncr == 0);
      assert(d >= dst);
      assert(d < dEnd);
      assert((d - dst) % dIncr == 0);

      val = func(oldVal = val, *s);
      *d = (oldVal == val) ? val : one_radR * oldVal + radR * val;
    }

    for (d = std::min(d, dEnd - dIncr); d >= dwBegin; d -= dIncr) {
      assert(d >= dst);
      assert(d < dEnd);
      assert((d - dst) % dIncr == 0);

      val = func(oldVal = val, 0);
      *d = (oldVal == val) ? val : one_radR * oldVal + radR * val;
    }

    // Compute suffixes
    swBegin = src + w * swIncr + srIncr,
    swEnd   = std::min(swBegin + swIncr + sIncr, sEnd);
    if (swBegin >= swEnd) continue;

    s = swBegin, d = dwBegin;

    val = *s;

    for (s += sIncr; (s < swEnd); s += sIncr, d += dIncr) {
      assert(s >= src);
      assert(s < sEnd);
      assert((s - src) % sIncr == 0);
      assert(d >= dst);
      assert(d < dEnd);
      assert((d - dst) % dIncr == 0);

      val = func(oldVal = val, *s);
      *d = func(*d, (oldVal == val) ? val : one_radR * oldVal + radR * val);
    }

    for (; d < dwEnd; d += dIncr) {
      assert(d >= dst);
      assert(d < dEnd);
      assert((d - dst) % dIncr == 0);

      val = func(oldVal = val, 0);
      *d = func(*d, (oldVal == val) ? val : one_radR * oldVal + radR * val);
    }
  }
}

//--------------------------------------------------------------

template <typename Pix, typename Chan>
void erodilate_chan(const TRasterPT<Pix> &src, const TRasterPT<Chan> &dst,
                    double radius, bool dilate) {
  assert(radius > 0.0);

  int radI    = tfloor(radius);
  double radR = radius - radI;

  // Using a temporary raster to keep intermediate results. This allows us to
  // perform a cache-friendly iteration in the separable/square kernel case
  int x, y, lx = src->getLx(), ly = src->getLy();

  // Perform rows erodilation
  TRasterPT<Chan> temp(ly, lx);  // Notice transposition plz

  {
    if (dilate)
      for (y = 0; y != ly; ++y)
        ::erodilate_row(lx, &src->pixels(y)->m, 4, temp->pixels(0) + y, ly,
                        radI, radR, MaxFunc<Chan>());
    else
      for (y = 0; y != ly; ++y)
        ::erodilate_row(lx, &src->pixels(y)->m, 4, temp->pixels(0) + y, ly,
                        radI, radR, MinFunc<Chan>());
  }

  // Perform columns erodilation
  {
    if (dilate)
      for (x = 0; x != lx; ++x)
        ::erodilate_row(ly, temp->pixels(x), 1, dst->pixels(0) + x,
                        dst->getWrap(), radI, radR, MaxFunc<Chan>());
    else
      for (x = 0; x != lx; ++x)
        ::erodilate_row(ly, temp->pixels(x), 1, dst->pixels(0) + x,
                        dst->getWrap(), radI, radR, MinFunc<Chan>());
  }
}

//--------------------------------------------------------------

template <typename Pix>
void rect_erodilate(const TRasterPT<Pix> &src, const TRasterPT<Pix> &dst,
                    double radius) {
  typedef typename Pix::Channel Chan;

  if (radius == 0.0) {
    // No-op case
    TRop::copy(dst, src);
    return;
  }

  bool dilate = (radius >= 0.0);

  // Perform columns erodilation
  TRasterPT<Chan> temp(src->getLx(), src->getLy());
  ::erodilate_chan(src, temp, fabs(radius), dilate);

  // Remember that we have just calculated the matte values. We still have to
  // apply them to the old RGB
  // values, which requires depremultiplying from source matte and
  // premultiplying with the new one.
  if (dilate)
    ::copyChannels_dilate(src, temp, dst);
  else
    ::copyChannels_erode(src, temp, dst);
}

}  // namespace

//********************************************************
//    EroDilate  round algorithm
//********************************************************

namespace {

template <typename Chan, typename Func>
void erodilate_quarters(int lx, int ly, Chan *src, int sIncrX, int sIncrY,
                        Chan *dst, int dIncrX, int dIncrY, double radius,
                        double shift, Func func) {
  double sqRadius     = sq(radius);
  double squareHeight = radius * M_SQRT1_2;
  int squareHeightI   = tfloor(squareHeight);

  // For every arc point
  int arcY;
  for (arcY = -squareHeightI; arcY <= squareHeightI; ++arcY) {
    // Calculate x and weights
    double sqArcY = sq(arcY);
    assert(sqRadius >= sqArcY);

    double x = shift + sqrt(sqRadius - sqArcY) - squareHeight;

    int arcX = tfloor(x);
    double w = x - arcX, one_w = 1.0 - w;

    // Build dst area influenced by the arc point. Func with 0 outside that.
    TRect bounds(0, 0, lx, ly);

    TRect dRect(bounds * (bounds + TPoint(-arcX, -arcY)));
    TRect sRect(bounds * (bounds + TPoint(arcX, arcY)));

    int sy, dy;

    // Func with 0 before dRect.y0
    for (dy = 0; dy < dRect.y0; ++dy) {
      Chan *d, *dBegin = dst + dy * dIncrY, *dEnd = dBegin + lx * dIncrX;
      for (d = dBegin; d != dEnd; d += dIncrX) {
        // assert(d >= dst); assert(d < dEnd); assert((d-dst) % dIncrX == 0);
        *d = func(*d, 0);
      }
    }

    // Func with 0 after dRect.y1
    for (dy = dRect.y1; dy < ly; ++dy) {
      Chan *d, *dBegin = dst + dy * dIncrY, *dEnd = dBegin + lx * dIncrX;
      for (d = dBegin; d != dEnd; d += dIncrX) {
        // assert(d >= dst); assert(d < dEnd); assert((d-dst) % dIncrX == 0);
        *d = func(*d, 0);
      }
    }

    // For every dst pixel in the area, Func with the corresponding pixel in src
    for (dy = dRect.y0, sy = sRect.y0; dy != dRect.y1; ++dy, ++sy) {
      Chan *d, *dLine = dst + dy * dIncrY, *dBegin = dLine + dRect.x0 * dIncrX;
      Chan *s, *sLine = src + sy * sIncrY, *sBegin = sLine + sRect.x0 * sIncrX,
               *sEnd = sLine + sRect.x1 * sIncrX;

      Chan *sLast = sEnd - sIncrX;  // sLast would lerp with sEnd

      for (d = dBegin, s = sBegin; s != sLast;
           d += dIncrX, s += sIncrX)  // hence we stop before it
      {
        // assert(s >= src); assert(s < sEnd); assert((s-src) % sIncrX == 0);
        // assert(d >= dst); assert(d < dEnd); assert((d-dst) % dIncrX == 0);

        *d = func(*d, *s * one_w + *(s + sIncrX) * w);
      }

      // assert(s >= src); assert(s < sEnd); assert((s-src) % sIncrX == 0);
      // assert(d >= dst); assert(d < dEnd); assert((d-dst) % dIncrX == 0);

      *d = func(*d, *s * one_w);  // lerp sLast with 0
    }
  }
}

//--------------------------------------------------------------

template <typename Pix>
void circular_erodilate(const TRasterPT<Pix> &src, const TRasterPT<Pix> &dst,
                        double radius) {
  typedef typename Pix::Channel Chan;

  if (radius == 0.0) {
    // No-op case
    TRop::copy(dst, src);
    return;
  }

  // Ok, the idea is: consider the maximal embedded square in our circular
  // structuring element.
  // Erodilating by it consists in the consecutive erodilation by rows and
  // columns with the same
  // 'square' radius. Now, it's easy to see that the square could be 'bent' so
  // that one of its
  // edges matches that of a 1/4 of the circle's edge, while remaining inside
  // the circle.
  // Erodilating by the bent square can be achieved by erodilating first by rows
  // or column for
  // the square edge radius, followed by perpendicular erodilationg with a
  // fourth of our
  // circumference. Sum the 4 erodilations needed to complete the circumference
  // - and it's done.

  // NOTE: Unfortunately, the above decomposition has lots of intersections
  // among the pieces - yet
  // it's simple enough and removes an O(radius) from the naive algorithm. Could
  // be done better?

  // First, build the various erodilation data
  bool dilate = (radius >= 0.0);
  radius      = fabs(radius);

  double inner_square_diameter = radius * M_SQRT2;

  double shift =
      0.25 *
      inner_square_diameter;  // Shift of the bent square SE needed to avoid
                              // touching the circumference on the other side
  double row_filter_radius = 0.5 * (inner_square_diameter - shift);
  double cseShift = 0.5 * shift;  // circumference structuring element shift

  int lx = src->getLx(), ly = src->getLy();

  TRasterPT<Chan> temp1(lx, ly), temp2(lx, ly);

  int radI    = tfloor(row_filter_radius);
  double radR = row_filter_radius - radI;

  if (dilate) {
    temp2->fill(0);  // Initialize with a Func-neutral value

    if (row_filter_radius > 0.0)
      for (int y = 0; y != ly; ++y)
        ::erodilate_row(lx, &src->pixels(y)->m, 4, temp1->pixels(y), 1, radI,
                        radR, MaxFunc<Chan>());
    else
      ::copyMatte(src, temp1);

    ::erodilate_quarters(lx, ly, temp1->pixels(0), 1, lx, temp2->pixels(0), 1,
                         lx, radius, cseShift, MaxFunc<Chan>());
    ::erodilate_quarters(lx, ly, temp1->pixels(0) + lx - 1, -1, lx,
                         temp2->pixels(0) + lx - 1, -1, lx, radius, cseShift,
                         MaxFunc<Chan>());

    if (row_filter_radius > 0.0)
      for (int x = 0; x != lx; ++x)
        ::erodilate_row(ly, &src->pixels(0)[x].m, 4 * src->getWrap(),
                        temp1->pixels(0) + x, lx, radI, radR, MaxFunc<Chan>());
    else
      ::copyMatte(src, temp1);

    ::erodilate_quarters(ly, lx, temp1->pixels(0), lx, 1, temp2->pixels(0), lx,
                         1, radius, cseShift, MaxFunc<Chan>());
    ::erodilate_quarters(ly, lx, temp1->pixels(0) + lx * ly - 1, -lx, -1,
                         temp2->pixels(0) + lx * ly - 1, -lx, -1, radius,
                         cseShift, MaxFunc<Chan>());
  } else {
    temp2->fill((std::numeric_limits<Chan>::max)());  // Initialize with a
                                                      // Func-neutral value

    if (row_filter_radius > 0.0)
      for (int y = 0; y != ly; ++y)
        ::erodilate_row(lx, &src->pixels(y)->m, 4, temp1->pixels(y), 1, radI,
                        radR, MinFunc<Chan>());
    else
      ::copyMatte(src, temp1);

    ::erodilate_quarters(lx, ly, temp1->pixels(0), 1, lx, temp2->pixels(0), 1,
                         lx, radius, cseShift, MinFunc<Chan>());
    ::erodilate_quarters(lx, ly, temp1->pixels(0) + lx - 1, -1, lx,
                         temp2->pixels(0) + lx - 1, -1, lx, radius, cseShift,
                         MinFunc<Chan>());

    if (row_filter_radius > 0.0)
      for (int x = 0; x != lx; ++x)
        ::erodilate_row(ly, &src->pixels(0)[x].m, 4 * src->getWrap(),
                        temp1->pixels(0) + x, lx, radI, radR, MinFunc<Chan>());
    else
      ::copyMatte(src, temp1);

    ::erodilate_quarters(ly, lx, temp1->pixels(0), lx, 1, temp2->pixels(0), lx,
                         1, radius, cseShift, MinFunc<Chan>());
    ::erodilate_quarters(ly, lx, temp1->pixels(0) + lx * ly - 1, -lx, -1,
                         temp2->pixels(0) + lx * ly - 1, -lx, -1, radius,
                         cseShift, MinFunc<Chan>());
  }

  // Remember that we have just calculated the matte values. We still have to
  // apply them to the old RGB
  // values, which requires depremultiplying from source matte and
  // premultiplying with the new one.
  if (dilate)
    ::copyChannels_dilate(src, temp2, dst);
  else
    ::copyChannels_erode(src, temp2, dst);
}

}  // namespace

//********************************************************
//    EroDilate  main functions
//********************************************************

void TRop::erodilate(const TRasterP &src, const TRasterP &dst, double radius,
                     ErodilateMaskType type) {
  assert(src->getSize() == dst->getSize());

  src->lock(), dst->lock();

  if ((TRaster32P)src && (TRaster32P)dst) switch (type) {
    case ED_rectangular:
      ::rect_erodilate<TPixel32>(src, dst, radius);
      break;
    case ED_circular:
      ::circular_erodilate<TPixel32>(src, dst, radius);
      break;
    default:
      assert(!"Unknown mask type");
      break;
    }
  else if ((TRaster64P)src && (TRaster64P)dst)
    switch (type) {
    case ED_rectangular:
      ::rect_erodilate<TPixel64>(src, dst, radius);
      break;
    case ED_circular:
      ::circular_erodilate<TPixel64>(src, dst, radius);
      break;
    default:
      assert(!"Unknown mask type");
      break;
    }
  else
    assert(!"Unsupported raster type!");

  src->unlock(), dst->unlock();
}