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

#include "DDImage/DeepOp.h"
#include "DDImage/Knobs.h"
#include "DDImage/Thread.h"
#include "DDImage/Iop.h"
#include "DDImage/Row.h"
#include "math.h"
#include "PxNukeDeepUtils/All.h"

using namespace DD::Image;
using namespace PxNukeDeepUtils;

#include "DeepReTarget.h"

#include "DDImage/DeepSample.h"
#include "DDImage/DeepComposite.h"

#include "Math/fnFunctions.h"
#include <memory>
#include <stdio.h>

#include <DDImage/DDMath.h>

static const char* CLASS = "DeepMerge";

static const char * const kMergeModes[] = { "combine", "holdout", 0};

class DeepMerge : public DeepOnlyOp
{
  enum MergeOp { eCombine, eHoldout };
  bool _dropHidden;
  int _operation;

public:
  DeepMerge(Node* node) : DeepOnlyOp(node)
  {
    _dropHidden = false;
    _operation = (int)eCombine;
  }

  virtual int minimum_inputs() const
  {
    return 2;
  }

  virtual Op* op()
  {
    return this;
  }

  virtual int maximum_inputs() const {
    return 999;
  }

  virtual bool test_input(int idx, Op*op) const {
    return dynamic_cast<DeepOp*>(op);
  }

  DeepOp* dinput(int n)
  {
    return dynamic_cast<DeepOp*>(Op::input(n));
  }

  const char* input_label(int n, char* buf) const
  {
    if (!n)
      return "B";
    if (n > 1 || inputs() > 2) {
      sprintf(buf, "A%d", n);
      return buf;
    }
    return "A";
  }

  const char* node_help() const
  {
    return "Merges two or more deep data inputs to produce a composite deep data stream.";
  }

  const char* Class() const
  {
    return CLASS;
  }

  void knobs(Knob_Callback f)
  {
    Enumeration_knob(f, &_operation, kMergeModes ,"operation", "operation");
    Tooltip(f, "Merge operation.");
    Tooltip(f, "How to merge the inputs."
               "<ul>"
               "<li>combine - Combine the samples from the A and B inputs."
               "<li>holdout - Holdout samples from the B input by the samples in the A input(s)"
               "</ul>" );

    Bool_knob(f, &_dropHidden, "drop_hidden", "drop hidden samples");
    Tooltip(f, "Drop samples that are behind other samples with alpha 1 (ie those that are entirely obscured)");
  }

  int knob_changed(Knob *k)
  {
    if ( &Knob::showPanel || k->is("operation") ) {
      knob( "drop_hidden")->enable( _operation == eCombine );
      return 1;
    }
    return 0;
  }


  //! get the extra channels needed apart from those being passed through.
  DD::Image::ChannelSet extraChans() const
  {
    if ( _operation == eHoldout )
      return Mask_DeepFront | Mask_DeepBack | Mask_Alpha;
    else if ( _dropHidden )
      //! If _dropHidden is on we need depth and alpha
      return Mask_DeepFront | Mask_Alpha;

    return Mask_None;
  }

  void _validate(bool real)
  {
    bool any = false;

    for (int i = 0; i < inputs(); i++) {
      if (DeepOp* dop = dinput(i)) {
        dop->validate(real);
        
        if (!any) {
          _deepInfo = dop->deepInfo();
          any = true;
        } else {
          _deepInfo.merge(dop->deepInfo());
        }
      }
    }
  }

  void getDeepRequests(DD::Image::Box box, const DD::Image::ChannelSet& channels, int count, std::vector<RequestData>& requests)
  {
    DD::Image::ChannelSet reqChans = channels + extraChans();

    for (int i = 0; i < inputs(); i++) {
      if (dinput(i)) {
        requests.push_back(RequestData(dinput(i), box, reqChans, count));
      }
    }
  }

  bool doDeepEngine(DD::Image::Box box, const ChannelSet& channels, DeepOutputPlane& plane)
  {
    plane = DeepOutputPlane(channels, box);

    bool done = true;

    std::vector<DeepPlane> inPlanes(inputs());
    
    DD::Image::ChannelSet getChans = channels + extraChans();

    for (int i = 0; i < inputs(); i++) {
      if (dinput(i)) {
        if (!dinput(i)->deepEngine(box, getChans, inPlanes[i]))
          return false;
      }
    }

    if ( _operation == eCombine )
      done = doDeepCombine(box, channels, inPlanes, plane);
    else if ( _operation == eHoldout )
      done = doDeepHoldout2(box, channels, inPlanes, plane);
    else
      abort();

    return done;
  }

  bool doDeepCombine(DD::Image::Box& box, const ChannelSet& channels, std::vector<DeepPlane> &inPlanes, DeepOutputPlane& plane )
  {
    bool done = true;
    int numChans = channels.size();

    for (DD::Image::Box::iterator it = box.begin(); it != box.end(); it++) {

      DeepOutPixel pixelData;

      pixelData.clear();

      double firstSolid = DBL_MAX;

      if (_dropHidden) {
        for (int i = 0; i < inputs(); i++)
          if (dinput(i)) {
            DeepPixel deepPixel = inPlanes[i].getPixel(it);
            for (size_t j = 0; j < deepPixel.getSampleCount(); j++)
              if (deepPixel.getUnorderedSample(j, Chan_Alpha) >= 1)
                firstSolid = std::min(firstSolid, (double)deepPixel.getUnorderedSample(j, Chan_DeepFront));
          }
      }

      for (int i = 0; i < inputs(); i++) {
        if (dinput(i)) {
          DeepPixel deepPixel = inPlanes[i].getPixel(it);

          pixelData.reserveMore(deepPixel.getSampleCount() * numChans);

          for (size_t j = 0; j < deepPixel.getSampleCount(); j++) {
            if (_dropHidden) {
              double front = deepPixel.getUnorderedSample(j, Chan_DeepFront);
              if (front > firstSolid)
                continue;
            }

            foreach (z, channels) {
              pixelData.push_back(deepPixel.getUnorderedSample(j, z));
            }
          }
        }
      }

      plane.addPixel(pixelData);
    }
    return done;
  }


  // Here's the algorithm for the DeepHoldout
  //
  // It works by inserting a new 'black' channel representing the holdout and a 'white' channel
  // representing the main object.
  //
  // It then composites them both working out what final target alpha of the combined black and white object is.
  // It also works what the final alpha of the main object is without the holdout.
  //
  // It then reweights the 'main' samples by a density ratio of final target alpha to the original main alpha
  //
  // Along the way it removes samples where the target alpha become zero, and keeps samples that aren't affected by the holdout.
  //
  //
  // JW 9/11/12

  bool doDeepHoldout2(DD::Image::Box& box, const ChannelSet& channelsRequired, std::vector<DeepPlane> &inPlanes, DeepOutputPlane& plane )
  {
    bool done = true;
    ChannelSet channels = channelsRequired;
    channels += Chan_DeepBack;
    channels += Chan_DeepFront;
    channels += Chan_Alpha;

    int numChans = channels.size();
    ChannelMap chanMap(channels);

    bool haveBackIn = false;
    bool haveBackHoldout = false;

    // check if we only have one 'holdout' input
    DeepPlane* holdoutInputPlane = 0;
    bool moreThanOneHoldout = false;
    for (int i = 1; i < inputs(); i++) {
      if ( holdoutInputPlane ) {
        moreThanOneHoldout = true;
        break;
      }
      if ( dinput(i) ) {
        holdoutInputPlane = &inPlanes[i];
        haveBackHoldout = dinput(i)->deepInfo().channels().contains( Chan_DeepBack );
      }
    }
    haveBackIn = dinput(0) ? dinput(0)->deepInfo().channels().contains( Chan_DeepBack ) : false;

    for (DD::Image::Box::iterator it = box.begin(); it != box.end(); it++) {

      DeepOutPixel pixelData;
      pixelData.clear();

      // must have at least the main B input
      if ( ! dinput(0) ) {
        plane.addPixel(pixelData);
        continue;
      }

      // no deep input plane, just return input0
      if ( ! holdoutInputPlane) {
        DeepPixel pixel = inPlanes[0].getPixel(it);
        plane.addPixel(pixel);
        continue;
      }

      DeepOutPixel holdoutData;
      holdoutData.clear();

      std::auto_ptr<DeepPixel> holdoutMerged;
      DeepPixel deepPixelIn = inPlanes[0].getPixel(it);
      DeepPixel holdoutSingle = holdoutInputPlane->getPixel(it);

      if ( moreThanOneHoldout ) {
        // loop over all the 'holdout' inputs and merge them into one pixel

        for (int i = 1; i < inputs(); i++) {
          if (dinput(i)) {
            DeepPixel deepPixel = inPlanes[i].getPixel(it);

            holdoutData.reserveMore(deepPixel.getSampleCount() * numChans);

            for (size_t j = 0; j < deepPixel.getSampleCount(); j++) {
              foreach (z, channels) {
                holdoutData.push_back(deepPixel.getUnorderedSample(j, z));
              }
            }
          }
        }
        holdoutMerged = std::auto_ptr<DeepPixel>( new DeepPixel( chanMap, holdoutData.data(), holdoutData.size(), DeepPixel::eUnordered ) );
      }

      DeepPixel& deepPixelHoldout = moreThanOneHoldout ? *holdoutMerged : holdoutSingle;

      // now we have all the holdout samples combined into one sample 'deepPixelHoldout'

      // count samples
      int holdoutSampleCount = deepPixelHoldout.getSampleCount();

      if ( holdoutSampleCount == 0 ) {
          // nothing to holdout - copy main input to output
          DeepPixel pixel = inPlanes[0].getPixel(it);
          plane.addPixel(pixel);
          continue;
      }

      // get front of holdout

      float zHoldoutFront = deepPixelHoldout.getOrderedSample( holdoutSampleCount-1, Chan_DeepFront);
      float zHoldoutBack = deepPixelHoldout.getOrderedSample( holdoutSampleCount-1,Chan_DeepBack );

      assert ( deepPixelHoldout.channels().contains( Chan_DeepFront ));
      assert ( deepPixelHoldout.channels().contains( Chan_DeepBack ));

      // now munge them together with holdout pixels Mask being 0.0 * alpha
      // and main holdout Mask pixels being 1.0 * alpha

      DeepOutPixel combinedData;
      combinedData.clear();

      assert (!channels.contains( Chan_Mask) );

      ChannelSet combinedChannels = channels;
      combinedChannels += Chan_Mask;
      ChannelMap combinedChanMap( combinedChannels );

      combinedData.reserveMore(deepPixelHoldout.getSampleCount() * ( numChans + 1 ) );
      for (size_t j = 0; j < deepPixelHoldout.getSampleCount(); j++) {
        foreach (z, combinedChannels) {
          if ( z == Chan_Alpha || z == Chan_DeepFront || z == Chan_DeepBack || z == Chan_Z )
            combinedData.push_back(deepPixelHoldout.getUnorderedSample(j, z));
          else
            // 0.0 * alpha - black in all the holdout channels so the color doesn't 'bleed' into the main input if the samples overlap
            combinedData.push_back(0.0);
        }
      }

      combinedData.reserveMore(deepPixelIn.getSampleCount() * ( numChans + 1 ) );

      for (size_t j = 0; j < deepPixelIn.getSampleCount(); j++) {
        float alpha = deepPixelIn.getUnorderedSample(j, Chan_Alpha);
        foreach (z, combinedChannels) {
          if ( z != Chan_Mask )
            combinedData.push_back(deepPixelIn.getUnorderedSample(j, z));
          else
            combinedData.push_back( alpha );  // 1.0 * alpha
         }
      }

      // anything to do ?
      if ( combinedData.size() == 0 ) {
        plane.addHole();
        continue;
      }

      // de-overlapp the result if required for volumetrics, will create a new possibly larger set of samples
      DeepPixel combinedPixel = combinedData.getPixel(combinedChanMap, DeepPixel::eUnordered);
      DeepOutPixel deoverlappedSamples;
      DeepOutPixel *pDeoverLappedSamples = &combinedData;

      if ( DetectOverlappingSamples(combinedPixel) ) {
        CombineOverlappingSamples(combinedChanMap, combinedPixel, deoverlappedSamples);
        pDeoverLappedSamples = &deoverlappedSamples;
      }


      // create two new sets of samples, before and after the holdout
      // samples that are not held out are left untouched

      DeepOutPixel notHeldOut;
      DeepOutPixel heldOut;

      bool atHoldout = false;
      float zHeldOutFront;
      DeepPixel deoverlappedPixel = pDeoverLappedSamples->getPixel(combinedChanMap, DeepPixel::eUnordered);
      int sampleCount = deoverlappedPixel.getSampleCount();

      for ( int samp = sampleCount-1 ; samp >= 0; samp--) {
        float front = deoverlappedPixel.getOrderedSample(samp, Chan_DeepFront) ;

        if ( !atHoldout && ( ( front > zHoldoutFront && front < zHoldoutBack )  ||
                             front == zHoldoutFront )  ) {
          atHoldout = true;
          zHeldOutFront = front;
        }
        if ( atHoldout ) {
          foreach(z, combinedChannels)
            heldOut.push_back( deoverlappedPixel.getOrderedSample(samp, z) );
        } else {
          foreach(z, combinedChannels)
            notHeldOut.push_back( deoverlappedPixel.getOrderedSample(samp, z) );
        }
      }

      // Now composite the Mask channel to work out the 'target alpha' ( held out )
      // and the original main alpha ( unheld out )

      // Also make a new set of main samples that will be reweighted later by the new target alpha
      // What we consider is a 'main' sample is if its mask is non-zero. ( eg not a holdout sample )

      // Note that at this stage the new set of main samples may be different to the orignal
      // input set as the deoverlapping process may have created new samples

      DeepPixel heldOutPixel = heldOut.getPixel( combinedChanMap, DeepPixel::eUnordered );
      DeepOutPixel filteredResult;
      filteredResult.clear();

      // Gather the viz's at 'main' each pixel, compute input and holdout viz.
      int numSamps =  heldOutPixel.getSampleCount();
      std::vector<double> vizs( ( size_t ) heldOutPixel.getSampleCount() );
      double inputViz = 1.0;
      double targetAlpha = 0.0;
      int filteredSampCount = 0;

      for ( int samp = 0; samp < numSamps; samp++ ) {
        double mask = heldOutPixel.getOrderedSample( samp, Chan_Mask );
        double alpha = heldOutPixel.getOrderedSample( samp, Chan_Alpha );
        bool isMainSample = Foundry::Math::IsNonZero( mask );
        double viz = ClampViz( 1.0 - double( alpha ) );

        if ( isMainSample ) {
          vizs[filteredSampCount++] = viz;
          inputViz *= viz;
          foreach(z, channels)
            filteredResult.push_back( heldOutPixel.getOrderedSample( samp, z) );
        }

        targetAlpha = targetAlpha * ( viz ) + mask;
        //std::cout << "Compositing " << heldOutPixel.getOrderedSample( samp, Chan_DeepFront ) << ": " << " alpha "
          //       << alpha << " " << " mask " << mask <<  " target " <<  targetAlpha << std::endl;
      }

      double inputAlpha = ClampAlpha( 1.0 - inputViz );

      // put the untouched non-heldout samples back into the output
      DeepOutPixel outputPixel;
      DeepPixel notHeldOutPixel = notHeldOut.getPixel(combinedChanMap, DeepPixel::eZAscending);
      int notHeldOutSampCount = notHeldOutPixel.getSampleCount();
      for( int samp = 0; samp < notHeldOutSampCount; samp++)
        foreach( z, channelsRequired )
          outputPixel.push_back( notHeldOutPixel.getUnorderedSample(samp, z));

      // if there's no pixels left to reweight we're done.
      if ( filteredSampCount == 0 ) {
        plane.addPixel(outputPixel);
        continue;
      }

      // Reweight the resulting 'main' input pixels and put them into the output.

      DeepPixel filteredResultPixel = filteredResult.getPixel( chanMap, DeepPixel::eZAscending );

      DeepReTarget::ReTarget(
                             filteredSampCount,
                             0, // no recoloring
                             inputAlpha,
                             targetAlpha,
                             vizs,
                             filteredResultPixel,
                             outputPixel,
                             channelsRequired,
                             true // drop zero alpha samples
                             );

      // and done.

      plane.addPixel( outputPixel );
    } // box iterator

    return done;
  }


  static const Description d;
  static const Description d2;
};

static Op* build(Node* node) { return new DeepMerge(node); }

const Op::Description DeepMerge::d(::CLASS, "Image/DeepMerge", build);
const Op::Description DeepMerge::d2("ToDeep", "Image/DeepMerge", build);