/*************************************************************************
 * Copyright(c) 1998, University of Connecticut, All rights reserved.    *
 * Author: Richard T. Jones, Assoc. 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.           *
 *************************************************************************/
//////////////////////////////////////////////////////////////////////////
//									//
// Map2D Class								//
//									//
// The Map2D class provides a few useful methods for manipulating and	//
// comparing topographical maps within the root framework.  Map2D is	//
// derived from the basic two dimensional histogram TH2D class in root	//
// and inherits all of the storage and visualization capabilities of	//
// the root framework.  The main thing that is missing in TH2D is any	//
// concept of empty or invalid regions within the plot.  Many standard	//
// image formats incorporate "matte" information (or alpha channel) to  //
// indicate regions within the image frame that are not actually part	//
// of the image.  This is a basic feature of maps, that the coverage	//
// region is complex, and cannot be reduced to a rectangle.  The Map2D	//
// class introduces an alpha channel by defining the value (double)0.0	//
// as the matte value, or undefined.  Operations on that value should	//
// leave it undefined, for example: 0 + 3.8 = 0.  Interpolation on a	//
// map in a region next to a matte pixel must be designed so as not to	//
// treat it as z=0, and preserve the sharpness of the boundary between	//
// matte and ordinary pixels.						//
//									//
// The class currently supports the following operations on maps:	//
//	* addition/subtraction of 2 maps,				//
//	* shifting a map in all 3 dimensions,				//
//	* rescaling a map in all 3 dimensions,				//
//	* rotating a map in all 3 dimensions,				//
//      * computing correlations between 2 maps,			//
//      * determining the optimum alignment between 2 similar maps,	//
// all of which respect the meaning of matte pixels.			//
//									//
// This package was developed at the University of Connecticut by	//
// Richard T. Jones                            				//
//									//
//////////////////////////////////////////////////////////////////////////

#include <iostream>
#include <TMath.h>
#include "Map2D.h"

ClassImp(Map2D);

#ifndef SQR_FUNC
#define SQR_FUNC 1
Double_t sqr(Double_t x) { return x*x; }
#endif

const Double_t Map2D::zunit2xy = 0.001;

Map2D::Map2D() : TH2D()
{
   SetStats(0);
}

Map2D::Map2D(const Map2D &m) : TH2D((TH2D&)m)
{
   SetStats(0);
}

Map2D::Map2D(const TH2D &h) : TH2D(h)
{
   SetStats(0);
}

Map2D::~Map2D()
{ }

Map2D &Map2D::Add(const Map2D *h1, Double_t c1)
{
   // Overloads TH2D::Add, has the same basic behaviour,
   // except that zero plus anything is zero.

   Int_t nxbins = h1->GetXaxis()->GetNbins();
   Int_t nybins = h1->GetYaxis()->GetNbins();
   if (nxbins != GetXaxis()->GetNbins() ||
       nybins != GetYaxis()->GetNbins()) {
      std::cerr << "Error in Map2D::Add - source and destination histograms"
                << " have different dimensions, cannot add them"
                << std::endl;
      return *this;
   }
   for (Int_t i=1; i < nxbins; ++i) {
      for (Int_t j=1; j < nybins; ++j) {
         Double_t z1 = h1->GetBinContent(i,j);
         Double_t z0 = GetBinContent(i,j);
         if (z0 == 0 || z1 == 0) {
            SetBinContent(i,j,0);
         }
         else {
            SetBinContent(i,j,z0+c1*z1);
         }
      }
   }
   return *this;
}

Map2D &Map2D::Add(const Map2D *h1, const Map2D *h2,
                          Double_t c1, Double_t c2)
{
   // Overloads TH2D::Add, has the same basic behaviour,
   // except that zero plus anything is zero.

   Int_t nxbins = h1->GetXaxis()->GetNbins();
   Int_t nybins = h1->GetYaxis()->GetNbins();
   if (nxbins != h2->GetXaxis()->GetNbins() ||
       nybins != h2->GetYaxis()->GetNbins()) {
      std::cerr << "Error in Map2D::Add - two source histograms"
                << " have different dimensions, cannot add them"
                << std::endl;
      return *this;
   }
   else if (nxbins != GetXaxis()->GetNbins() ||
            nybins != GetYaxis()->GetNbins()) {
      Double_t xmin = h1->GetXaxis()->GetXmin();
      Double_t xmax = h1->GetXaxis()->GetXmax();
      Double_t ymin = h1->GetYaxis()->GetXmin();
      Double_t ymax = h1->GetYaxis()->GetXmax();
      SetBins(nxbins,xmin,xmax,nybins,ymin,ymax);
   }
   for (Int_t i=1; i < nxbins; ++i) {
      for (Int_t j=1; j < nybins; ++j) {
         Double_t z1 = h1->GetBinContent(i,j);
         Double_t z2 = h2->GetBinContent(i,j);
         if (z1 == 0 || z2 == 0) {
            SetBinContent(i,j,0);
         }
         else {
            SetBinContent(i,j,c1*z1+c2*z2);
         }
      }
   }
   return *this;
}

Map2D &Map2D::Shift(Double_t dx, Double_t dy, Double_t dz)
{
   // Translates the map image by a distance (dx,dy) in map coordinates,
   // and shifts all pixels up or down by vertical distance dz in
   // height units.  Image regions shifted outside the range of the
   // map are lost.  Image regions shift into the range of the map
   // from outside are zero.  Original contents are overwritten.

   Map2D *tmp = shift(dx,dy,dz);
   *this = *tmp;
   delete tmp;
   return *this;
}

Map2D *Map2D::shift(Double_t dx, Double_t dy, Double_t dz) const
{
   // Same as Shift(), except a new object is created to hold
   // the results and returned, and the original object is not
   // touched.  The caller must delete the returned object.

   Int_t nxbins = GetXaxis()->GetNbins();
   Int_t nybins = GetYaxis()->GetNbins();
   Double_t xmin = GetXaxis()->GetXmin();
   Double_t xmax = GetXaxis()->GetXmax();
   Double_t ymin = GetYaxis()->GetXmin();
   Double_t ymax = GetYaxis()->GetXmax();
   Double_t dbinx = dx/(xmax-xmin)*nxbins;
   Double_t dbiny = dy/(ymax-ymin)*nybins;
   Int_t di = int(dbinx);
   Int_t dj = int(dbiny);
   Int_t i0 = (di > 0)? di+1 : 1;
   Int_t j0 = (dj > 0)? dj+1 : 1;
   Int_t i1 = (di < 0)? di+nxbins : nxbins;
   Int_t j1 = (dj < 0)? dj+nybins : nybins;
   Double_t fx1 = dbinx-di;
   Double_t fy1 = dbiny-dj;
   Double_t fx0 = 1-fx1;
   Double_t fy0 = 1-fy1;

   Map2D *dst = new Map2D(*this);
   dst->Reset();
   for (Int_t i=i0; i < i1; ++i) {
      Int_t isrc = i-di;
      Double_t z00 = GetBinContent(isrc,j0-dj);
      Double_t z10 = GetBinContent(isrc+1,j0-dj);
      Double_t zint0 = (z00 == 0)? ((fx0 > 0.5)? 0 : z10) :
                       (z10 == 0)? ((fx1 > 0.5)? 0 : z00) :
                       z00*fx0 + z10*fx1;
      for (Int_t j=j0; j < j1; ++j) {
         Int_t jsrc = j-dj;
         Double_t z01 = GetBinContent(isrc,jsrc+1);
         Double_t z11 = GetBinContent(isrc+1,jsrc+1);
         Double_t zint1 = (z01 == 0)? ((fx0 > 0.5)? 0 : z11) :
                          (z11 == 0)? ((fx1 > 0.5)? 0 : z01) :
                          z01*fx0 + z11*fx1;
         Double_t zint = (zint0 == 0)? ((fy0 > 0.5)? 0 : zint1) :
                         (zint1 == 0)? ((fy1 > 0.5)? 0 : zint0) :
                         zint0*fy0 + zint1*fy1;
         dst->SetBinContent(i,j,(zint == 0)? 0 : zint+dz);
         zint0 = zint1;
      }
   }
   return dst;
}

Map2D &Map2D::Rotate(Double_t phi, Double_t theta, Double_t psi, Int_t degrees)
{
   // Rotates the map image by Euler angles phi,theta,psi about the center
   // of the map surface.  Angles may be in degrees (degrees=1) or
   // radians (degrees=0).  The angle phi determines the direction of
   // the theta rotation axis in the xy plane, and does not generate
   // an independent rotation.  Map rotations only make sense for small
   // theta.  The theta rotation is implemented as a rotation about the
   // x axis by angle theta*cos(phi) followed by a rotation about the y
   // axis by angle theta*sin(phi), with errors entering at the level of
   // theta^3.  Rotations in the xy plane (angle psi) may be of any size,
   // and may result in clipping of the image.  Regions of the map that
   // are rotated outside the map range are lost.  Original contents of
   // the map are overwritten.

   Double_t xmin = GetXaxis()->GetXmin();
   Double_t xmax = GetXaxis()->GetXmax();
   Double_t ymin = GetYaxis()->GetXmin();
   Double_t ymax = GetYaxis()->GetXmax();
   Double_t x0 = (xmin+xmax)/2;
   Double_t y0 = (ymin+ymax)/2;
   Int_t nxbins = GetXaxis()->GetNbins();
   Int_t nybins = GetYaxis()->GetNbins();
   Int_t i0 = nxbins/2;
   Int_t j0 = nybins/2;
   Int_t peri=0;
   Double_t sum0=1;
   Double_t sum1=GetBinContent(i0,j0);
   while (sum0 < nxbins*nybins/10) {
      ++peri;
      for (Int_t i=-peri; i < peri; ++i) {
          Double_t z;
          sum1 += ((z=GetBinContent(i0+i,j0-peri)) != 0)? ++sum0,z : 0;
          sum1 += ((z=GetBinContent(i0+peri,j0+i)) != 0)? ++sum0,z : 0;
          sum1 += ((z=GetBinContent(i0-i,j0+peri)) != 0)? ++sum0,z : 0;
          sum1 += ((z=GetBinContent(i0-peri,j0-i)) != 0)? ++sum0,z : 0;
      }
   }
   Double_t z0 = sum1/sum0;
   Double_t toradians = (degrees)? M_PI/180 : 1;
   Double_t thetax = theta*cos(phi*toradians);
   Double_t thetay = theta*sin(phi*toradians);
   Map2D *tmp1 = this;
   if (thetax != 0) {
      tmp1 = rotateX(thetax*toradians,y0,z0);
   }
   Map2D *tmp2 = tmp1;
   if (thetay != 0) {
      tmp2 = tmp1->rotateY(thetay*toradians,x0,z0);
   }
   Map2D *tmp3 = tmp2;
   if (psi != 0) {
      tmp3 = tmp2->rotateZ(psi*toradians,x0,y0);
   }
   if (tmp3 != this) {
      *this = *tmp3;
   }
   if (tmp1 != this) {
      delete tmp1;
   }
   if (tmp2 != tmp1) {
      delete tmp2;
   }
   if (tmp3 != tmp2) {
      delete tmp3;
   }
   return *this;
}

Map2D *Map2D::rotateX(Double_t thetarad, Double_t y0, Double_t z0) const
{
   // Rotates the map image about the x axis by angle thetarad in radians.
   // The rotation axis is located at coordinate y0 and height z0.
   // The map must remain single-valued after the rotation, which places
   // an upper limit on the magnitude of theta that depends on the
   // height profile of the map relative to z0.  An error message is
   // printed if this constraint is violated.  A new Map2D object is
   // created to hold the result, and the original object is untouched.
   // The caller must delete the returned object.

   Int_t nxbins = GetXaxis()->GetNbins();
   Int_t nybins = GetYaxis()->GetNbins();
   Double_t ymin = GetYaxis()->GetXmin();
   Double_t ymax = GetYaxis()->GetXmax();
   Double_t dy = (ymax-ymin)/nybins;
   Double_t costheta = cos(thetarad);
   Double_t sintheta = sin(thetarad);
   Map2D *dst = new Map2D(*this);
   dst->Reset();
   for (Int_t j=1; j < nybins; ++j) {
      Double_t y = ymin+(j-0.5)*(ymax-ymin)/nybins - y0;
      for (Int_t i=1; i < nxbins; ++i) {
         Double_t z = GetBinContent(i,j);
         if (z == 0) {
            dst->SetBinContent(i,j,0);
         }
         else {
            dst->SetBinContent(i,j,y*sintheta/zunit2xy+(z-z0)*costheta+z0);
         }
         if (fabs((z-z0)*sintheta)*zunit2xy > dy) {
            std::cerr << "Warning in Map2D::rotateX - loss of accuracy"
                      << " trying to rotate too far away from normal."
                      << std::endl;
         }
      }
   }
   Map2D *redst = dst->rescale(1,costheta,1,0,y0,0);
   delete dst;
   return redst;
}

Map2D *Map2D::rotateY(Double_t thetarad, Double_t x0, Double_t z0) const
{
   // Rotates the map image about the y axis by angle thetarad in radians.
   // The rotation axis is located at coordinate x0 and height z0.
   // The map must remain single-valued after the rotation, which places
   // an upper limit on the magnitude of theta that depends on the
   // height profile of the map relative to z0.  An error message is
   // printed if this constraint is violated.  A new Map2D object is
   // created to hold the result, and the original object is untouched.
   // The caller must delete the returned object.

   Int_t nxbins = GetXaxis()->GetNbins();
   Int_t nybins = GetYaxis()->GetNbins();
   Double_t xmin = GetXaxis()->GetXmin();
   Double_t xmax = GetXaxis()->GetXmax();
   Double_t dx = (xmax-xmin)/nxbins;
   Double_t costheta = cos(thetarad);
   Double_t sintheta = sin(thetarad);
   Map2D *dst = new Map2D(*this);
   dst->Reset();
   for (Int_t i=1; i < nxbins; ++i) {
      Double_t x = xmin+(i-0.5)*(xmax-xmin)/nxbins - x0;
      for (Int_t j=1; j < nybins; ++j) {
         Double_t z = GetBinContent(i,j);
         if (z == 0) {
            dst->SetBinContent(i,j,0);
         }
         else {
            dst->SetBinContent(i,j,-x*sintheta/zunit2xy+(z-z0)*costheta+z0);
         }
         if (fabs((z-z0)*sintheta)*zunit2xy > dx) {
            std::cerr << "Warning in Map2D::rotateY - loss of accuracy"
                      << " trying to rotate too far away from normal."
                      << std::endl;
         }
      }
   }
   Map2D *redst = dst->rescale(costheta,1,1,x0,0,0);
   delete dst;
   return redst;
}

Map2D *Map2D::rotateZ(Double_t phirad, Double_t x0, Double_t y0) const
{
   // Rotates the map image about the z axis by angle thetarad in radians.
   // The rotation axis is located at (x0,y0) in map coordinates.  There
   // are no assumed bounds on phirad.  Regions of the image that move
   // outside the map range are clipped and lost.  Regions of the image
   // that move into the range from outside are zero.  A new Map2D object
   // is created to hold the result, and the original object is untouched.
   // The caller must delete the returned object.

   Int_t nxbins = GetXaxis()->GetNbins();
   Int_t nybins = GetYaxis()->GetNbins();
   Double_t xmin = GetXaxis()->GetXmin();
   Double_t xmax = GetXaxis()->GetXmax();
   Double_t ymin = GetYaxis()->GetXmin();
   Double_t ymax = GetYaxis()->GetXmax();
   Double_t dx = (xmax-xmin)/nxbins;
   Double_t dy = (ymax-ymin)/nybins;
   Double_t cosphi = cos(phirad);
   Double_t sinphi = sin(phirad);

   Map2D *dst = new Map2D(*this);
   dst->Reset();
   for (Int_t i=1; i <= nxbins; ++i) {
      Double_t x = xmin+(i-0.5)*dx - x0;
      for (Int_t j=1; j <= nybins; ++j) {
         Double_t y = ymin+(j-0.5)*dy - y0;
         Double_t xsrc = x*cosphi+y*sinphi + x0;
         Double_t ysrc = y*cosphi-x*sinphi + y0;
         Double_t xbins = (xsrc-xmin)/dx;
         Double_t ybins = (ysrc-ymin)/dy;
         if (xbins < 0 || xbins > nxbins ||
             ybins < 0 || ybins > nybins) {
            continue;
         }
         Int_t i0 = int(xbins+0.5);
         Int_t j0 = int(ybins+0.5);
         i0 = (i0 < 1)? 1 : i0;
         j0 = (j0 < 1)? 1 : j0;
         Int_t i1 = (i0 >= nxbins)? i0 : i0+1;
         Int_t j1 = (j0 >= nybins)? j0 : j0+1;
         Double_t fx0 = xbins+0.5-i0;
         Double_t fy0 = ybins+0.5-j0;
         Double_t fx1 = 1-fx0;
         Double_t fy1 = 1-fy0;
         Double_t z00 = GetBinContent(i0,j0);
         Double_t z10 = GetBinContent(i1,j0);
         Double_t z01 = GetBinContent(i0,j1);
         Double_t z11 = GetBinContent(i1,j1);
         Double_t zint0 = (z00 == 0)? ((fx0 > 0.5)? 0 : z10) :
                          (z10 == 0)? ((fx1 > 0.5)? 0 : z00) :
                          z00*fx0 + z10*fx1;
         Double_t zint1 = (z01 == 0)? ((fx0 > 0.5)? 0 : z11) :
                          (z11 == 0)? ((fx1 > 0.5)? 0 : z01) :
                          z01*fx0 + z11*fx1;
         Double_t zint = (zint0 == 0)? ((fy0 > 0.5)? 0 : zint1) :
                         (zint1 == 0)? ((fy1 > 0.5)? 0 : zint0) :
                         zint0*fy0 + zint1*fy1;
         dst->SetBinContent(i,j,zint);
      }
   }
   return dst;
}

Map2D &Map2D::Rescale(Double_t sx, Double_t sy, Double_t sz)
{
   // Rescales the map image along all three axes by independent scale
   // factors.  The origin of the rescaling is taken to be the center
   // of the map in x and y, and z=0.  The original map is overwritten.
   // Note: the special case of sx=-1,sy=1,sz=-1 or sx=1,sy=-1,sz=-1
   // can be used to invert the surface, effectively flipping it over.

   Double_t xmin = GetXaxis()->GetXmin();
   Double_t xmax = GetXaxis()->GetXmax();
   Double_t ymin = GetYaxis()->GetXmin();
   Double_t ymax = GetYaxis()->GetXmax();
   Map2D *tmp = rescale(sx,sy,sz,(xmax+xmin)/2,(ymax-ymin)/2,0);
   *this = *tmp;
   delete tmp;
   return *this;
}

Map2D *Map2D::rescale(Double_t sx, Double_t sy, Double_t sz,
                      Double_t x0, Double_t y0, Double_t z0) const
{
   // Rescales the map image along all three axes by independent scale
   // factors.  The origin of the rescaling is (x0,y0) in map coordinates,
   // and z0 in height units.  A new object is created to hold the result
   // and the original object is untouched.  The caller must delete the
   // returned object.

   Int_t nxbins = GetXaxis()->GetNbins();
   Int_t nybins = GetYaxis()->GetNbins();
   Double_t xmin = GetXaxis()->GetXmin();
   Double_t xmax = GetXaxis()->GetXmax();
   Double_t ymin = GetYaxis()->GetXmin();
   Double_t ymax = GetYaxis()->GetXmax();
   Double_t dx = (xmax-xmin)/nxbins;
   Double_t dy = (ymax-ymin)/nybins;

   Map2D *dst = new Map2D(*this);
   dst->Reset();
   for (Int_t i=1; i <= nxbins; ++i) {
      Double_t x = (xmin+(i-0.5)*dx-x0)/sx + x0;
      for (Int_t j=1; j <=nybins; ++j) {
         Double_t y = (ymin+(j-0.5)*dy-y0)/sy + y0;
         Double_t xbins = (x-xmin)/dx;
         Double_t ybins = (y-ymin)/dx;
         Int_t i0 = int(xbins+0.5);
         Int_t j0 = int(ybins+0.5);
         i0 = (i0 < 1)? 1 : i0;
         j0 = (j0 < 1)? 1 : j0;
         Int_t i1 = (i0 >= nxbins)? i0 : i0+1;
         Int_t j1 = (j0 >= nybins)? j0 : j0+1;
         Double_t fx0 = i1-xbins-0.5;
         Double_t fy0 = j1-ybins-0.5;
         Double_t fx1 = 1-fx0;
         Double_t fy1 = 1-fy0;
         Double_t z00 = GetBinContent(i0,j0);
         Double_t z10 = GetBinContent(i1,j0);
         Double_t z01 = GetBinContent(i0,j1);
         Double_t z11 = GetBinContent(i1,j1);
         Double_t zint0 = (z00 == 0)? ((fx0 > 0.5)? 0 : z10) :
                          (z10 == 0)? ((fx1 > 0.5)? 0 : z00) :
                          z00*fx0 + z10*fx1;
         Double_t zint1 = (z01 == 0)? ((fx0 > 0.5)? 0 : z11) :
                          (z11 == 0)? ((fx1 > 0.5)? 0 : z01) :
                          z01*fx0 + z11*fx1;
         Double_t zint = (zint0 == 0)? ((fy0 > 0.5)? 0 : zint1) :
                         (zint1 == 0)? ((fy1 > 0.5)? 0 : zint0) :
                         zint0*fy0 + zint1*fy1;
         dst->SetBinContent(i,j,(zint-z0)*sz+z0);
      }
   }
   return dst;
}

Double_t Map2D::Correlation(const Map2D *m1, Double_t contrast)
{
   // Computes the height correlation with another map, useful for
   // aligning two images of the same terrain.  The contrast
   // argument serves as a weight factor on terms in the sum
   // where one or the other of the two maps is zero.  This allows
   // control over how zero regions are treated, such that
   //	contrast=0 : ignore zero regions
   //	contrast=1 : treat them like plateaus with z=0
   //   contrast=BIG : boundaries dominate, maps are B/W
   // Agreement between their coordinate ranges is not required,
   // but the two maps must have the same dimensions. 

   Int_t nxbins = GetXaxis()->GetNbins();
   Int_t nybins = GetYaxis()->GetNbins();
   if (m1->GetXaxis()->GetNbins() != nxbins ||
       m1->GetYaxis()->GetNbins() != nybins) {
      std::cerr << "Error in Map2D::Correlation - "
                << "maps must have the same dimensions."
                << std::endl;
      return 0;
   }
   Double_t sum00=0;
   Double_t sum10=0;
   Double_t sum01=0;
   Double_t sum11=0;
   Double_t sum20=0;
   Double_t sum02=0;
   for (Int_t i=1; i <= nxbins; ++i) {
      for (Int_t j=1; j <= nybins; ++j) {
         Double_t z0 = GetBinContent(i,j);
         Double_t z1 = m1->GetBinContent(i,j);
         Double_t wgt = (z0*z1 == 0)? contrast : 1;
         sum00 += wgt;
         sum10 += wgt*z0;
         sum01 += wgt*z1;
         sum11 += wgt*z0*z1;
         sum20 += wgt*z0*z0;
         sum02 += wgt*z1*z1;
      }
   }
   Double_t covar = sum11/sum00-(sum10/sum00)*(sum01/sum00);
   Double_t var0 = sum20/sum00-(sum10/sum00)*(sum10/sum00);
   Double_t var1 = sum02/sum00-(sum01/sum00)*(sum01/sum00);
   return covar/sqrt((var0*var1 > 0)? var0*var1 : 1e-100);
}
Map2D &Map2D::Align(const Map2D *h1, Double_t contrast)
{
   std::cerr << "Map2D::Align -- not yet implemented"
             << std::endl;
   return *this;
}