#include "wavelet.h"

wavelet::wavelet()
{
  data = NULL;
}

wavelet::~wavelet()
{
  del();
}

int wavelet::del()
{

  if(data != NULL)
    del(data);
  levelptr.clear();
  data = NULL;
  return 0;
}

int wavelet::del( wavelettree *d )
{
  if( d->hl != NULL )
    delete d->hl;
  if( d->lh != NULL )
    delete d->lh;
  if( d->hh != NULL )
    delete d->hh;
  if( d->next != NULL )
    del(d->next);
  delete d;
  return 0;
}

int wavelet::load( wavelettree *d )
{
  del();
  data = d;
  wavelettree *t;
  t = d;
  int lc;
  lc = 0;
  while(t != NULL)
    {
      lc++;
      t = t->next;
    }
  levelptr.resize(lc);
  t = data;
  while( t != NULL )
    {
      lc--;
      levelptr[lc] = t;
      t = t->next;
    }
  return 0;
}

int wavelet::analyze( image2d &im, int type )
{
  del();
  realszx = im.getheight();
  realszy = im.getwidth();
  image2d myim;
  myim.load( &im );
  if( realszx % 8 != 0 || realszy % 8 != 0 )
    {
      //padd image
      int nx, ny;
      nx = (im.getheight()/8+1)*8;
      ny = (im.getwidth()/8+1)*8;
      myim.padzeros( nx, ny );
    }
  if(type == 1)
    {
#ifdef __DEBUG__
      std::cout << "using integer to integer transform\n";
#endif
      iti = 1;
    }
  else
    {
#ifdef __DEBUG__
      std::cout << "using floating point transform\n";
#endif
      iti = 0;
    }
  type = 0;
  strans( myim );
  return 0;
}

int wavelet::istrans( image2d &im, int type)
{
  wavelettree *wlt;
  wlt = data->next;
  image2d ll1, ll2;
  image2d *ll, *nll, *t2d;
  ll = &ll1;
  nll = &ll2;
  if(type == 1)
    {
#ifdef __DEBUG__
      std::cout << "using integer to integer itransform\n";
#endif
      iti = 1;
    }
  else
    {
#ifdef __DEBUG__
      std::cout << "using floating point itransform\n";
#endif
      iti = 0;
    }
  ll->load( data->hl );
  while( wlt != NULL )
    {
#ifdef __DEBUG__
      std::cout << "starting level reconstruction ";
      std::cout << (long unsigned int)wlt;
      std::cout << "\n";
#endif

      istrans( ll, wlt->hl, wlt->lh,
               wlt->hh, nll);
#ifdef __DEBUG__
      std::cout << " size: " << nll->getheight() << "\n";
#endif
      t2d = ll;
      ll = nll;
      nll = t2d;
      wlt = wlt->next;
    }
  im.load( ll );
  if( realszx != im.getheight() || realszy != im.getwidth() )
    {
      im.crop( realszx, realszy );
    }
  return 0;
}

int wavelet::calcgradvector( int level, image2d &im )
{
  wavelettree *wl = levelptr[level];

  im.load( wl->hl );
  for( int i = 0; i < im.getheight(); i++ )
    {
      for( int j = 0; j < im.getwidth(); j++ )
        {
          double a, b;
          a = wl->hl->getpixel( i, j );
          b = wl->lh->getpixel( i, j );
          im.setpixel( i, j, atan2( a, b ) );
        }
    }

  return 0;
}

int wavelet::calcgradmag( int level, image2d &im )
{
  wavelettree *wl = levelptr[level];

  im.load( wl->hl );
  for( int i = 0; i < im.getheight(); i++ )
    {
      for( int j = 0; j < im.getwidth(); j++ )
        {
          double a, b;
          a = wl->hl->getpixel( i, j );
          b = wl->lh->getpixel( i, j );
          im.setpixel( i, j, sqrt(pow(a, 2) + pow(b, 2)) );
        }
    }

  return 0;
}

int wavelet::strans( image2d &im )
{
  /// create two image for buffer swapping
  image2d tim, tim1;
#ifdef __DEBUG__
  image2d check;
#endif
  image2d *timptr, *timptr1, *timptr2;
  wavelettree *wlt;

  tim.load( &im );
  timptr = &tim;  
  timptr1 = &tim1;

  while( timptr->getheight() % 2 == 0 &&
         timptr->getwidth() % 2 == 0 && 
         timptr->getheight() > 64 &&
         timptr->getwidth() > 64)
    {
      wlt = new wavelettree;
      wlt->hl = new image2d;
      wlt->hh = new image2d;
      wlt->lh = new image2d;

      strans( timptr, timptr1, wlt->hl,
              wlt->lh, wlt->hh );
#ifdef __DEBUG__
      std::cout << "level size: " 
                << wlt->hl->getheight() << " "
                << (long unsigned int) wlt << "\n";

      istrans( timptr1, wlt->hl, wlt->lh,
               wlt->hh, &check );
#endif
      // swap the buffers
      timptr2 = timptr;
      timptr = timptr1;
      timptr1 = timptr2;

      // link the node to the wavelet tree
      wlt->next = data;
      data = wlt;
      levelptr.push_back( data );
    }
#ifdef __DEBUG__
  std::cout << "all levels created\n";
  std::cout << "final lowpass size: " 
            << timptr->getheight() << " " 
            << timptr->getwidth() << "\n";
#endif
  // create the final low pass data
  wlt = new wavelettree;
  wlt->hl = new image2d;
  wlt->hl->load( timptr );
  wlt->hh = wlt->lh = NULL;
  wlt->next = data;
  data = wlt;
  levelptr.push_back( data );
#ifdef __DEBUG__
  std::cout << "copy done\n";
#endif
  return 0;
}

int wavelet::strans( image2d *in, image2d *ll,
                     image2d *hl, image2d *lh,
                     image2d *hh)
{
  int origh, origw;  // original width and heigh
  origh = in->getheight();
  origw = in->getwidth();

  // delete any previous stuff
  ll->clear();
  hl->clear();
  lh->clear();
  hh->clear();

  ll->create( origh/2, origw/2);
  hl->create( origh/2, origw/2);
  lh->create( origh/2, origw/2);
  hh->create( origh/2, origw/2);

  double *buffer, *sbuf, *sbufo;
  buffer = new double[origh*origw];
  sbuf = new double[ origh ];
  sbufo = new double[ origh ];

  // filter in x direction
  for(int i = 0; i < origw; i++)
    {
      for( int j = 0; j < origh; j++)
        {
          sbuf[j] = in->getpixel(j, i);
          if( isnan(sbuf[j]) != 0 )
            {
              std::cout << "ERROR not a number strans12\n";
              exit(0);
            }
        }

      strans(sbuf, sbufo, origh);

      for( int j = 0; j < origh; j++)
        {
          if( isnan( sbufo[j] ) != 0 )
            {
              std::cout << "wavelet is nan error1\n";
              std::cout << i << " " << j << " " << sbufo[j] << "\n";
              exit(0);
            }

          buffer[j*origw+i] = sbufo[j];
        }
    }
  delete [] sbuf;
  delete [] sbufo;

  // filter in y direction

  sbufo = new double[origw];
  for(int i = 0; i < origh; i++)
    {
      strans( buffer+(i*origw), sbufo, origw );
      for(int j = 0; j < origw; j++)
        buffer[i*origw+j] = sbufo[j];
      for(int j = 0; j < origw; j++)
        {
          double p;
          p = sbufo[j];
          if( isnan(p) != 0 )
            {
              std::cout << "wavelet is nan error2\n";
              std::cout << i << " " << j << " "<< p << "\n";
              std::cout << origh << " " << origw << "\n";
              exit(0);
            }

          if( i < origh/2 && j < origw/2 )
            ll->setpixel(i, j, p);
          else if( i < origh/2 && j >= origw/2)
            lh->setpixel(i, j-(origw/2), p);
          else if( i >= origh/2 && j >= origw/2)
            hh->setpixel(i-(origh/2), j-(origw/2), p);
          else if( i >= origh/2 && j < origw/2 )
            hl->setpixel(i-(origh/2), j, p);
          else
            {
              std::cout << "EROR strans\n";
              exit(0);
            }
        }
    }

#ifdef __DEBUG__
  double *obuffer;
  obuffer = new double[origw*origh];
  istrans(buffer, obuffer, origw, origh);

  for(int i = 0; i < origh; i++)
    {
      for(int j = 0; j < origw; j++)
        {
          if(obuffer[i*origw+j] != in->getpixel(i,j) )
            {
              std::cout << "ERROR ERROR ERROR strans check\n";
            }
        }
    }
  delete [] obuffer;
#endif

  delete [] buffer;
  delete [] sbufo;

#ifdef __DEBUG__
  std::cout << "wavelet done\n";
#endif

  return 0;
}

int wavelet::istrans( double *in, double *out, int h, int w)
{
  double *buf, *ibuf, *obuf;
  buf = new double[h*w];
  // inverse transform in y dir
  ibuf = new double[w];
  obuf = new double[w];
  for(int i = 0; i < h; i++)
    istrans( in+(i*w), in+(i*w+w/2), buf+(i*w), w/2 );

  // inverse transform in x dir
  delete [] ibuf;
  delete [] obuf;
  ibuf = new double[h];
  obuf = new double[h];
  for(int i = 0; i < w; i++)
    {
      for(int j = 0; j < h; j++)
        ibuf[j] = buf[j*w+i];
      istrans( ibuf, ibuf+(h/2), obuf, h/2 );
      for(int j = 0; j < h; j++)
        out[j*w+i] = obuf[j];
    }
  delete [] ibuf;
  delete [] obuf;
  delete [] buf;
  return 0;
}

int wavelet::istrans( image2d *ll, image2d *hl,
                      image2d *lh, image2d *hh,
                      image2d *out)
{
  int origh, origw;
  origh = ll->getheight();
  origw = ll->getwidth();
  out->create( origh*2, origw*2 );
  double *buffer;
  buffer = new double[origh*2*origw*2];
  for(int i = 0; i < origh*2; i++)
    {
      for(int j = 0; j < origw*2; j++)
        {
          double p;
          if( i < origh && j < origw )
            p = ll->getpixel(i, j);
          else if( i < origh && j >= origw)
            p = lh->getpixel(i, j-(origw));
          else if( i >= origh && j >= origw)
            p = hh->getpixel(i-(origh), j-(origw));
          else if( i >= origh && j < origw )
            p = hl->getpixel(i-(origh), j);
          else
            {
              std::cout << "EROR istrans\n";
              exit(0);
            }
          buffer[i*(origw*2)+j] = p;
        }
    }
  istrans(buffer, out->getimageptr(), origh*2, origw*2);
  delete [] buffer;
  return 0;
}

int wavelet::istrans( double *l, double *h,
                      double *out, int size )
{
  for(int i = 0; i < size; i++)
    {
      if(iti == 1)
        {
          out[2*i] = l[i] - floor(h[i]/2.0);
          double ip, fp;
          fp = modf( l[i], &ip );
          if( fp > 0.000001 )
            std::cout << "ERROR istrans " << fp << "\n";
          fp = modf( h[i], &ip );
          if( fp > 0.000001 )
            std::cout << "ERROR istrans " << fp << "\n";
          fp = modf( out[2*i], &ip );
          if( fp > 0.000001 )
            std::cout << "ERROR istrans " << fp << "\n";
          out[2*i] = round( out[2*i] );
        }
      else
        out[2*i] = l[i] - h[i]/2.0;
      out[2*i+1] = round( h[i] + out[2*i] );
      if( iti == 1 )
        {
          double ip, fp;
          fp = modf( out[2*i+1], &ip );
          if( fp > 0.000001 )
            std::cout << "ERROR istrans " << fp << "\n";
        }
    }
  return 0;
}

int wavelet::strans( double *in, double *out, int size)
{
  // calculate the detail coefficients
  double *d;
  d = out+(size/2);
  for(int i = 0; i < size/2; i++)
    d[i] = round( in[2*i+1] - in[2*i] );

  // calculate the low pass
  if( iti == 1)
    {
      for(int i = 0; i < size/2; i++)
        {
          out[i] = in[2*i] +floor( d[i]/2.0);
          double ip, fp;
          fp = modf( d[i], &ip );
          if( fp > 0.000001 )
            {
              std::cout << "ERROR strans input " << fp << "\n";
              for( int i = 0; i < size; i++ )
                std::cout << in[i] << " ";
              std::cout << "\n";
              exit(0);
            }
          fp = modf( in[2*i], &ip );
          if( fp > 0.000001 )
            {
              std::cout << "ERROR strans input2 " << fp << "\n";
              for( int i = 0; i < size; i++ )
                std::cout << in[i] << " ";
              std::cout << "\n";
              exit(0);
            }
          fp = modf( out[i], &ip );
          if( fp > 0.000001 )
            {
              std::cout << "ERROR strans output " << fp << "\n";
              exit(0);
            }
          out[i] = round( out[i] );
        }
    }
  else
    {
      for(int i = 0; i < size/2; i++)
        out[i] = in[2*i] + d[i]/2.0;
    }

#ifdef __DEBUG__
  double *tmpbuf;
  tmpbuf = new double[size];
  istrans( out, d, tmpbuf, size/2 );
  for(int i = 0; i < size; i++)
    if( in[i] != tmpbuf[i])
      std::cout << "ERROR: " << in[i] << " " << tmpbuf[i] << "\n";
  delete [] tmpbuf;
#endif

  return 0;
}

int wavelet::gethplevel( image2d &im, int level )
{
  im.create( levelptr[0]->hl->getheight()*2,
             levelptr[0]->hl->getwidth()*2 );
  for( int i = 0; i < im.getheight(); i = i + 1)
    {
      for( int j = 0; j < im.getwidth(); j = j + 1)
        {
          double nv;
          double di;
          double dj;
          
          di = ((double)i)/( (double) im.getheight());
          dj = ((double)j)/( (double) im.getwidth());
          nv = 0;
          nv = nv + fabs(levelptr[level]->hl->getpixeld( di, dj ));
          
          nv = nv + fabs(levelptr[level]->lh->getpixeld( di, dj ));
          nv = nv + fabs(levelptr[level]->hh->getpixeld( di, dj ));
          
          im.setpixel( i, j, nv);
        }
    }
  return 0;
}

int wavelet::gethpmap( image2d &im )
{
  im.create( levelptr[0]->hl->getheight() * 2,
             levelptr[0]->hl->getwidth() * 2 );
#ifdef __DEBUG__
  std::cout << "generating high pass map\n";
#endif
  for(unsigned int l = 1; l < 4; l++ )
    {
      for( int i = 0; i < im.getheight(); i = i + 1)
        {
          for( int j = 0; j < im.getwidth(); j = j + 1)
            {
              double nv;
              double di;
              double dj;

              di = ((double)i)/( (double) im.getheight());
              dj = ((double)j)/( (double) im.getwidth());
              nv =  im.getpixel( i, j );
              if( isnan(nv) != 0 )
                {
                  std::cout << "ERROR level " << l << "\n";
                  exit(0);
                }
              nv = nv + fabs(levelptr[l]->hl->getpixeld( di, dj ));
              if( isnan(nv) != 0 )
                {
                  std::cout << "ERROR level " << l << " hl\n";
                  exit(0);
                }
              nv = nv + fabs(levelptr[l]->lh->getpixeld( di, dj ));
              if( isnan(nv) != 0 )
                {
                  std::cout << "ERROR level " << l << " lh\n";
                  exit(0);
                }
              nv = nv + fabs(levelptr[l]->hh->getpixeld( di, dj ));
              if( isnan(nv) != 0 )
                {
                  std::cout << "ERROR level " << l << " hh\n";
                  exit(0);
                }
              im.setpixel( i, j, nv);
            }
        }
    }

  return 0;
}
