//---Author:  Richard T. Jones  07/12/2007
//---Version:  1.0

/*************************************************************************
 * Copyright(c) 1998, University of Connecticut, All rights reserved.    *
 * Author: Richard T. Jones, Asst. Prof. of Physics                      *
 *                                                                       *
 * Permission to use, copy, modify and distribute this software and its  *
 * documentation for non-commercial purposes is hereby granted without   *
 * fee, provided that the above copyright notice appears in all copies   *
 * and that both the copyright notice and this permission notice appear  *
 * in the supporting documentation. The author makes no claims about the *
 * suitability of this software for any purpose.                         *
 * It is provided "as is" without express or implied warranty.           *
 *************************************************************************/
//////////////////////////////////////////////////////////////////////////
//									//
// Michelson Interferometric Topography Analysis Package		//
//									//
// The present package implements tools for modeling the shape of a     //
// general 2D surface in terms of its Legendre expansion.		//
//									//
//////////////////////////////////////////////////////////////////////////
 
#include <iostream>
#include <fstream>

using namespace std;

#include <TASImage.h>
#include <TArrayD.h>
#include <Math/SpecFuncMathMore.h>
#include <TBitmap.h>

ClassImp(TBitmap)


TBitmap::TBitmap(UInt_t width, UInt_t height)
 : fWidth(width), fHeight(height), fDeg(0)
{
   fImage = new Double_t[width*height];
   fTimage = 0;
   fHasModel = 0;
}

TBitmap::TBitmap(const TBitmap& map)
 : fWidth(map.fWidth), fHeight(map.fHeight), fDeg(0)
{
   fImage = new Double_t[fWidth*fHeight];
   for (UInt_t i=0; i<fWidth*fHeight; i++) {
      fImage[i] = map.fImage[i];
   }
   fTimage = 0;
   fHasModel = 0;
   if (map.fDeg) {
      fDeg = map.fDeg;
      fXbasis = new Double_t[fDeg*fWidth];
      fYbasis = new Double_t[fDeg*fHeight];
      for (UInt_t i=0,k=0; i<fWidth; i++)
         for (UInt_t j=0; j<fDeg; j++,k++)
            fXbasis[k] = map.fXbasis[k];
      for (UInt_t i=0,k=0; i<fHeight; i++)
         for (UInt_t j=0; j<fDeg; j++,k++)
            fYbasis[k] = map.fYbasis[k];
         
   }
   if (map.fHasModel) {
      fHasModel = new TModel(*(map.fHasModel));
   }
}

TBitmap::~TBitmap()
{
   delete [] fImage;
   if (fDeg) {
      delete [] fXbasis;
      delete [] fYbasis;
   }
   if (fTimage)
      delete fTimage;
   if (fHasModel)
      delete fHasModel;
}

TBitmap& TBitmap::operator=(const TBitmap& source)
{
   delete [] fImage;
   if (fTimage) {
      delete fTimage;
      fTimage = 0;
   }
   if (fHasModel) {
      delete fHasModel;
      fHasModel = 0;
   }
   fWidth = source.fWidth;
   fHeight = source.fHeight;
   fImage = new Double_t[fWidth*fHeight];
   for (UInt_t i=0; i<fWidth*fHeight; i++) {
      fImage[i] = source.fImage[i];
   }
   if (source.fDeg)
      Precompute(fDeg);
   if (source.fHasModel)
      fHasModel = new TModel(*(source.fHasModel));
   return *this;
}

TBitmap& TBitmap::operator+=(const TBitmap &source)
{
   if (source.fWidth != fWidth || source.fHeight != fHeight) {
      std::cerr << "Error in TBitmap::operator+= : source and destination maps"
                << " have different dimensions!" << std::endl;
      return *this;
   }

   UInt_t k=0;
   for (UInt_t m=0; m<fHeight; m++) {
      for (UInt_t n=0; n<fWidth; n++,k++) {
         fImage[k] += source.fImage[k];
      }
   }
   if (fDeg)
      Precompute(fDeg);
   if (fTimage) {
      delete fTimage;
      fTimage = 0;
   }
   if (fHasModel) {
      delete fHasModel;
      fHasModel = 0;
   }
   return *this;
}

TBitmap& TBitmap::operator-=(const TBitmap &source)
{
   if (source.fWidth != fWidth || source.fHeight != fHeight) {
      std::cerr << "Error in TBitmap::operator-= : source and destination maps"
                << " have different dimensions!" << std::endl;
      return *this;
   }

   UInt_t k=0;
   for (UInt_t m=0; m<fHeight; m++) {
      for (UInt_t n=0; n<fWidth; n++,k++) {
         fImage[k] -= source.fImage[k];
      }
   }
   if (fDeg)
      Precompute(fDeg);
   if (fTimage) {
      delete fTimage;
      fTimage = 0;
   }
   if (fHasModel) {
      delete fHasModel;
      fHasModel = 0;
   }
   return *this;
}

TBitmap& TBitmap::operator*=(const TBitmap &source)
{
   if (source.fWidth != fWidth || source.fHeight != fHeight) {
      std::cerr << "Error in TBitmap::operator*= : source and destination maps"
                << " have different dimensions!" << std::endl;
      return *this;
   }

   UInt_t k=0;
   for (UInt_t m=0; m<fHeight; m++) {
      for (UInt_t n=0; n<fWidth; n++,k++) {
         fImage[k] *= source.fImage[k];
      }
   }
   if (fDeg)
      Precompute(fDeg);
   if (fTimage) {
      delete fTimage;
      fTimage = 0;
   }
   if (fHasModel) {
      delete fHasModel;
      fHasModel = 0;
   }
   return *this;
}

TBitmap& TBitmap::operator/=(const TBitmap &source)
{
   if (source.fWidth != fWidth || source.fHeight != fHeight) {
      std::cerr << "Error in TBitmap::operator/= : source and destination maps"
                << " have different dimensions!" << std::endl;
      return *this;
   }

   UInt_t k=0;
   for (UInt_t m=0; m<fHeight; m++) {
      for (UInt_t n=0; n<fWidth; n++,k++) {
         fImage[k] /= (source.fImage[k] + 1e-30);
      }
   }
   if (fDeg)
      Precompute(fDeg);
   if (fTimage) {
      delete fTimage;
      fTimage = 0;
   }
   if (fHasModel) {
      delete fHasModel;
      fHasModel = 0;
   }
   return *this;
}

TBitmap& TBitmap::operator*=(Double_t factor)
{
   UInt_t k=0;
   for (UInt_t m=0; m<fHeight; m++) {
      for (UInt_t n=0; n<fWidth; n++,k++) {
         fImage[k] *= factor;
      }
   }
   if (fDeg)
      Precompute(fDeg);
   if (fTimage) {
      delete fTimage;
      fTimage = 0;
   }
   if (fHasModel) {
      delete fHasModel;
      fHasModel = 0;
   }
   return *this;
}

TBitmap& TBitmap::operator/=(Double_t factor)
{
   UInt_t k=0;
   for (UInt_t m=0; m<fHeight; m++) {
      for (UInt_t n=0; n<fWidth; n++,k++) {
         fImage[k] /= (factor+1e-30);
      }
   }
   if (fDeg)
      Precompute(fDeg);
   if (fTimage) {
      delete fTimage;
      fTimage = 0;
   }
   if (fHasModel) {
      delete fHasModel;
      fHasModel = 0;
   }
   return *this;
}

TBitmap TBitmap::operator-() const
{
   TBitmap map(*this);
   UInt_t k=0;
   for (UInt_t m=0; m<fHeight; m++) {
      for (UInt_t n=0; n<fWidth; n++,k++) {
         fImage[k] = -fImage[k];
      }
   }
   if (fDeg)
      map.Precompute(fDeg);
   return map;
}

Bool_t TBitmap::operator==(const TBitmap &other) const
{
   if (other.fWidth != fWidth || other.fHeight != fHeight) {
      return kFALSE;
   }

   UInt_t k=0;
   for (UInt_t m=0; m<fHeight; m++) {
      for (UInt_t n=0; n<fWidth; n++,k++) {
         if (fImage[k] != other.fImage[k]) {
            return kFALSE;
         }
      }
   }
   return kTRUE;
}

void TBitmap::ComputeFast(const TModel& model)
{
   // A faster version of Compute()
   // Assumes Precompute(deg) is called first with a degree deg at least
   // as large as the degree of model, otherwise the following
   // implementation truncates model at degree fDeg.

   // std::cout << "+" << std::endl;
   if (fTimage) {
      delete fTimage;
      fTimage = 0;
   }

   // If fHasModel is set, we can compute changes since the last call
   // instead of computing the entire sum.  This is more efficient but
   // can result in accummulated errors.  We check this later on below.

   UInt_t deg=model.GetDegree();
   if (fHasModel && fHasModel->GetDegree() != deg) {
      delete fHasModel;
      fHasModel = 0;
   }

   if (fHasModel) {
      ++fHasModelCtr;
      Double_t coeff[deg*deg];
      UInt_t ci[deg*deg];
      UInt_t cj[deg*deg];
      UInt_t nchanged=0;
      for (UInt_t i=0; i<deg; i++) {
         for (UInt_t j=0; j<deg; j++) {
            ci[nchanged] = i;
            cj[nchanged] = j;
            coeff[nchanged] = model[j][i]-(*fHasModel)[j][i];
            if (coeff[nchanged] != 0)
               ++nchanged;
         }
      }
      UInt_t km=0;
      for (UInt_t m=0; m<fHeight; m++, km+=fDeg) {
         UInt_t kn=0;
         for (UInt_t n=0; n<fWidth; n++, kn+=fDeg) {
            UInt_t pix = m*fWidth+n;
            Double_t f = fImage[pix];
            for (UInt_t c=0; c<nchanged; c++) {
               UInt_t ki = km+ci[c];
               UInt_t kj = kn+cj[c];
               f += coeff[c]*fXbasis[kj]*fYbasis[ki];
            }
            fImage[pix] = f;
         }
      }
   }

   else {
      UInt_t km=0;
      for (UInt_t m=0; m<fHeight; m++, km+=fDeg) {
         UInt_t kn=0;
         for (UInt_t n=0; n<fWidth; n++, kn+=fDeg) {
            Double_t f=0;
            UInt_t ki=km;
            UInt_t kj=kn;
            for (UInt_t i=0; i<fDeg; i++,ki++,kj=kn) {
               for (UInt_t j=0; j<fDeg; j++,kj++) {
                  if (i < deg && j < deg)
                     f += model[j][i]*fXbasis[kj]*fYbasis[ki];
               }
            }
            fImage[m*fWidth+n] = f;
         }
      }
   }

   if (fHasModel && fHasModelCtr > 1000) {
      UInt_t k=0;
      for (UInt_t m=0; m<fHeight; m++) {
         Double_t y=(2.0*m-fHeight+1)/(fHeight-1.0);
         for (UInt_t n=0; n<fWidth; n++,k++) {
            Double_t x=(2.0*n-fWidth+1)/(fWidth-1.0);
            Double_t expected = model(x,y);
            if (fabs(fImage[k]-expected) > 1e-10) {
               std::cerr << "Warning in TBitmap::ComputeFast - "
                         << "algorithm finds " << fImage[k]
                         << ", expected " << expected
                         << ", difference = " << fImage[k]-expected
                         << std::endl;
            }
            fImage[k] = expected;
         }
      }
      fHasModelCtr = 0;
   }

   if (fHasModel)
      delete fHasModel;
   fHasModel = new TModel(model);
}

void TBitmap::Compute(const TModel& model)
{
   // Compute the value of the model at the lattice sites
   // of the bitmap and store the results in fImage.

   if (fDeg) {
      ComputeFast(model);
      return;
   }

   std::cout << "+c+" << std::endl;
   if (fTimage) {
      delete fTimage;
      fTimage = 0;
   }
   if (fHasModel) {
      delete fHasModel;
      fHasModel = 0;
   }
   UInt_t k=0;
   for (UInt_t m=0; m<fHeight; m++) {
      Double_t y=(2.0*m-fHeight+1)/(fHeight-1.0);
      for (UInt_t n=0; n<fWidth; n++,k++) {
         Double_t x=(2.0*n-fWidth+1)/(fWidth-1.0);
         fImage[k] = model(x,y);
      }
   }
   std::cout << "-c-" << std::endl;
}

void TBitmap::Precompute(UInt_t degree)
{
   // Precompute the Legendre basis functions at the center of each pixel

   if (fDeg) {
      delete [] fXbasis;
      delete [] fYbasis;
   }
   if (degree == 0) {
      fDeg = 0;
      return;
   }
   fXbasis = new Double_t[fWidth*degree];
   fYbasis = new Double_t[fHeight*degree];
   fDeg = degree;
   UInt_t k=0;
   for (UInt_t n=0; n<fWidth; n++) {
      Double_t x=(2.0*n-fWidth+1)/(fWidth-1.0);
      for (UInt_t m=0; m<degree; m++,k++) {
         fXbasis[k] = ROOT::Math::legendre(m,x);
      }
   }
   k=0;
   for (UInt_t n=0; n<fHeight; n++) {
      Double_t y=(2.0*n-fHeight+1)/(fHeight-1.0);
      for (UInt_t m=0; m<degree; m++,k++) {
         fYbasis[k] = ROOT::Math::legendre(m,y);
      }
   }
}

void TBitmap::Draw(Option_t* option)
{
   if (fTimage == 0) {
      fTimage = new TASImage(fWidth,fHeight);
      fTimage->SetImage(fImage,fWidth,fHeight);
   }
   fTimage->Gray();
   fTimage->Draw(option);
}

void TBitmap::ReadBMP(const char* filename)
{
#include <BMPformat.h>
   struct bmp_header_ bmp_header;
   struct dib_header_ dib_header;
   std::ifstream infile(filename);
   if (! infile.is_open()) {
      std::cout << "ReadBMP error - file %s cannot be opened for reading.\n"
                << std::endl;
      return;
   }
   infile.read((char*)&bmp_header,sizeof(struct bmp_header_));
   if (bmp_header.magic != 0x4d42) {
      std::cout << "ReadBMP error - file %s does not have a valid bmp file header.\n"
                << std::endl;
      return;
   }
   infile.read((char*)&dib_header,sizeof(struct dib_header_));
   if (dib_header.hdrlen != 40) {
      std::cout << "ReadBMP error - file %s does not have a recognized bmp file header.\n"
                << std::endl;
      return;
   }
   unsigned int *values = new unsigned int[dib_header.width*dib_header.height];
   infile.read((char*)values,dib_header.size);
   infile.close();

   delete [] fImage;
   fWidth = dib_header.width;
   fHeight = dib_header.height;
   fImage = new Double_t[fWidth*fHeight];
   for (UInt_t i=0; i<fWidth*fHeight; i++) {
      unsigned char *rgba = (unsigned char*)&values[i];
      fImage[i] = rgba[0]+rgba[1]+rgba[2];
   }
   if (fTimage) {
      delete fTimage;
      fTimage = 0;
   }
   if (fHasModel) {
      delete fHasModel;
      fHasModel = 0;
   }
   delete [] values;
}

void TBitmap::WriteBMP(const char* filename)
{
#include <BMPformat.h>
   struct bmp_header_ bmp_header;
   struct dib_header_ dib_header;
   std::ofstream outfile(filename);
   if (! outfile.is_open()) {
      std::cout << "WriteBMP error - file %s cannot be opened for writing.\n"
                << std::endl;
      return;
   }
   bmp_header.magic = 0x4d42;
   bmp_header.fsize = sizeof(struct bmp_header_) + 
                      sizeof(struct dib_header_) +
                      fWidth * fHeight * 4;
   bmp_header.reserved[0] = 0;
   bmp_header.reserved[1] = 0;
   bmp_header.offset = 54;
   outfile.write((char*)&bmp_header,sizeof(struct bmp_header_));
   dib_header.hdrlen = 40;
   dib_header.width = fWidth;
   dib_header.height = fHeight;
   dib_header.planes = 1;
   dib_header.depth = 32;
   dib_header.compression = 0;
   dib_header.size = fWidth * fHeight * 4;
   dib_header.hresol = 0;
   dib_header.vresol = 0;
   dib_header.ncolors = 0;
   dib_header.nimpcolors = 0;
   outfile.write((char*)&dib_header,sizeof(struct dib_header_));

   unsigned int *values = new unsigned int[fWidth*fHeight];
   double maxval=-1e30;
   double minval=1e30;
   for (unsigned int i=0; i<fWidth*fHeight; i++) {
      maxval = (fImage[i] > maxval)? fImage[i] : maxval;
      minval = (fImage[i] < minval)? fImage[i] : minval;
   }
   double slope = 3*255/(maxval-minval+1e-30);
   for (unsigned int i=0; i<fWidth*fHeight; i++) {
      int value = (int)((fImage[i] - minval) * slope);
      unsigned char *rgba = (unsigned char*)&values[i];
      rgba[0] = value/3;
      rgba[1] = (value-rgba[0])/2;
      rgba[2] = value-rgba[0]-rgba[1];
      rgba[3] = 0;
   }
   outfile.write((char*)values,dib_header.size);
   outfile.close();
   delete [] values;
}

void TBitmap::InterfereFast(const TModel& refa,
                            const TModel& sig1a,
                            const TModel& sig2a,
                            const TModel& sig1ph,
                            const TModel& sig2ph)
{
   // A faster version of Interfere()
   // Assumes Precompute(deg) is called first with a degree deg at least
   // as large as the highest-order model in the argument list, otherwise
   // the following implementation truncates all models at degree fDeg.

   // std::cout << "+" << std::endl;
   if (fTimage) {
      delete fTimage;
      fTimage = 0;
   }
   if (fHasModel) {
      delete fHasModel;
      fHasModel = 0;
   }
   UInt_t drefa=refa.GetDegree();
   UInt_t dsig1a=sig1a.GetDegree();
   UInt_t dsig2a=sig2a.GetDegree();
   UInt_t dsig1ph=sig1ph.GetDegree();
   UInt_t dsig2ph=sig2ph.GetDegree();
   UInt_t km=0;
   for (UInt_t m=0; m<fHeight; m++, km+=fDeg) {
      UInt_t kn=0;
      for (UInt_t n=0; n<fWidth; n++, kn+=fDeg) {
         Double_t rrefa=0;
         Double_t rsig1a=0;
         Double_t rsig2a=0;
         Double_t rsig1ph=0;
         Double_t rsig2ph=0;
         UInt_t ki=km;
         UInt_t kj=kn;
         for (UInt_t i=0; i<fDeg; i++,ki++,kj=kn) {
            for (UInt_t j=0; j<fDeg; j++,kj++) {
               Double_t basis = fXbasis[kj]*fYbasis[ki];
               if (i < drefa && j < drefa)
                  rrefa += refa[j][i]*basis;
               if (i < dsig1a && j < dsig1a)
                  rsig1a += sig1a[j][i]*basis;
               if (i < dsig2a && j < dsig2a)
                  rsig2a += sig2a[j][i]*basis;
               if (i < dsig1ph && j < dsig1ph)
                  rsig1ph += sig1ph[j][i]*basis;
               if (i < dsig2ph && j < dsig2ph)
                  rsig2ph += sig2ph[j][i]*basis;
            }
         }
         fImage[m*fWidth+n] = rrefa*rrefa
                            + rsig1a*rsig1a
                            + rsig2a*rsig2a
                            + 2*rrefa*rsig1a*cos(rsig1ph)
                            + 2*rrefa*rsig2a*cos(rsig2ph)
                            + 2*rsig1a*rsig2a*cos(rsig1ph-rsig2ph);
      }
   }
   //std::cout << "=" << std::endl;
}

void TBitmap::Interfere(const TModel& refa,
                        const TModel& sig1a,
                        const TModel& sig2a,
                        const TModel& sig1ph,
                        const TModel& sig2ph)
{
   // Form a Michelson Interferometer image using
   //     refa: amplitude map of reference beam
   //     sig1a: amplitude map of signal beam 1
   //     sig2a: amplitude map of signal beam 2
   //     sig1ph: phase map of signal beam 1 (relative to ref beam)
   //     sig2ph: phase map of signal beam 2 (relative to ref beam)

   if (fDeg) {
      InterfereFast(refa,sig1a,sig2a,sig1ph,sig2ph);
      return;
   }

   std::cout << "++" << std::endl;
   if (fTimage) {
      delete fTimage;
      fTimage = 0;
   }
   if (fHasModel) {
      delete fHasModel;
      fHasModel = 0;
   }
   for (UInt_t m=0; m<fHeight; m++) {
      Double_t y=(2.0*m-fHeight+1)/(fHeight-1.0);
      for (UInt_t n=0; n<fWidth; n++) {
         Double_t x=(2.0*n-fWidth+1)/(fWidth-1.0);
         Double_t drefa=refa(x,y);
         Double_t dsig1a=sig1a(x,y);
         Double_t dsig2a=sig2a(x,y);
         Double_t dsig1ph=sig1ph(x,y);
         Double_t dsig2ph=sig2ph(x,y);
         fImage[m*fWidth+n] = drefa*drefa
                            + dsig1a*dsig1a
                            + dsig2a*dsig2a
                            + 2*drefa*dsig1a*cos(dsig1ph)
                            + 2*drefa*dsig2a*cos(dsig2ph)
                            + 2*dsig1a*dsig2a*cos(dsig1ph-dsig2ph);
      }
   }
   std::cout << "==" << std::endl;
}

void TBitmap::Interfere(TBitmap& refa,
                        TBitmap& sig1a,
                        TBitmap& sig2a,
                        TBitmap& sig1ph,
                        TBitmap& sig2ph)
{
   // Form a Michelson Interferometer image using
   //     refa: amplitude map of reference beam
   //     sig1a: amplitude map of signal beam 1
   //     sig2a: amplitude map of signal beam 2
   //     sig1ph: phase map of signal beam 1 (relative to ref beam)
   //     sig2ph: phase map of signal beam 2 (relative to ref beam)

   if (fTimage) {
      delete fTimage;
      fTimage = 0;
   }
   if (fHasModel) {
      delete fHasModel;
      fHasModel = 0;
   }
   for (UInt_t m=0; m<fHeight; m++) {
      for (UInt_t n=0; n<fWidth; n++) {
         Double_t drefa=refa[m][n];
         Double_t dsig1a=sig1a[m][n];
         Double_t dsig2a=sig2a[m][n];
         Double_t dsig1ph=sig1ph[m][n];
         Double_t dsig2ph=sig2ph[m][n];
         fImage[m*fWidth+n] = drefa*drefa
                            + dsig1a*dsig1a
                            + dsig2a*dsig2a
                            + 2*drefa*dsig1a*cos(dsig1ph)
                            + 2*drefa*dsig2a*cos(dsig2ph)
                            + 2*dsig1a*dsig2a*cos(dsig1ph-dsig2ph);
      }
   }
}

Double_t TBitmap::ChiSquared(const TBitmap& map) const
{
   // Form a chi-square comparison between this image and the map argument

   if (map.fWidth != fWidth || map.fHeight != fHeight) {
      std::cerr << "Error in TBitmap::ChiSquared : this and map images"
                << " have different dimensions!" << std::endl;
      return 0;
   }
   Double_t chi2 = 0;
   for (UInt_t n=0; n<fWidth*fHeight; n++) {
      Double_t diff = fImage[n]-map.fImage[n];
      chi2 += diff*diff;
   }
   return chi2;
}

void TBitmap::Streamer(TBuffer &buf)
{
   // Put/get image data to/from stream buffer buf.

   if (buf.IsReading()) {
      buf >> fWidth >> fHeight;
      buf.ReadStaticArray(fImage);
   } else {
      buf << fWidth << fHeight;
      buf.WriteArray(fImage, fWidth*fHeight);
   }
}