//--------------------------------------------------------------
// This Fx is based on & modified from the source code by Zhanping Liu, licensed
// with the terms as follows:

//----------------[Start of the License Notice]-----------------
////////////////////////////////////////////////////////////////////////////
///          Line Integral Convolution for Flow Visualization            ///
///                           Initial  Version                           ///
///                             May 15, 1999                             ///
///                           by  Zhanping Liu                           ///
///                      (zhanping@erc.msstate.edu)                      ///
///                    while with  Graphics  Laboratory,                 ///
///                    Peking  University, P. R. China                   ///
///----------------------------------------------------------------------///
///                           Later  Condensed                           ///
///                             May 4,  2002                             ///
///                                                                      ///
///          VAIL (Visualization Analysis & Imaging Laboratory)          ///
///                  ERC  (Engineering Research Center)                  ///
///                     Mississippi State University                     ///
////////////////////////////////////////////////////////////////////////////

////////////////////////////////////////////////////////////////////////////
/// This code was developed  based on  the original algorithm  proposed by///
/// Brian Cabral  and  Leith (Casey) Leedom  in the paper  "Imaging Vector///
/// Fields Using Line Integral Convolution", published  in  Proceedings of///
/// ACM  SigGraph 93,  Aug 2-6,  Anaheim,  California,  pp. 263-270, 1993.///
/// Permission to use, copy, modify, distribute and sell this code for any///
/// purpose is hereby granted without fee,  provided that the above notice///
/// appears in all copies  and  that both that notice and  this permission///
/// appear in supporting documentation. The  developer  of this code makes///
/// no  representations  about  the  suitability  of  this  code  for  any///
/// purpose.  It is provided "as is"  without express or implied warranty.///
////////////////////////////////////////////////////////////////////////////
//-----------------[End of the License Notice]------------------

//--------------------------------------------------------------
#include "iwa_flowblurfx.h"
#include <QThreadPool>

namespace {
const double LINE_SQUARE_CLIP_MAX = 100000.0;
const double VECTOR_COMPONENT_MIN = 0.05;

inline double clamp01(double val) {
  return (val > 1.0) ? 1.0 : (val < 0.0) ? 0.0 : val;
}

// convert sRGB color space to power space
template <typename T = double>
inline T to_linear_color_space(T nonlinear_color, T exposure, T gamma) {
  return std::pow(nonlinear_color, gamma) / exposure;
}
// convert power space to sRGB color space
template <typename T = double>
inline T to_nonlinear_color_space(T linear_color, T exposure, T gamma) {
  return std::pow(linear_color * exposure, T(1) / gamma);
}

}  // namespace
//------------------------------------------------------------
template <typename RASTER, typename PIXEL>
void Iwa_FlowBlurFx::setSourceTileToBuffer(const RASTER srcRas, double4 *buf,
                                           bool isLinear, double gamma) {
  double4 *buf_p = buf;
  for (int j = 0; j < srcRas->getLy(); j++) {
    PIXEL *pix = srcRas->pixels(j);
    for (int i = 0; i < srcRas->getLx(); i++, pix++, buf_p++) {
      (*buf_p).x = double(pix->r) / double(PIXEL::maxChannelValue);
      (*buf_p).y = double(pix->g) / double(PIXEL::maxChannelValue);
      (*buf_p).z = double(pix->b) / double(PIXEL::maxChannelValue);
      (*buf_p).w = double(pix->m) / double(PIXEL::maxChannelValue);
      if (isLinear && (*buf_p).w > 0.0) {
        (*buf_p).x =
            to_linear_color_space((*buf_p).x / (*buf_p).w, 1.0, gamma) *
            (*buf_p).w;
        (*buf_p).y =
            to_linear_color_space((*buf_p).y / (*buf_p).w, 1.0, gamma) *
            (*buf_p).w;
        (*buf_p).z =
            to_linear_color_space((*buf_p).z / (*buf_p).w, 1.0, gamma) *
            (*buf_p).w;
      }
    }
  }
}
//------------------------------------------------------------

template <typename RASTER, typename PIXEL>
void Iwa_FlowBlurFx::setFlowTileToBuffer(const RASTER flowRas, double2 *buf,
                                         double *refBuf) {
  double2 *buf_p   = buf;
  double *refBuf_p = refBuf;
  for (int j = 0; j < flowRas->getLy(); j++) {
    PIXEL *pix = flowRas->pixels(j);
    for (int i = 0; i < flowRas->getLx(); i++, pix++, buf_p++) {
      double val = double(pix->r) / double(PIXEL::maxChannelValue);
      (*buf_p).x = val * 2.0 - 1.0;
      val        = double(pix->g) / double(PIXEL::maxChannelValue);
      (*buf_p).y = val * 2.0 - 1.0;
      if (refBuf != nullptr) {
        *refBuf_p = double(pix->b) / double(PIXEL::maxChannelValue);
        refBuf_p++;
      }
    }
  }
}

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

template <typename RASTER, typename PIXEL>
void Iwa_FlowBlurFx::setReferenceTileToBuffer(const RASTER refRas,
                                              double *buf) {
  double *buf_p = buf;
  for (int j = 0; j < refRas->getLy(); j++) {
    PIXEL *pix = refRas->pixels(j);
    for (int i = 0; i < refRas->getLx(); i++, pix++, buf_p++) {
      // Value = 0.3R 0.59G 0.11B
      (*buf_p) = (double(pix->r) * 0.3 + double(pix->g) * 0.59 +
                  double(pix->b) * 0.11) /
                 double(PIXEL::maxChannelValue);
    }
  }
}

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

template <typename RASTER, typename PIXEL>
void Iwa_FlowBlurFx::setOutputRaster(double4 *out_buf, const RASTER dstRas,
                                     bool isLinear, double gamma) {
  double4 *buf_p = out_buf;
  for (int j = 0; j < dstRas->getLy(); j++) {
    PIXEL *pix = dstRas->pixels(j);
    for (int i = 0; i < dstRas->getLx(); i++, pix++, buf_p++) {
      if (!isLinear || (*buf_p).w == 0.0) {
        pix->r = (typename PIXEL::Channel)(clamp01((*buf_p).x) *
                                           (double)PIXEL::maxChannelValue);
        pix->g = (typename PIXEL::Channel)(clamp01((*buf_p).y) *
                                           (double)PIXEL::maxChannelValue);
        pix->b = (typename PIXEL::Channel)(clamp01((*buf_p).z) *
                                           (double)PIXEL::maxChannelValue);
      } else {
        double val;
        val = to_nonlinear_color_space((*buf_p).x / (*buf_p).w, 1.0, gamma) *
              (*buf_p).w;
        pix->r = (typename PIXEL::Channel)(clamp01(val) *
                                           (double)PIXEL::maxChannelValue);
        val    = to_nonlinear_color_space((*buf_p).y / (*buf_p).w, 1.0, gamma) *
              (*buf_p).w;
        pix->g = (typename PIXEL::Channel)(clamp01(val) *
                                           (double)PIXEL::maxChannelValue);
        val    = to_nonlinear_color_space((*buf_p).z / (*buf_p).w, 1.0, gamma) *
              (*buf_p).w;
        pix->b = (typename PIXEL::Channel)(clamp01(val) *
                                           (double)PIXEL::maxChannelValue);
      }

      pix->m = (typename PIXEL::Channel)(clamp01((*buf_p).w) *
                                         (double)PIXEL::maxChannelValue);
    }
  }
}
//------------------------------------------------------------

void FlowBlurWorker::run() {
  auto sourceAt = [&](int x, int y) {
    if (x < 0 || x >= m_dim.lx || y < 0 || y >= m_dim.ly) return double4();
    return m_source_buf[y * m_dim.lx + x];
  };
  auto flowAt = [&](int x, int y) {
    if (x < 0 || x >= m_dim.lx || y < 0 || y >= m_dim.ly) return double2();
    return m_flow_buf[y * m_dim.lx + x];
  };

  // ロジスティック分布の確率密度関数を格納する(s = 1/3、0-1の範囲で100分割)
  double logDist[101];
  if (m_filterType == Gaussian) {
    double scale = 1.0 / 3.0;
    for (int i = 0; i <= 101; i++) {
      double x   = (double)i / 100.0;
      logDist[i] = std::tanh(x / (2.0 * scale));
    }
  }
  // 0-1に正規化した変数の累積値を返す
  auto getCumulative = [&](double pos) {
    if (pos > 1.0) return 1.0;
    if (m_filterType == Linear)
      return 2.0 * pos - pos * pos;
    else if (m_filterType == Gaussian)
      return logDist[(int)(std::ceil(pos * 100.0))];
    else  // Flat
      return pos;
  };

  int max_ADVCTS = int(m_krnlen * 3);  // MAXIMUM number of advection steps per
                                       // direction to break dead loops

  double4 *out_p = &m_out_buf[m_yFrom * m_dim.lx];

  double *ref_p =
      (m_reference_buf) ? &m_reference_buf[m_yFrom * m_dim.lx] : nullptr;

  // for each pixel in the 2D output LIC image
  for (int j = m_yFrom; j < m_yTo; j++) {
    for (int i = 0; i < m_dim.lx; i++, out_p++) {
      double4 t_acum[2];            // texture accumulators zero-initialized
      double w_acum[2] = {0., 0.};  // weight accumulators zero-initialized

      double blurLength = (ref_p) ? (*ref_p) * m_krnlen : m_krnlen;

      if (blurLength == 0.0) {
        (*out_p) = sourceAt(i, j);
        if (ref_p) ref_p++;
        continue;
      }

      // for either advection direction
      for (int advDir = 0; advDir < 2; advDir++) {
        // init the step counter, curve-length measurer, and streamline seed///
        int advcts =
            0;  // number of ADVeCTion stepS per direction (a step counter)
        double curLen = 0.0;  // CURrent   LENgth of the streamline
        double clp0_x =
            (double)i + 0.5;  // x-coordinate of CLiP point 0 (current)
        double clp0_y =
            (double)j + 0.5;  // y-coordinate of CLiP point 0	(current)

        // until the streamline is advected long enough or a tightly  spiralling
        // center / focus is encountered///

        double2 pre_vctr = flowAt(int(clp0_x), int(clp0_y));
        if (advDir == 1) {
          pre_vctr.x = -pre_vctr.x;
          pre_vctr.y = -pre_vctr.y;
        }

        while (curLen < blurLength && advcts < max_ADVCTS) {
          // access the vector at the sample
          double2 vctr = flowAt(int(clp0_x), int(clp0_y));

          // in case of a critical point
          if (vctr.x == 0.0 && vctr.y == 0.0) {
            w_acum[advDir] = (advcts == 0) ? 1.0 : w_acum[advDir];
            break;
          }

          // negate the vector for the backward-advection case
          if (advDir == 1) {
            vctr.x = -vctr.x;
            vctr.y = -vctr.y;
          }

          // assuming that the flow field vector is bi-directional
          // if the vector is at an acute angle to the previous vector, negate
          // it.
          if (vctr.x * pre_vctr.x + vctr.y * pre_vctr.y < 0) {
            vctr.x = -vctr.x;
            vctr.y = -vctr.y;
          }

          // clip the segment against the pixel boundaries --- find the shorter
          // from the two clipped segments replace  all  if-statements  whenever
          // possible  as  they  might  affect the computational speed
          double tmpLen;
          double segLen = LINE_SQUARE_CLIP_MAX;  // SEGment   LENgth
          segLen        = (vctr.x < -VECTOR_COMPONENT_MIN)
                              ? (std::floor(clp0_x) - clp0_x) / vctr.x
                              : segLen;
          segLen =
              (vctr.x > VECTOR_COMPONENT_MIN)
                  ? (std::floor(std::floor(clp0_x) + 1.5) - clp0_x) / vctr.x
                  : segLen;
          segLen = (vctr.y < -VECTOR_COMPONENT_MIN)
                       ? (((tmpLen = (std::floor(clp0_y) - clp0_y) / vctr.y) <
                           segLen)
                              ? tmpLen
                              : segLen)
                       : segLen;
          segLen = (vctr.y > VECTOR_COMPONENT_MIN)
                       ? (((tmpLen = (std::floor(std::floor(clp0_y) + 1.5) -
                                      clp0_y) /
                                     vctr.y) < segLen)
                              ? tmpLen
                              : segLen)
                       : segLen;

          // update the curve-length measurers
          double prvLen = curLen;
          curLen += segLen;
          segLen += 0.0004;

          // check if the filter has reached either end
          segLen =
              (curLen > m_krnlen) ? ((curLen = m_krnlen) - prvLen) : segLen;

          // obtain the next clip point
          double clp1_x = clp0_x + vctr.x * segLen;
          double clp1_y = clp0_y + vctr.y * segLen;

          // obtain the middle point of the segment as the texture-contributing
          // sample
          double2 samp;
          samp.x = (clp0_x + clp1_x) * 0.5;
          samp.y = (clp0_y + clp1_y) * 0.5;

          // obtain the texture value of the sample
          double4 texVal = sourceAt(int(samp.x), int(samp.y));

          // update the accumulated weight and the accumulated composite texture
          // (texture x weight)
          double W_ACUM = getCumulative(curLen / blurLength);
          // double W_ACUM = curLen;
          // W_ACUM = wgtLUT[int(curLen * len2ID)];

          double smpWgt  = W_ACUM - w_acum[advDir];
          w_acum[advDir] = W_ACUM;
          t_acum[advDir].x += texVal.x * smpWgt;
          t_acum[advDir].y += texVal.y * smpWgt;
          t_acum[advDir].z += texVal.z * smpWgt;
          t_acum[advDir].w += texVal.w * smpWgt;

          // update the step counter and the "current" clip point
          advcts++;
          clp0_x = clp1_x;
          clp0_y = clp1_y;

          pre_vctr = vctr;
          // check if the streamline has gone beyond the flow field///
          if (clp0_x < 0.0 || clp0_x >= double(m_dim.lx) || clp0_y < 0.0 ||
              clp0_y >= double(m_dim.ly))
            break;
        }
      }

      // normalize the accumulated composite texture
      (*out_p).x = (t_acum[0].x + t_acum[1].x) / (w_acum[0] + w_acum[1]);
      (*out_p).y = (t_acum[0].y + t_acum[1].y) / (w_acum[0] + w_acum[1]);
      (*out_p).z = (t_acum[0].z + t_acum[1].z) / (w_acum[0] + w_acum[1]);
      (*out_p).w = (t_acum[0].w + t_acum[1].w) / (w_acum[0] + w_acum[1]);

      if (ref_p) ref_p++;
    }
  }
}

Iwa_FlowBlurFx::Iwa_FlowBlurFx()
    : m_length(5.0)
    , m_linear(false)
    , m_gamma(2.2)
    , m_filterType(new TIntEnumParam(Linear, "Linear"))
    , m_referenceMode(new TIntEnumParam(REFERENCE, "Reference Image")) {
  addInputPort("Source", m_source);
  addInputPort("Flow", m_flow);
  addInputPort("Reference", m_reference);

  bindParam(this, "length", m_length);
  bindParam(this, "linear", m_linear);
  bindParam(this, "gamma", m_gamma);
  bindParam(this, "filterType", m_filterType);
  bindParam(this, "referenceMode", m_referenceMode);

  m_length->setMeasureName("fxLength");
  m_length->setValueRange(0., 100.0);
  m_gamma->setValueRange(0.2, 5.0);

  m_filterType->addItem(Gaussian, "Gaussian");
  m_filterType->addItem(Flat, "Flat");

  m_referenceMode->addItem(FLOW_BLUE_CHANNEL, "Blue Channel of Flow Image");
}

void Iwa_FlowBlurFx::doCompute(TTile &tile, double frame,
                               const TRenderSettings &settings) {
  // do nothing if Source or Flow port is not connected
  if (!m_source.isConnected() || !m_flow.isConnected()) {
    tile.getRaster()->clear();
    return;
  }

  // output size
  TDimension dim = tile.getRaster()->getSize();

  double fac     = sqrt(fabs(settings.m_affine.det()));
  double krnlen  = fac * m_length->getValue(frame);
  int max_ADVCTS = int(krnlen * 3);  // MAXIMUM number of advection steps per
                                     // direction to break dead loops

  if (max_ADVCTS == 0) {
    // No blur will be done. The underlying fx may pass by.
    m_source->compute(tile, frame, settings);
    return;
  }

  bool isLinear = m_linear->getValue();
  double gamma  = m_gamma->getValue(frame);

  // obtain Source memory buffer (RGBA)
  TTile sourceTile;
  m_source->allocateAndCompute(sourceTile, tile.m_pos, dim, tile.getRaster(),
                               frame, settings);
  // allocate buffer
  double4 *source_buf;
  TRasterGR8P source_buf_ras(dim.lx * dim.ly * sizeof(double4), 1);
  source_buf_ras->lock();
  source_buf       = (double4 *)source_buf_ras->getRawData();
  TRaster32P ras32 = tile.getRaster();
  TRaster64P ras64 = tile.getRaster();
  if (ras32)
    setSourceTileToBuffer<TRaster32P, TPixel32>(sourceTile.getRaster(),
                                                source_buf, isLinear, gamma);
  else if (ras64)
    setSourceTileToBuffer<TRaster64P, TPixel64>(sourceTile.getRaster(),
                                                source_buf, isLinear, gamma);

  // obtain Flow memory buffer (XY)
  TTile flowTile;
  m_flow->allocateAndCompute(flowTile, tile.m_pos, dim, tile.getRaster(), frame,
                             settings);
  // allocate buffer
  double2 *flow_buf;
  TRasterGR8P flow_buf_ras(dim.lx * dim.ly * sizeof(double2), 1);
  flow_buf_ras->lock();
  flow_buf = (double2 *)flow_buf_ras->getRawData();

  double *reference_buf = nullptr;
  TRasterGR8P reference_buf_ras;
  if (m_referenceMode->getValue() == FLOW_BLUE_CHANNEL ||
      m_reference.isConnected()) {
    reference_buf_ras = TRasterGR8P(dim.lx * dim.ly * sizeof(double), 1);
    reference_buf_ras->lock();
    reference_buf = (double *)reference_buf_ras->getRawData();
  }

  if (ras32)
    setFlowTileToBuffer<TRaster32P, TPixel32>(flowTile.getRaster(), flow_buf,
                                              reference_buf);
  else if (ras64)
    setFlowTileToBuffer<TRaster64P, TPixel64>(flowTile.getRaster(), flow_buf,
                                              reference_buf);

  if (m_referenceMode->getValue() == REFERENCE && m_reference.isConnected()) {
    TTile referenceTile;
    m_reference->allocateAndCompute(referenceTile, tile.m_pos, dim,
                                    tile.getRaster(), frame, settings);
    if (ras32)
      setReferenceTileToBuffer<TRaster32P, TPixel32>(referenceTile.getRaster(),
                                                     reference_buf);
    else if (ras64)
      setReferenceTileToBuffer<TRaster64P, TPixel64>(referenceTile.getRaster(),
                                                     reference_buf);
  }

  // buffer for output raster
  double4 *out_buf;
  TRasterGR8P out_buf_ras(dim.lx * dim.ly * sizeof(double4), 1);
  out_buf_ras->lock();
  out_buf = (double4 *)out_buf_ras->getRawData();

  int activeThreadCount = QThreadPool::globalInstance()->activeThreadCount();
  // use half of the available threads
  int threadAmount       = std::max(1, activeThreadCount / 2);
  FILTER_TYPE filterType = (FILTER_TYPE)m_filterType->getValue();
  QList<QThread *> threadList;
  int tmpStart = 0;
  for (int t = 0; t < threadAmount; t++) {
    int tmpEnd =
        (int)std::round((float)(dim.ly * (t + 1)) / (float)threadAmount);

    FlowBlurWorker *worker =
        new FlowBlurWorker(source_buf, flow_buf, out_buf, reference_buf, dim,
                           krnlen, tmpStart, tmpEnd, filterType);
    worker->start();
    threadList.append(worker);
    tmpStart = tmpEnd;
  }

  for (auto worker : threadList) {
    worker->wait();
    delete worker;
  }

  source_buf_ras->unlock();
  flow_buf_ras->unlock();
  if (reference_buf) reference_buf_ras->unlock();

  // render vector field with red & green channels
  if (ras32)
    setOutputRaster<TRaster32P, TPixel32>(out_buf, ras32, isLinear, gamma);
  else if (ras64)
    setOutputRaster<TRaster64P, TPixel64>(out_buf, ras64, isLinear, gamma);

  out_buf_ras->unlock();
}

bool Iwa_FlowBlurFx::doGetBBox(double frame, TRectD &bBox,
                               const TRenderSettings &info) {
  bBox = TConsts::infiniteRectD;
  return true;
}

bool Iwa_FlowBlurFx::canHandle(const TRenderSettings &info, double frame) {
  return true;
}

FX_PLUGIN_IDENTIFIER(Iwa_FlowBlurFx, "iwa_FlowBlurFx")