//
// Programmer:    Craig Stuart Sapp <craig@ccrma.stanford.edu>
// Creation Date: Sat May 12 04:36:30 PDT 2007 (derived from hicor)
// Last Modified: Sat May 12 04:36:45 PDT 2007
// Filename:      ...sig/examples/all/hicor.cpp
// Web Address:   http://sig.sapp.org/examples/museinfo/archetype/archetype.cpp
// Syntax:        C++; museinfo
//
// Description:   Generate a hierarchical correlation plots compared to
//                input data.
//

#include "Array.h"
#include "Options.h"
#include "PixelColor.h"

#include <string.h>
#include <stdio.h>
#include <math.h>

#ifndef OLDCPP
   using namespace std;
   #include <iostream>
   #include <fstream>
#else
   #include <iostream.h>
   #include <fstream.h>
#endif


#define COLOR_BW  0
#define COLOR_HUE 1

#define SHAPE_TRIANGLE  0
#define SHAPE_RECTANGLE 1


// function declarations:
void     checkOptions(Options& opts, int argc, char** argv);
void     example(void);
void     usage(const char* command);
void     calculateAverage(Array<Array<double> >& correlations, 
                                Array<double>& a);
void     calculateCorrelations(Array<Array<double> >& correlations,
                                Array<double>& a, 
                                Array<double>& b);
double   pearsonCorrelation(int size, double* a, double* b);
double   getMean(int size, double* a);
void     printImageTriangle(Array<Array<double> >& correlation);
void     printImageRectangle(Array<Array<double> >& correlation);
void     printCorrelationPixel(double value, int style);
double   average(int size, double* data);
void     normalizeAverageData(Array<Array<double> >& correlations);
double   getStandardDeviation(double mean, Array<double>& data);
int      logScale(int level, int maxlevel, double factor);
void     calculateAverageDifference(Array<Array<double> >& correlations, 
                                Array<double>& a, 
                                Array<double>& b);

void     makeSineArch(Array<double>& arch, int size);
void     makeRamp(Array<double>& ramp, int size);
void     calculateArchCorrelations(Array<Array<double> >& correlations, 
                                Array<double>& a);
void     calculateRampCorrelations(Array<Array<double> >& correlations, 
                                Array<double>& a);




// User interface variables:
Options   options;
int       dataq     = 0;         // debug: display input data only with -d
int       correlq   = 0;         // debug: display correlations with -c
int       colortype = COLOR_BW;  // used with various color options
int       plotshape = SHAPE_TRIANGLE;  // used with -r option
int       averageq  = 0;         // used with -a option
int       diffq     = 0;         // display average differences with -d option
int       condiffq  = 0;         // display contoured differences, --cd option
int       column    = 0;         // used with -n option (for averaging)
int       logq      = 0;         // used with -g option
double    logfactor = 1500;      // used with -f option
int       flipq     = 0;         // flip the data from tempo to duration
int       polyq     = 0;         // used with -p option
int       archq     = 0;         // used with -A option
int       rampq     = 0;         // used with -R option
int       maxheight = -1;        // used with --height option
int       gradientq = 0;         // used with -G option

//////////////////////////////////////////////////////////////////////////

int main(int argc, char** argv) {
   // process the command-line options
   checkOptions(options, argc, argv);

   double aval;
   double bval;

   Array<double> a;
   Array<double> b;

   a.setSize(100000);
   a.setGrowth(100000);
   a.setSize(0);

   b.setSize(100000);
   b.setGrowth(100000);
   b.setSize(0);

   ifstream infile;
   const char* filename = "";

   if (options.getArgCount() <= 0) {
      filename = "<STDIN>";
   
   } else {
      filename = options.getArg(1);

      #ifndef OLDCPP
         infile.open(filename, ios::in);
      #else
         infile.open(filename, ios::in | ios::nocreate);
      #endif
   }

   if (!infile.is_open()) {
      cerr << "Cannot open file for reading: " << filename << endl;
      exit(1);
   }

   char templine[4096];

   if (averageq) {

      while (!infile.eof()) {
         infile.getline(templine, 4096, '\n');
         if (infile.eof() && (strcmp(templine, "") == 0)) {
            break;
         } else if (sscanf(templine, "%lf\t%lf", &aval, &bval) == 2) {
            if (flipq) {
               aval = 60.0 / aval;
               bval = 60.0 / bval;
            }
            a.append(aval);
            b.append(bval);
         } else if (sscanf(templine, "%lf", &aval) == 1) {
            if (flipq) {
               aval = 60.0 / aval;
            }
            a.append(aval);
            b.append(aval);
         }
      }

   } else {

      while (!infile.eof()) {
         infile.getline(templine, 4096, '\n');
         if (infile.eof() && (strcmp(templine, "") == 0)) {
            break;
         } else {
            if (sscanf(templine, "%lf\t%lf", &aval, &bval) == 2) {
               if (flipq) {
                  aval = 60.0 / aval;
                  bval = 60.0 / bval;
               }
               a.append(aval);
               b.append(bval);
            }
         }
      }

   }

   int i, j;
   if (dataq) {
      for (i=0; i<a.getSize(); i++) {
         cout << a[i] << "\t" << b[i] << endl;
      }
      exit(0);
   } 

   Array<Array<double> > correlations;
   correlations.setSize(a.getSize());
   for (i=0; i<a.getSize(); i++) {
      correlations[i].setSize(i+1);
      correlations[i].allowGrowth(0);
   }

   if (averageq) {
      if (column == 0) {
         calculateAverage(correlations, a);
      } else {
         calculateAverage(correlations, b);
      }
   } else if (archq) {
      calculateArchCorrelations(correlations, a);
   } else if (rampq) {
      calculateRampCorrelations(correlations, a);
   } else if (condiffq || diffq) {
      calculateAverageDifference(correlations, a, b);
   } else {
      calculateCorrelations(correlations, a, b);
   }

   if (correlq) {
      for (i=0; i<correlations.getSize(); i++) {
         for (j=0; j<correlations[i].getSize(); j++) {
            cout << correlations[i][j];
            if (j < correlations[i].getSize()-1) {
               cout << '\t';
            }
         }
         cout << '\n';
      }
      exit(0);
   }

   switch (plotshape) {
      case SHAPE_RECTANGLE:
         printImageRectangle(correlations);
         break;
      case SHAPE_TRIANGLE:
      default:
         printImageTriangle(correlations);
   }


   return 0;
}


//////////////////////////////////////////////////////////////////////////


//////////////////////////////
//
// printImageTriangle --
//

void printImageTriangle(Array >& correlation) {
   const char* background = options.getString("background");

   int maxcolumn = correlation.getSize() * 2;
   int maxrow    = (correlation.getSize()) * 2;

   int startrow = 0;
   if (maxheight > 0 && maxheight < correlation.getSize()) {
      startrow = correlation.getSize() - maxheight;
      maxrow = (maxheight) * 2;
   }

   cout << "P3\n";
   cout << maxcolumn << " " << maxrow << "\n";
   cout << "255\n";

   int i, j;
   int ii;
   int datawidth;
   int bgwidth;
   for (i=startrow; i<correlation.getSize(); i++) {
      ii = i;
      if (logq) {
         ii = logScale(i, correlation.getSize(), logfactor);
      }
      datawidth = correlation[ii].getSize();
      bgwidth = (maxcolumn - datawidth * 2)/2;

      for (j=0; j<bgwidth; j++) {
         cout << " " << background << " ";
      }
      for (j=0; j<datawidth; j++) {
         printCorrelationPixel(correlation[ii][j], colortype);
      }
      for (j=0; j<bgwidth; j++) {
         cout << " " << background << " ";
      }
      cout << "\n";

      // repeat the row to double it (to make image square)      
      for (j=0; j<bgwidth; j++) {
         cout << " " << background << " ";
      }
      for (j=0; j<datawidth; j++) {
         printCorrelationPixel(correlation[ii][j], colortype);
      }
      for (j=0; j<bgwidth; j++) {
         cout << " " << background << " ";
      }
      cout << "\n";
   }
}



//////////////////////////////
//
// logScale -- 
//

int logScale(int level, int maxlevel, double factor) {
   level = maxlevel - level - 1;
   double dummy = factor;
   double minval = log10(10.0);
   double valrange = log10(10.0 + dummy) - minval;
   int    templevel = maxlevel - level;
   double tempval = (log10(10.0 + dummy * templevel/maxlevel)-minval)/valrange;
   int    newlevel = (int)(tempval * maxlevel + 0.5);
   // newlevel = maxlevel - newlevel;
   newlevel = newlevel - 1;
   return newlevel;
}



//////////////////////////////
//
// printImageRectangle --
//

void printImageRectangle(Array >& correlation) {

   int maxcolumn = correlation.getSize() * 2;
   int maxrow    = (correlation.getSize()-1) * 2;
   if (maxheight > 0 && maxheight < maxrow) {
      maxrow = maxheight;
   }

   cout << "P3\n";
   cout << maxcolumn << " " << maxrow << "\n";
   cout << "255\n";

   int i, j;
   int datawidth;
   int jstretch;

   for (i=0; i<correlation.getSize()-1; i++) {
      datawidth = correlation[i].getSize();

      for (j=0; j<maxcolumn/2; j++) {
         jstretch = int(double(j)/maxcolumn*2*datawidth);
         printCorrelationPixel(correlation[i][jstretch], colortype);
      }
      cout << "\n";

      // repeat for double pixel size
      for (j=0; j<maxcolumn/2; j++) {
         jstretch = int(double(j)/maxcolumn*2*datawidth);
         printCorrelationPixel(correlation[i][jstretch], colortype);
      }
      cout << "\n";

   }
}


 
//////////////////////////////
//
// printCorrelationPixel --
//

void printCorrelationPixel(double value, int style) {

   PixelColor pc;

   switch (style) {
      case COLOR_HUE:
         pc.setHue(-((value+1)/2.0)*0.80 - 0.20); // -1=purple, 
	                                          // 0=green +1=red
         //pc.setHue(value);	 // -1=purple, 0=green +1=red
         break;
      case COLOR_BW:
      default:
         pc.setGrayNormalized((value + 1.0)/2.0);
   }

   if (diffq) {
      if (value > 0) {
         pc.setColor(255,0,0);
      } else if (value < 0) {
         pc.setColor(0,0,255);
      } else {
         pc.setColor(255,255,255);
      }
   } else if (condiffq) {
      if      (value >   0.210)  {  pc.setColor(252,41,41);   } // red
      else if (value >   0.150)  {  pc.setColor(252,41,210);  } // pink
      else if (value >   0.090)  {  pc.setColor(255,144,41);  } // orange
      else if (value >   0.030)  {  pc.setColor(223,250,35);  } // yellow
      else if (value >  -0.030)  {  pc.setColor(239,239,239); } // whiteish
      else if (value >  -0.090)  {  pc.setColor(148,247,138); } // green
      else if (value >  -0.150)  {  pc.setColor(82,246,221);  } // lt blue
      else if (value >  -0.210)  {  pc.setColor(36,59,227);   } // dk blue
      else if (value <= -0.210)  {  pc.setColor(80,4,184);    } // violet
   }

   cout << " ";
   pc.writePpm3(cout);
   cout << " ";
	   
   // repeat the cell for a square image
   cout << " ";
   pc.writePpm3(cout);
   cout << " ";

}



//////////////////////////////
//
// calculateAverage --
//

void calculateAverage(Array >& correlations, Array& a) {

   int i, j;
   double *aa;
   int rowsize;
   int corelsize;

   for (i=0; i<correlations.getSize()-1; i++) {
      rowsize = correlations[i].getSize();
      corelsize = a.getSize() - i;
      for (j=0; j<rowsize; j++) {
         aa = a.getBase() + j;
         correlations[i][j] = average(corelsize, aa);
      }
   }

   // set the bottom row to data
   i = correlations.getSize()-1;
   for (j=0; j<correlations[i].getSize(); j++) {
      correlations[i][j] = a[j];
   }

   // normalized the data into the range from -1 to 1 so that
   // the colorization of the pixels will work.

   normalizeAverageData(correlations);

}



//////////////////////////////
//
// normalizeAverageData -- transform the data range into -1 to +1 
// using the sigmoid function positions at the mean of the bottom row
// (i.e. the top item), and transition with equal to 4 standard 
// deviations (-2s to +2s) of the bottom row.
//

#define SIGMOID(x,c,w) (1.0/(1.0+exp(-(((x)-(c))/((w)/8)))))

void normalizeAverageData(Array >& correlations) {

   int i, j;
   double mean = 0.0;
   for (i=0; i<correlations[correlations.getSize()-1].getSize(); i++) {
      mean += correlations[correlations.getSize()-1][i];
   }
   mean = mean / correlations.getSize();

   double sd = getStandardDeviation(mean,correlations[correlations.getSize()-1]);

   double& center = mean;
   double  width  = 4 * sd;

   for (i=0; i<correlations.getSize(); i++) {
      for (j=0; j<correlations[i].getSize(); j++) {
         correlations[i][j] = SIGMOID(correlations[i][j], center, width)*2-1;
      }
   } 
}


//////////////////////////////
//
// getStandardDeviation --
//

double getStandardDeviation(double mean, Array& data) {
   double sum = 0.0;
   double value;
   int i;
   for (i=0; i<data.getSize(); i++) {
      value = data[i] - mean;
      sum += value * value;
   }
   
   return sqrt(sum / data.getSize());
}



//////////////////////////////
//
// calculateAverageDifference --
//

void calculateAverageDifference(Array<Array<double> >& correlations, 
      Array<double>& a, Array<double>& b) {

   int i, j;
   double *aa;
   double *bb;
   double aval = 1.0;
   double bval = 1.0;
   int rowsize;
   int corelsize;

   for (i=0; i<correlations.getSize(); i++) {
      rowsize = correlations[i].getSize();
      corelsize = a.getSize() - i;
      for (j=0; j<rowsize; j++) {
         aa = a.getBase() + j;
         bb = b.getBase() + j;
         aval = average(corelsize, aa);
         bval = average(corelsize, bb);
         correlations[i][j] = bval - aval;    // switch to b-a later

         if (condiffq) {
            // segment by tempo bands.
            correlations[i][j] = correlations[i][j] / aval;
         }
      }

   }


   // normalized the data into the range from -1 to 1 so that
   // the colorization of the pixels will work.

   //normalizeAverageData(correlations);


}



//////////////////////////////
//
// calculateCorrelations --
//

void calculateCorrelations(Array<Array<double> >& correlations,
      Array<double>& a, Array<double>& b) {

   int i, j;
   double *aa;
   double *bb;
   int rowsize;
   int corelsize;

   for (i=0; i<correlations.getSize()-1; i++) {
      rowsize = correlations[i].getSize();
      corelsize = a.getSize() - i;
      for (j=0; j<rowsize; j++) {
         aa = a.getBase() + j;
         bb = b.getBase() + j;
         correlations[i][j] = pearsonCorrelation(corelsize, aa, bb);
      }
   }

   // set the bottom row to 1.0 correlations
   i = correlations.getSize()-1;
   for (j=0; j<correlations[i].getSize(); j++) {
      correlations[i][j] = 1.0;
   }

}



//////////////////////////////
//
// average -- calculate the average of the input values.
//

double average(int size, double* data) {
   double sum = 0.0;

   int i;
   for (i=0; i<size; i++) {
      sum += data[i];
   }

   return sum / size;
}



//////////////////////////////
//
// pearsonCorrelation --
//

double pearsonCorrelation(int size, double* a, double* b) {
   double meana = getMean(size, a);
   double meanb = getMean(size, b);

   double topsum     = 0.0;
   double bottomasum = 0.0;
   double bottombsum = 0.0;

   int i;
   for (i=0; i<size; i++) {
      topsum += (a[i] - meana) * (b[i] - meanb);
      bottomasum += (a[i] - meana) * (a[i] - meana);
      bottombsum += (b[i] - meanb) * (b[i] - meanb);
   }

   if (bottomasum == 0.0 || bottombsum == 0.0) {
      return 0.0;
   }
   return topsum / sqrt(bottomasum * bottombsum);
}



//////////////////////////////
//
// getMean --
//

double getMean(int size, double* a) {
   if (size <= 0) {
      return 0.0;
   }

   int i;
   double sum = 0.0;
   for (i=0; i<size; i++) {
      sum += a[i];
   }

   return sum / size;
}



//////////////////////////////
//
// checkOptions -- 
//

void checkOptions(Options& opts, int argc, char* argv[]) {
   opts.define("g|log=b", "display vertical axis on a log scale");
   opts.define("flip=b", "flip the input data from tempo to duration");
   opts.define("d|difference=b", "display average difference plots");
   opts.define("A|arch=b", "calculate arch correlations");
   opts.define("R|ramp=b", "calculate ramp correlations");
   opts.define("cd|condiff=b", "display contoured difference plots");
   opts.define("f|logfactor=d:1500.0", "factor for strenth of log scaling");
   opts.define("n|column=i:1", "which column to average");
   opts.define("a|average=b", "display data average rather than correlation");
   opts.define("x|data=b",    "display input data and then quit");
   opts.define("c|correlation=b", "display correlation data and then quit");
   opts.define("h|hue=b", "display correlations in hue coloring");
   opts.define("r|rectangle=b", "display correlations in rectangular plot");
   opts.define("p|poly=b", "poly correlation plot");
   opts.define("b|bg|background=s:255 255 255", "background color");
   opts.define("height=i:-1", "Maximum height scope");

   opts.define("author=b",  "author of program"); 
   opts.define("version=b", "compilation info");
   opts.define("example=b", "example usages");   
   opts.define("help=b",  "short description");
   opts.process(argc, argv);
   
   // handle basic options:
   if (opts.getBoolean("author")) {
      cout << "Written by Craig Stuart Sapp, "
           << "craig@ccrma.stanford.edu, Jul 2006" << endl;
      exit(0);
   } else if (opts.getBoolean("version")) {
      cout << argv[0] << ", version: 8 Jul 2006" << endl;
      cout << "compiled: " << __DATE__ << endl;
      exit(0);
   } else if (opts.getBoolean("help")) {
      usage(opts.getCommand());
      exit(0);
   } else if (opts.getBoolean("example")) {
      example();
      exit(0);
   }

   logq      = opts.getBoolean("log");
   logfactor = opts.getDouble("logfactor");
   flipq     = opts.getBoolean("flip");
   averageq  = opts.getBoolean("average");
   diffq     = opts.getBoolean("difference");
   condiffq  = opts.getBoolean("condiff");
   column    = opts.getInteger("column") - 1;
   polyq     = opts.getBoolean("poly");
   archq     = opts.getBoolean("arch");
   maxheight = opts.getInteger("height");
   rampq     = opts.getBoolean("ramp");
   if (column < 0) {
      column = 0;
   } else if (column > 1) {
      column = 1;
   }
   dataq     = opts.getBoolean("data");
   correlq   = opts.getBoolean("correlation");
   if (opts.getBoolean("hue")) {
      colortype = COLOR_HUE;
   }
   
   if (opts.getBoolean("rectangle"))  {
      plotshape = SHAPE_RECTANGLE;
   }
}



//////////////////////////////
//
// makeSineArch --  generate a sqrt sine half-cycle arch of the
//     specified length.
//

void makeSineArch(Array& arch, int size) {
   arch.setSize(size);
   arch.allowGrowth(0);

   if (size == 1) {
      arch[0] = 0;
      return;
   }

   int i;
   double freq = 2.0 * M_PI / ((size-1) * 2.0);
   for (i=0; i<size-1; i++) {
      arch[i] = sin(freq * i);
   }
   arch[size-1] = 0.0;

}



//////////////////////////////
//
// makeRamp --
//

void makeRamp(Array& ramp, int size) {
   ramp.setSize(size);
   ramp.allowGrowth(0);

   if (size == 1) {
      ramp[0] = 0;
   }

   double factor = 1.0 / (size - 1.0);
   int i;
   for (i=0; i<size; i++) {
      ramp[i] = factor * i;
   }

}



//////////////////////////////
//
// calculateArchCorrelations --
//

void calculateArchCorrelations(Array<Array<double> >& correlations, 
      Array<double>& a) {

   Array<double> arch;

   int i, j;
   double *aa;
   double *bb;
   int rowsize;
   int corelsize;

   for (i=0; i<correlations.getSize()-1; i++) {
      makeSineArch(arch, i+1);
      rowsize = correlations[i].getSize();
      corelsize = a.getSize() - i;
      makeSineArch(arch, corelsize);
      for (j=0; j<rowsize; j++) {
         aa = a.getBase() + j;
         bb = arch.getBase();
         correlations[i][j] = pearsonCorrelation(corelsize, aa, bb);
      }
   }

   // set the bottom row to 1.0 correlations
   i = correlations.getSize()-1;
   for (j=0; j<correlations[i].getSize(); j++) {
      correlations[i][j] = 1.0;
   }

}



//////////////////////////////
//
// calculateRampCorrelations --
//

void calculateRampCorrelations(Array<Array<double> >& correlations, 
      Array<double>& a) {

   Array<double> ramp;

   int i, j;
   double *aa;
   double *bb;
   int rowsize;
   int corelsize;

   for (i=0; i<correlations.getSize()-1; i++) {
      rowsize = correlations[i].getSize();
      corelsize = a.getSize() - i;
      makeRamp(ramp, corelsize);
      for (j=0; j<rowsize; j++) {
         aa = a.getBase() + j;
         bb = ramp.getBase();
         correlations[i][j] = pearsonCorrelation(corelsize, aa, bb);
      }
   }

   // set the bottom row to 1.0 correlations
   i = correlations.getSize()-1;
   for (j=0; j<correlations[i].getSize(); j++) {
      correlations[i][j] = 1.0;
   }
}




//////////////////////////////
//
// example --
//

void example(void) {


}



//////////////////////////////
//
// usage --
//

void usage(const char* command) {

}



// md5sum: c3eecd281967728bcde5ba0c6b42d012 archetype.cpp [20080227]