// Copyright (C) 1996  Eleftherios Gkioulekas <lf@amath.washington.edu>
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.

#include <iostream.h>
#include <iomanip.h>
#include <fstream.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <assert.h>
#include "qx.h"

//
// Color conversion functions
//

// Specify the color in (r,g,b) variables in the 0..1 range.
// You are returned an Xlib XColor structure
XColor QxColor(double r,double g,double b)
{
 XColor col;
 col.red   = (unsigned short int) (65535.*r);
 col.green = (unsigned short int) (65535.*g);
 col.blue  = (unsigned short int) (65535.*b);
 return col;
}

double QxLum(double r,double g,double b)
{ return 0.299*r+0.587*g+0.114*b; }

// Red the red,green,blue fields of an Xlib XColor.
// Values returned in the 0..1 range
int QxColorRed  (XColor col) { return int(col.red)/256; }
int QxColorGreen(XColor col) { return int(col.green)/256; }
int QxColorBlue (XColor col) { return int(col.blue)/256; }

// Abort in case of error
void QxError(char *error) 
{
 cerr << "qx: " << error;
 abort();
}

                           ////////////////////
                           // Our colormaps! //
                           ////////////////////

// How to make your own colormap:
// If you intend to use the PutDataOnWindow methods then your colormaps
// should follow the following conventions:
// col[0] = black (i.e (0,0,0) )
// col[1] = "white" (ex. (1,1,1) )
// col[2..N+2] = your colormap
// Note that white can be other color if your colormap includes white.
// For example in the grey-scale colormap, white is red+blue
// In general white should be a color well-contrasted by your colormap.
// You may want to draw axises or boundary on top of your colored image.

// How to use:
// XColor *colormap = QxMakeGreyscaleColormap(200);
// server.ColorRequest(200,colormap);

// Here is a greyscale colormap
XColor *QxMakeGreyscaleColormap(int N,double cor)
{
 XColor *col = new XColor[N];
 col[0] = QxColor(0,0,0); col[1] = QxColor(0,1,0);
 double lum;
 for (int i=2;i<N;i++)
 { lum = pow(double(i-2)/double(N-2),cor); col[i] = QxColor(lum,lum,lum); }
 return col;
}

// A rainbow colormap. Probably needs some improvement
XColor *QxMakeRainbowColormap(int N,double cor)
{
 XColor *col = new XColor[N];
 col[0] = QxColor(0,0,0); col[1] = QxColor(1,1,1);
 double r,g,b,lum;
 for (int i=2;i<N;i++)
 {
  lum = pow(double(i-2)/double(N-2),cor);
  r = (lum < 1./3.) ? 0     : ((lum < 2./3.) ? 3*(lum-1./3.) : 1);
  g = (lum < 0.5)   ? 2*lum : ((lum < 2./3.) ? 1 : 1-3*(lum-2./3.));
  b = (lum < 1./4.) ? 1     : ((lum < 2./3.) ? 1-(lum-1./4)/(2./3.-1./4.) : 0);
  col[i] = QxColor(r,g,b);
 }
 return col;
}
                           ////////////////////////
                           // The QxServer class //
                           ////////////////////////

// Connects to the X server and gets some basic information from it 
QxServer::QxServer()
{
 // Alloc and connect to the server 
 display = XOpenDisplay(NULL); 
 if (display == NULL) QxError("Can not connect to X server\n");
 
 // Now get some basic information from the server 
 screen   = XDefaultScreen(display);
 gc       = XDefaultGC(display,screen);
 visual   = XDefaultVisual(display,screen);
 rtwin    = XRootWindow(display,screen);
 rtwidth  = XDisplayWidth(display,screen);
 rtheight = XDisplayHeight(display,screen);
 cmap     = XDefaultColormap(display,screen);
 
 // The following should be set by a call to QxColorRequest 
 Ncols = 0;
 cols = NULL;
}

// Disconnects from the X server 
QxServer::~QxServer()
{
 if (cols != NULL)     delete[] cols;
 XCloseDisplay(display);
}

// Attempts to request the number of colors specified by the user
// Returns 'true' is success and 'false' if failure
bool QxServer::ColorRequest(int N) 
{
 // Make sure this is the first time
 assert(Ncols == 0 && "Allocating color-cells AGAIN?! DON'T!");
 // Useless shit required by X
 long unsigned int plane_masks[1]; 
 // Attempt to allocate the colors
 pixels = new (long unsigned int)[N];
 if (XAllocColorCells(display,cmap,False,plane_masks,0,pixels,N))
 { Ncols = N; cols  = new XColor[N]; return true; }
 else 
 { delete pixels; pixels = NULL; return false; }
}

// Get all the colorcells we can possibly get
// Returns the number of colors obtained
int QxServer::SuckAllColors(void)
{
 assert(Ncols == 0 && "Allocating color-cells AGAIN?! DON'T!");
 int N = 255;
 while(N-- > 0) { if (ColorRequest(N)) break; }
 return N;
}

// Set the color for a colorcell
QxServer* QxServer::SetColor(int pxl,XColor col)
{
 assert(pxl >= 0 && pxl < Ncols);
 cols[pxl] = col;
 cols[pxl].pixel = pixels[pxl];
 cols[pxl].flags = DoRed | DoGreen | DoBlue;
 return this;
}

// This function will take an array of 'n' colors 'col' and set consecutive
// colorcells with these colors, starting from 'start'
QxServer* QxServer::SetColor(int start,XColor *col,int n)
{
 assert(start+n-1 < Ncols);
 for (int i=0;i<n;i++) SetColor(start+i,col[i]);
 return this;
}

// Similar to previous function. Only start = 0
QxServer* QxServer::SetColor(XColor *col,int n) 
{ SetColor(0,col,n); return this; }

// This function will flush any changes we made to the color-cells
QxServer* QxServer::FlushColors(void)
{ 
 XStoreColors(display,cmap,cols,Ncols); 
 XFlush(display);
 return this; 
}

                     ////////////////////////////////////
                     //        QxImageWindow class     //
                     // Windows for putting in images  //
                     ////////////////////////////////////

// Constructor -- Create a window
QxImageWindow::QxImageWindow(int width,int height,char *n):QxServer()
{
 QxImageWindow *w;
 XSetWindowAttributes winattr;
 XSizeHints size_hints;
 XTextProperty text_prop;
 XEvent report;

 // Set the easy stuff 
 SzX = width; SzY = height;
 // Make an image for the window 
 ImD = new char[SzX*SzY];
 Im  = XCreateImage(display,visual,8,ZPixmap,0,ImD,SzX,SzY,8,SzX);
 // Make the window itself 
 win = XCreateSimpleWindow(display,rtwin,0,0,SzX,SzY,2,0,0);
 XSelectInput(display,win, ExposureMask | ButtonPressMask );
 // Enable backing store 
 winattr.backing_store = Always;
 XChangeWindowAttributes(display,win,CWBackingStore,&winattr);
 // Make the window non-resizable 
 size_hints.flags  = USSize | PMinSize | PMaxSize;
 size_hints.width  = size_hints.min_width  = size_hints.max_width  = SzX;
 size_hints.height = size_hints.min_height = size_hints.max_height = SzY; 
 XSetWMNormalHints(display,win,&size_hints);
 // Set the window name 
 name = new char[1+strlen(n)]; strcpy(name,n);
 XStringListToTextProperty(&name,1,&text_prop);
 XSetWMName(display,win,&text_prop);
 // Wait for the window to show up 
 XMapWindow(display,win);
 for (XNextEvent(display,&report);
      report.type != Expose;
      XNextEvent(display,&report));
}

// Destructor -- Destroy a window
QxImageWindow::~QxImageWindow()
{
 XDestroyWindow(display,win);
 XDestroyImage(Im);
 delete[] name;
 XFlush(display);
}

// Set a pixel to a color
QxImageWindow* QxImageWindow::SetPixel(int x,int y,int color)
{ 
 if (0<=x && x<SzX && 0<=y && y<SzY) ImD[x+(SzY-1-y)*SzX] = cols[color].pixel; 
 return this; 
}

// Same but specify coordinates in a funny coordinate system
QxImageWindow* QxImageWindow::SetPixel(QxCoords c,double dx,double dy,int color)
{
 SetPixel(c.ox + int(dx*c.scale), 
          c.oy + int(dy*c.scale),
          color);
 return this;
}

// An integer algorithm for drawing lines. 
// Originally written by Egon (egon@ugcs.caltech.edu)
// Egon is cool! Long live Egon :)
QxImageWindow* QxImageWindow::DrawLine(int x0,int y0,int x1,int y1,int color)
{
 int dx,dy,sx,sy,ax,ay,dec;
 dx = x1-x0; dy = y1-y0;
 sx = (dx>0 ? 1 : (dx<0 ? -1:0));  sy = (dy>0 ? 1 : (dy<0 ? -1:0));
 dx = abs(dx); dy = abs(dy); ax = dx*2; ay = dy*2;
 if (dy<dx)
 {
  dec=ay-dx; 
  while(1)
  {
   SetPixel(x0,y0,color); 
   if (x0==x1) break; if (dec >= 0) { dec -= ax; y0 += sy; }
   x0 += sx; dec += ay;
  }
 }
 else
 {
  dec = ax-dy;
  while(1)
  {
   SetPixel(x0,y0,color);
   if (y0==y1) break; if (dec >= 0) { dec -= ay; x0 += sx; }
   y0 += sy; dec += ax;
  }
 }
 return this;
}

QxImageWindow* QxImageWindow::DrawLine
(QxCoords c,double dx0,double dy0,double dx1,double dy1,int color)
{
 DrawLine(c.ox + int(dx0*c.scale), c.oy + int(dy0*c.scale),
          c.ox + int(dx1*c.scale), c.oy + int(dy1*c.scale),
          color);
 return this;
}

// You probably want to use this to clear the window
QxImageWindow* QxImageWindow::FillWindow(int color)
{ 
 for (int x=0;x<SzX;x++) for (int y=0;y<SzY;y++) SetPixel(x,y,color); 
 return this; 
}

// Update the window
QxImageWindow* QxImageWindow::FlushWindow(void)
{
 XPutImage(display,win,gc,Im,0,0,0,0,SzX,SzY);
 XFlush(display);
 return this;
}

// Update only part of the window. Useful when speed is critical
QxImageWindow* QxImageWindow::FlushPartOfWindow(int ux,int uy,int lx,int ly)
{
 XPutImage(display,win,gc,Im,ux,uy,lx,ly,SzX,SzY);
 XFlush(display);
 return this;
}

// Wait for the user to click mouse on the window
QxImageWindow* QxImageWindow::WaitMouseClick(void)
{
 XEvent report;

 for (XNextEvent(display,&report);
      report.type != ButtonPress ||
      (report.type == ButtonPress && report.xbutton.window != win);
      XNextEvent(display,&report));
 return this;
}

// Save the current contents of our window to a file
QxImageWindow* QxImageWindow::WritePPM(ostream &output)
{
 // Put out the PPM header
 output << "P3" << endl;
 output << SzX << "  " << SzY << endl;
 output << "255" << endl;
 
 // Create a PixelValue -> Index mapping
 int mapping[256]; for (int i=0;i<256;i++) mapping[i] = 0; 
 for (int i=0;i<Ncols;i++) mapping[cols[i].pixel] = i;

 // Now dump the image
 int r,g,b;
 XColor col;
 unsigned char colpixel;
 for (int y=0;y<SzY;y++) for (int x=0;x<SzX;x++)
 { 
  colpixel = (unsigned char) ImD[x+y*SzX];
  col = cols[mapping[int(colpixel)]];
  r = QxColorRed(col);
  g = QxColorGreen(col);
  b = QxColorBlue(col);
  output << r << "  " << g << "  " << b << "  " << endl;
 }
 return this;
}

// For those who are too lazy to open and close their own files :)
QxImageWindow* QxImageWindow::WritePPM(char *filename)
{
 ofstream file(filename);
 if (file.bad()) 
 {
  cerr << "Can't open file " << filename << endl; 
  return this;
 }
 WritePPM(file);
 file << flush;
 file.close();
}

// This function takes a double 'array' whose values are supposed to be
// in the range min,max, scales it to our colormap, then sends it into
// an area in the window. The data array is assumed to be same size as
// the window
QxImageWindow* QxImageWindow::PutDataOnWindow
(double **data,double min,double max)
{
 int i,j;
 int col;
 double val;

 for (i=0;i<SzX;i++) for (j=0;j<SzY;j++)  
 {
  val = (data[i][j]-min)/(max-min);               // Scale to 0..1 
  col = 2+int((Ncols-2)*val);                     // Scale to 2..Ncols-1 
  if (col<2) col=2; if (col>=Ncols) col=Ncols-1;  // Clip if needed 
  SetPixel(i,j,col);
 }
 return this;
}

// This function takes a data value 'data' located at (i,j) and places
// it on our window. The pixel color used is scaled from the range (min,max)
// to our colormap.
QxImageWindow* QxImageWindow::PutDataOnWindow
(int i,int j,double data,double min,double max)
{
 double val = (data-min)/(max-min);
 int    col = 2+int((Ncols-2)*val);
 if (col<2) col=2; if (col>=Ncols) col=Ncols-1;
 SetPixel(i,j,col);
}

// This function takes a data value 'data' located at (i,j) and places
// it on our window. The pixel color used is scaled from the range (min,max)
// to our colormap. However a big part of the colormap is blacked out.
// As a result, we will obtain a contour plot. The variable 'b'
// determines how much to black out. Can take values from 1 to 100
QxImageWindow* QxImageWindow::PutDataOnWindow
(int i,int j,double data,double min,double max,int b)
{
 double val  = (data-min)/(max-min);
 double dcol = 2+(Ncols-2)*val; double ddcol = dcol-int(dcol);
 int col = (0.5-b/200. < ddcol && ddcol < 0.5+b/200.) ? int(dcol) : 0;
 if (col<2) col=2; if (col>=Ncols) col=Ncols-1;
 SetPixel(i,j,col);
}
 
// This function does the same as the previous one, but it is meant to
// be used for these cases where the array 'data' is very small and we
// want to magnify it. We assume that the scaling of the array is
// the same as the scaling of the window and that "ratio" is the size of
// an array cell in pixel by pixel.
QxImageWindow* QxImageWindow::PutDataOnWindow
(double **data,int ratio,double min,double max)
{
 int col;
 double val;

 for (int i=0;i<SzX/ratio;i++) for (int j=0;j<SzY/ratio;j++)
 {
  val = (data[i][j]-min)/(max-min);               // Scale to 0..1
  col = 2+int((Ncols-2)*val);                     // Scale to 2..Ncols-1
  if (col<2) col=2; if (col>=Ncols) col=Ncols-1;  // Clip if needed
  for (int x=0;x<ratio;x++) for (int y=0;y<ratio;y++)
  { SetPixel(ratio*i+x,ratio*j+y,col); }
 }
 return this;
}