//
// 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 "Options.h"
#include "PixelColor.h"

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

#include <iostream>
#include <fstream>
#include <string>

using namespace std;

#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 string& command);
void     calculateAverage      (vector<vector<double> >& correlations,
                                vector<double>& a);
void     calculateCorrelations (vector<vector<double> >& correlations,
                                vector<double>& a,
                                vector<double>& b);
double   pearsonCorrelation    (int size, double* a, double* b);
double   getMean               (int size, double* a);
void     printImageTriangle    (vector<vector<double> >& correlation);
void     printImageRectangle   (vector<vector<double> >& correlation);
void     printCorrelationPixel (double value, int style);
double   average               (int size, double* data);
void     normalizeAverageData  (vector<vector<double> >& correlations);
double   getStandardDeviation  (double mean, vector<double>& data);
int      logScale              (int level, int maxlevel, double factor);
void     calculateAverageDifference(vector<vector<double> >& correlations,
                                vector<double>& a,
                                vector<double>& b);

void     makeSineArch          (vector<double>& arch, int size);
void     makeRamp              (vector<double>& ramp, int size);
void     calculateArchCorrelations(vector<vector<double> >& correlations,
                                vector<double>& a);
void     calculateRampCorrelations(vector<vector<double> >& correlations,
                                vector<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;

	vector<double> a;
	vector<double> b;

	a.reserve(100000);
	b.reserve(100000);

	ifstream infile;
	string 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.push_back(aval);
				b.push_back(bval);
			} else if (sscanf(templine, "%lf", &aval) == 1) {
				if (flipq) {
					aval = 60.0 / aval;
				}
				a.push_back(aval);
				b.push_back(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.push_back(aval);
					b.push_back(bval);
				}
			}
		}

	}

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

	vector<vector<double> > correlations;
	correlations.resize(a.size());
	for (i=0; i<(int)a.size(); i++) {
		correlations[i].resize(i+1);
	}

	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<(int)correlations.size(); i++) {
			for (j=0; j<(int)correlations[i].size(); j++) {
				cout << correlations[i][j];
				if (j < (int)correlations[i].size()-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(vector<vector<double> >& correlation) {
	string background = options.getString("background");

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

	int startrow = 0;
	if (maxheight > 0 && maxheight < (int)correlation.size()) {
		startrow = (int)correlation.size() - 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<(int)correlation.size(); i++) {
		ii = i;
		if (logq) {
			ii = logScale(i, correlation.size(), logfactor);
		}
		datawidth = correlation[ii].size();
		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(vector<vector<double> >& correlation) {

	int maxcolumn = correlation.size() * 2;
	int maxrow    = ((int)correlation.size()-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<(int)correlation.size()-1; i++) {
		datawidth = correlation[i].size();

		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(vector<vector<double> >& correlations, vector<double>& a) {

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

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

	// set the bottom row to data
	i = (int)correlations.size()-1;
	for (j=0; j<(int)correlations[i].size(); 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(vector<vector<double> >& correlations) {

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

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

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

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


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

double getStandardDeviation(double mean, vector<double>& data) {
	double sum = 0.0;
	double value;
	int i;
	for (i=0; i<(int)data.size(); i++) {
		value = data[i] - mean;
		sum += value * value;
	}

	return sqrt(sum / data.size());
}



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

void calculateAverageDifference(vector<vector<double> >& correlations,
		vector<double>& a, vector<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<(int)correlations.size(); i++) {
		rowsize = correlations[i].size();
		corelsize = (int)a.size() - i;
		for (j=0; j<rowsize; j++) {
			aa = a.data() + j;
			bb = b.data() + 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(vector<vector<double> >& correlations,
		vector<double>& a, vector<double>& b) {

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

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

	// set the bottom row to 1.0 correlations
	i = (int)correlations.size()-1;
	for (j=0; j<(int)correlations[i].size(); 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(vector<double>& arch, int size) {
	arch.resize(size);

	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(vector<double>& ramp, int size) {
	ramp.resize(size);

	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(vector<vector<double> >& correlations,
		vector<double>& a) {

	vector<double> arch;

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

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

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

}



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

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

	vector<double> ramp;

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

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

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




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

void example(void) {


}



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

void usage(const string& command) {

}



// md5sum: c3eecd281967728bcde5ba0c6b42d012 archetype.cpp [20080227]