/*
 * Implementation of class ImageDenoiser
 *
 * $Date$
 * $Revision$
 *
 */

#include "WTools/ImageDenoiser.hh"
#include "Wave/PyramidTransform.hh"
#ifdef _WIN32_WCE
#undef DELETE
#endif
#include <cassert>
#include <stdexcept>
#include <cmath>
#include "WImage/miscdefs.h"

#if !(defined __WATCOMC__ && __WATCOMC__ < 1230)
using namespace std;
#endif


ImageDenoiser::ImageDenoiser (Image &img, unsigned areas, int fromSub, 
                              int toSub, double alpha, FilterSet &filter,
                              int isSingleSignificant, int replacementFunction)
{
  if (areas == 0)
    {
      throw invalid_argument ("ImageDenoiser: areas cannot be zero.");
    }
  if (fromSub < 1 || toSub < 1)
    {
      throw invalid_argument ("ImageDenoiser: invalid subbands.");
    }
  if (fromSub <= toSub)
    {
      throw invalid_argument ("ImageDenoiser: less than two subbands.");
    }
  if (isSingleSignificant < 0 || isSingleSignificant > 2
      || replacementFunction < 0 || replacementFunction > 1)
    {
      throw invalid_argument ("ImageDenoiser: invalid function.");
    }
  
  m_image = &img;
  m_areas = areas;
  m_from = fromSub;
  m_to = toSub;
  m_alpha = alpha;
  m_filter = &filter;
  m_significanceFunction = isSingleSignificant;
  m_replacementFunction = replacementFunction;
  m_transform = NEW (PyramidTransform (*m_image, *m_filter));
}

ImageDenoiser::~ImageDenoiser (void)
{
  DELETE (m_transform);
}

coeff 
ImageDenoiser::threshold (int subband, area channel)
{
  coeff avg = m_transform->aaverage (subband, channel);
  coeff sdev = m_transform->sdeviation (subband, channel, true);
  return fabs (avg - sdev) * m_alpha;
}

void 
ImageDenoiser::denoise (void)
{
  m_transform->analysis (m_from);

  static unsigned areaBits [] = { 0x00, DENOISE_HL, DENOISE_LH, DENOISE_HH };
  int yoffs, xoffs, ysize, xsize;
  int depth = m_from - m_to + 1;
  assert (depth > 0);
  coeff *thresholds = NEW (coeff [depth]);
  PyramidTree *tree = NULL;

  for (int i = 1; i < (int)areaINVALID; i++)
    {
      if (m_areas & areaBits[i])
        {
          for (int j = 0; j < depth; j++)
            {
              thresholds[j] = threshold (j + m_to, areas[i]);
              //cout << "j: " << j << ", subband: " << j + m_to 
              //     << ", threshold: " << thresholds[j] << endl;
            }
          m_transform->where (areas[i], m_from, 
                              yoffs, xoffs, ysize, xsize);

          tree = NEW (PyramidTree (*m_transform, yoffs, xoffs, -1, NULL, depth));
          
          for (int y = yoffs; y < yoffs + ysize; y++)
            {
              for (int x = xoffs; x < xoffs + xsize; x++)
                {
                  tree->moveTo (y, x);
                  processFrom (*tree, thresholds, areas[i], m_from);
                }
            }
          DELETE (tree);
        }
    }

  DELETE (thresholds);
  m_transform->synthesis (m_from, m_from);
}

void 
ImageDenoiser::processFrom (PyramidTree &tree, coeff *thresholds, 
                            area channel, int level)
{
  coeff threshold = thresholds[level - m_to];

  bool isSingleSignificant = false;
  bool stop = false;
  switch (m_significanceFunction)
    {
    case SIGNIFICANT_COEFF:
      isSingleSignificant = isCoeffSingleSignificant (tree.data (), channel, 
                                                      threshold, stop);
      break;
    case SIGNIFICANT_REGION:
      isSingleSignificant = isSingleSignificantInRegion (tree.data (), channel,
                                                         threshold, stop);
      break;
    case SIGNIFICANT_CHANNEL:
      isSingleSignificant = isSingleSignificantInChannel (tree.data (), 
                                                          channel, threshold, 
                                                          stop);
      break;
    default: 
      // should have been caught by the constructor!
      assert (0);
    }
  
  if (isSingleSignificant && level != m_from)
    {
      //cout << " (not root)" << endl;
      // @todo adaptive threshold?
      //cout << "changing at: " 
      //     << (tree.data ()).ypos () << ", " 
      //     << (tree.data ()).xpos () << endl;
      switch (m_replacementFunction)
        {
        case REPLACE_SIMPLE:
          replaceSimple ((tree.data ()).ypos (), (tree.data ()).xpos (),
                         channel, threshold);
          break;
        case REPLACE_CHANNEL:
          replaceChannel ((tree.data ()).ypos (), (tree.data ()).xpos (),
                          channel, threshold);
          break;
        default: 
          // can't happen.
          assert (0);
        }
    }

  if (level <= m_to || stop)
    {
      return;
    }
  
  // Either insignificant in the root node or anywhere in the leaves
  for (int i = 0; i < tree.aryness (); i++)
    {
      processFrom ((PyramidTree&)tree.childAt (i), thresholds, 
                   channel, level - 1);
    }
}

bool 
ImageDenoiser::isCoeffSingleSignificant (CoeffInformation &c, area channel, 
                                         coeff threshold, bool &stop)
{
  stop = false;
  return fabs (c.val ()) >= threshold;
}

bool 
ImageDenoiser::isSingleSignificantInRegion (CoeffInformation &c, area channel,
                                            coeff threshold, bool &stop)
{
  stop = false;
  if (fabs (c.val ()) < threshold)
    {
      return false;
    }
  
  int fromY = c.ypos () - 1;
  int fromX = c.xpos () - 1;
  int toY = fromY + 3;
  int toX = fromX + 3;

  fromY = (fromY < 0? 0: fromY);
  fromX = (fromX < 0? 0: fromX);
  toY = (toY >= (m_transform->image ()).rows ()? 
         (m_transform->image ()).rows (): toY);
  toX = (toX >= (m_transform->image ()).cols ()? 
         (m_transform->image ()).cols (): toX);
  
  int insignificant = 0;
  int significant = 0;

  for (int y = fromY; y < toY; y++)
    {
      for (int x = fromX; x < toX; x++)
        {
          if (fabs ((m_transform->image ()).at (y, x)) >= threshold)
            {
              ++significant;
            }
          else
            {
              ++insignificant;
            }
        }
    }
  
  stop = significant >= insignificant;
  return !stop;
}

bool 
ImageDenoiser::isSingleSignificantInChannel (CoeffInformation &c, area channel,
                                             coeff threshold, bool &stop)
{
  stop = false;
  if (fabs (c.val ()) < threshold)
    {
      return false;
    }
  
  int fromY = c.ypos () - 1;
  int fromX = c.xpos () - 1;
  int toY = fromY + 3;
  int toX = fromX + 3;

  fromY = (fromY < 0? 0: fromY);
  fromX = (fromX < 0? 0: fromX);
  toY = (toY >= (m_transform->image ()).rows ()? 
         (m_transform->image ()).rows (): toY);
  toX = (toX >= (m_transform->image ()).cols ()? 
         (m_transform->image ()).cols (): toX);
  
  int insignificant = 0;
  int significant = 0;

  switch (channel)
    {
    case HH:
      // Look at the diagonals:
      for (int y = fromY; y < toY; y++)
        {
          if (y == c.ypos ())
            {
              continue;
            }
          for (int x = fromX; x < toX; x++)
            {
              if (x == c.xpos ())
                {
                  continue;
                }
              if (fabs ((m_transform->image ()).at (y, x)) >= threshold)
                {
                  ++significant;
                }
              else
                {
                  ++insignificant;
                }
            }
        }
      // the center is signficant as we already know!
      significant++;
      break;
    case HL:
      // Look at one column:
      for (int y = fromY; y < toY; y++)
        {
          if (fabs ((m_transform->image ()).at (y, c.xpos ())) >= threshold)
            {
              ++significant;
            }
          else
            {
              ++insignificant;
            }
        }
      break;
    case LH:
      // Look at one row:
      for (int x = fromX; x < toX; x++)
        {
          if (fabs ((m_transform->image ()).at (c.ypos (), x)) >= threshold)
            {
              ++significant;
            }
          else
            {
              ++insignificant;
            }
        }
      break;
      // can't happen
    case LL:
    case areaINVALID:
      assert (0);
    }
  stop = significant >= insignificant;
  return !stop;
}

void 
ImageDenoiser::replaceSimple (int y, int x, area channel, coeff threshold)
{
  //cout << "Replacing " << (m_transform->image ()).at (y, x) << " by ";
  (m_transform->image ()).to (y, x, ((m_transform->image ()).at (y, x) < 0? 
                                     -1: 1) * 0.99 * threshold);
  //cout << (m_transform->image ()).at (y, x) << endl;
}

void 
ImageDenoiser::replaceChannel (int y, int x, area channel, coeff threshold)
{
  int fromY = y - 1;
  int fromX = x - 1;
  int toY = fromY + 3;
  int toX = fromX + 3;

  fromY = (fromY < 0? 0: fromY);
  fromX = (fromX < 0? 0: fromX);
  toY = (toY >= (m_transform->image ()).rows ()? 
         (m_transform->image ()).rows (): toY);
  toX = (toX >= (m_transform->image ()).cols ()? 
         (m_transform->image ()).cols (): toX);
  
  coeff value = 0.0;
  int nvals = 0;

  switch (channel)
    {
    case HH:
      // Look at the diagonals:
      for (int row = fromY; row < toY; row++)
        {
          if (row == y)
            {
              continue;
            }
          for (int col = fromX; col < toX; col++)
            {
              if (col == x)
                {
                  continue;
                }
              value += SQUARE ((m_transform->image ()).at (row, col));
              nvals++;
            }
        }
      // This one also goes into consideration!
      value += SQUARE ((m_transform->image ()).at (y, x));
      nvals++;
      break;
    case HL:
      // Look at one column:
      for (int row = fromY; row < toY; row++)
        {
          value += SQUARE ((m_transform->image ()).at (row, x));
          nvals++;
        }
      break;
    case LH:
      // Look at one row:
      for (int col = fromX; col < toX; col++)
        {
          value += SQUARE ((m_transform->image ()).at (y, col));
          nvals++;
        }
      break;
      // can't happen
    case LL:
    case areaINVALID:
      assert (0);
    }
  value = sqrt (value/nvals);
  (m_transform->image ()).to (y, x, ((m_transform->image ()).at (y, x) < 0? 
                                     -1: 1) * value);
}

