// JNI_PMI.cpp 
//

#include "pmi_MI.h"
#include <iostream>
#include <cmath>
#include <stdio.h>
#include <cstdlib>
#include <algorithm>
#include <vector>
//#include <kdtree++\kdtree.hpp>


using namespace std;

static double* Fai = NULL;

static void destory();
int countIfSmaller(double*, double, int);
int* sort_idx(double*, int);
double calculatePMI(double**, double**, double**, int, int, int, int, int);
double** getDistanceMatrix(double**, int, int);
double getDist(double*, double**, int);

//template <typename V, typename K>
bool compare_pair(const pair<double,int> & a, const pair<double,int> & b){
	return (a.first < b.first);
}

inline double max(double a, double b){
	return a < b ? b : a;
}
inline double min(double a, double b){
	return a < b ? a : b;
}

double calcDist(double**, int, int, int, int);




/**
 * 	public static node[] sort_index(double[] a){
		node[] ar = new node[a.length];
		for (int i = 0; i <  a.length; i++){
			ar[i] = new node(a[i],i);
		}

		Arrays.sort(ar, (x,y)->{double c = x.v-y.v;if (c>0) return 1; else return -1;});

		return ar;
	}
 * */


int countIfSmaller(double* a, double b, int length){
	int c = 0;
	for (int i = 0; i < length; i++){
		if (a[i] < b){
			c++;
		}
	}
	return c;
}

int* sort_idx(double* a, int length){
	vector<pair<double,int> > ar;

	for (int i = 0; i < length; i++){

		pair<double,int> ts(a[i], i);
		ar.push_back(ts);

	}
	sort(ar.begin(), ar.end(), compare_pair/*<double, int>*/);
	int* ret = new int[ar.size()];
	int ccc = 0;
	for (vector<pair<double, int> >::iterator it = ar.begin(); it != ar.end(); it++){
		ret[ccc++] = it->second;
	}

	return ret;
}

double getDist(double* po1, double** arr , int n, int width){
	double dff = 0.0;
	for (int i = 0; i < width; i++){
		double k = abs(arr[i][n] - po1[i]);
		dff = k > dff ? k : dff;
	}
	return dff;
}

double* getPos(double** a1, double** a2, int n, int da1, int da2){
	
	double* ret = new double[da1+da2];
	for (int i = 0; i < da1; i++){
		ret[i] = a1[i][n];
	}
	for (int i = da1; i < da1+da2; i++){
		ret[i] = a2[i-da1][n];
	}
	return ret;
}


double calculatePMI(double** a, double** b, double** c, int k, int N, int da, int db, int dc){
	//cout << "k=" << k << ", N=" << N << ", da=" << da << ", db=" << db << ", dc=" << dc << endl;
	//cout << a[da - 1][N - 1] << ", " << b[db - 1][N - 1] << ", " << c[dc - 1][N - 1] << endl;
	double** Mabc = new double*[da+db+dc];
	double** Mac = new double*[da+dc];
	double** Mbc = new double*[db+dc];
	double** Mc = c;

	int mabc_c = 0;
	int mac_c = 0;
	int mbc_c = 0;
	for (int i = 0; i < da; i++){
		Mabc[mabc_c++] = a[i];
		Mac[mac_c++] = a[i];
	}
	for (int i = 0; i < db; i++){
		Mabc[mabc_c++] = b[i];
		Mbc[mbc_c++] = b[i];
	}
	for (int i = 0; i < dc; i++){
		Mabc[mabc_c++] = c[i];
		Mac[mac_c++] = c[i];
		Mbc[mbc_c++] = c[i];
	}
	double** dMac = getDistanceMatrix(Mac,da+dc, N);
	double** dMbc = getDistanceMatrix(Mbc,db+dc, N);
	double** dMc = getDistanceMatrix(Mc, dc, N);
	double** dMabc = getDistanceMatrix(Mabc,da+dc+db, N);

	double* epiac = new double[N];
	double* epibc = new double[N];
	double* epic  = new double[N];
	double* epsilon = new double[N];

	for (int i = 0; i < N; i++){
		int* pos = sort_idx(dMabc[i], N);
		int posK = pos[k];
		epsilon[i] = dMabc[i][posK];
		epiac[i] = dMac[i][posK];
		epibc[i] = dMbc[i][posK];
		epic[i] = dMc[i][posK];
		delete[] pos;
	}

	double aveac = 0,
			avebc = 0,
			avec = 0;

	if (Fai == NULL){
		Fai = new double[N+5]; //remember to delete it!!!!!!!!
		Fai[1] = Fai[0] = -0.57721566490153286060651209008240243104215933593992;
		for (int i = 2; i < N+5; i++){
			Fai[i] = Fai[i-1] + 1.0 / (i-1);
		}
		cout << "Fai Calculation" << endl;

	}

	for (int i = 0; i < N; i++){
		aveac += Fai[countIfSmaller(dMac[i], epiac[i], N)] / N;
		avebc += Fai[countIfSmaller(dMbc[i], epibc[i], N)] / N;
		avec  += Fai[countIfSmaller( dMc[i], epic[i], N)] / N;
	}
	delete[] epsilon;
	delete[] epiac;
	delete[] epibc;
	delete[] epic;
        //Mabc, Mac, Mbc, dMabc, dMac, dMbc, dMc,	
	/*	for (int i = 0; i < da+db+dc; i++)
	  delete[] Mabc[i];
	for (int i = 0; i < da+dc; i++)
	  delete[] Mac[i];
	for (int i = 0; i < db+dc; i++)
	delete[] Mbc[i];*/
	for (int i = 0; i < N; i++){
	  delete[] dMabc[i]; delete[] dMac[i];delete[] dMbc[i];delete[] dMc[i];
	}

	delete[] Mabc;delete[] Mac;delete[] Mbc;delete[] dMabc;delete[] dMac;delete[] dMbc;delete[] dMc;
	return Fai[k] - aveac - avebc + avec;


}
double** getDistanceMatrix(double** a,int dimension, int length){
	double** ret = new double*[length];
	for (int i = 0; i < length; i++){
		ret[i] = new double[length];
	}

	for (int i = 0; i < length; i++){
		ret[i][i] = 0;
		for (int j = i+1; j < length; j++){
			ret[i][j] = calcDist(a, i, j, dimension, length);
			ret[j][i] = ret[i][j];
		}
	}


	return ret;
}

inline double calcDist(double** a, int i, int j, int dimension, int length){
	double r = 0.0;
	for (int k = 0; k < dimension; k++){
		r = max(r, abs(a[k][i] - a[k][j]));
//		r += (a[k][i] - a[k][j])*(a[k][i] - a[k][j]);
	}
	return r;
//	return sqrt(r);
}

static void destory(void){
	delete[] Fai;
}










JNIEXPORT jdouble JNICALL Java_pmi_MI_partialMutualInformation(JNIEnv *env, jclass Jclass, jobjectArray a, jobjectArray b, jobjectArray c, jint k){
	const double bad_value = (jdouble)(99999);

	if (a == NULL || b == NULL || c == NULL || k < 1){
		return bad_value; //very big value
	}

	int lengthA = env->GetArrayLength(a);
	int lengthB = env->GetArrayLength(b);
	int lengthC = env->GetArrayLength(c);

	if (lengthA == 0 || lengthB == 0 || lengthC == 0){
		return bad_value;
	}

	jdoubleArray* jarrayA = new jdoubleArray[lengthA];
	jdoubleArray* jarrayB = new jdoubleArray[lengthB];
	jdoubleArray* jarrayC = new jdoubleArray[lengthC];
	for (int i = 0; i < lengthA; i++){
		jarrayA[i] = static_cast<jdoubleArray>(env->GetObjectArrayElement(a, i));
	}
	jsize widthA = env->GetArrayLength(jarrayA[0]);

	for (int i = 0; i < lengthB; i++){
		jarrayB[i] = static_cast<jdoubleArray>(env->GetObjectArrayElement(b, i));
	}
	jsize widthB = env->GetArrayLength(jarrayB[0]);

	for (int i = 0; i < lengthC; i++){
		jarrayC[i] = static_cast<jdoubleArray>(env->GetObjectArrayElement(c, i));
	}
	jsize widthC = env->GetArrayLength(jarrayC[0]);



	if (widthA != widthB || widthA != widthC || widthB != widthC){
		return bad_value;
	}
	double** A = new double*[lengthA];
	double** B = new double*[lengthB];
	double** C = new double*[lengthC];

	//for (int i = 0; i < lengthA; i++){
	//	A[i] = new double[widthA];
	//}
	//for (int i = 0; i < lengthB; i++){
	//	B[i] = new double[widthB];
	//}
	//for (int i = 0; i < lengthC; i++){
	//	C[i] = new double[widthC];
	//}
	jboolean* Aflag = new jboolean[lengthA];//flags of whether I need to release the memory manually
	jboolean* Bflag = new jboolean[lengthB];
	jboolean* Cflag = new jboolean[lengthC];
	
	fill(&Aflag[0], &Aflag[lengthA], JNI_FALSE);
	fill(&Bflag[0], &Bflag[lengthB], JNI_FALSE);
	fill(&Cflag[0], &Cflag[lengthC], JNI_FALSE);
	for (int i = 0; i < lengthA; i++){		
		A[i] = env->GetDoubleArrayElements(jarrayA[i], &Aflag[i]);
	}
	for (int i = 0; i < lengthB; i++){		
		B[i] = env->GetDoubleArrayElements(jarrayB[i], &Bflag[i]);
	}
	for (int i = 0; i < lengthC; i++){		
		C[i] = env->GetDoubleArrayElements(jarrayC[i], &Cflag[i]);
	}
	//jdouble retVal = static_cast<jdouble>(0);
	jdouble retVal = (jdouble)calculatePMI(A, B, C, (int)k, widthA, lengthA, lengthB, lengthC);
	//int maxRow = max(lengthA, max(lengthB, lengthC));
	//
	//switch (maxRow){
	//case 1:
	//	retVal = static_cast<jdouble>(calculatePMI2<1>(A, B, C, (int)k, widthA, lengthA, lengthB, lengthC));
	//	break;
	//case 2:
	//	retVal = static_cast<jdouble>(calculatePMI2<2>(A, B, C, (int)k, widthA, lengthA, lengthB, lengthC));
	//	break;
	//default:
	//	if (maxRow < 5)
	//		retVal = static_cast<jdouble>(calculatePMI2<5>(A, B, C, (int)k, widthA, lengthA, lengthB, lengthC));
	//	else if (maxRow < 10)
	//		retVal = static_cast<jdouble>(calculatePMI2<10>(A, B, C, (int)k, widthA, lengthA, lengthB, lengthC));
	//	else
	//		retVal = static_cast<jdouble>(calculatePMI2<100>(A, B, C, (int)k, widthA, lengthA, lengthB, lengthC));
	//}

	/////////////////////////////////////////////
	//it seems that they are freed automatically

	for (int i = 0; i < lengthA; i++){
		env->ReleaseDoubleArrayElements(jarrayA[i], A[i], JNI_ABORT);
	}
	for (int i = 0; i < lengthB; i++){
		env->ReleaseDoubleArrayElements(jarrayB[i], B[i], JNI_ABORT);
	}
	for (int i = 0; i < lengthC; i++){
		env->ReleaseDoubleArrayElements(jarrayC[i], C[i], JNI_ABORT);
	}
	/////////////////////////////////////////////
	for (int i = 0; i < lengthA; i++){		
		env->DeleteLocalRef(jarrayA[i]);
	}
	for (int i = 0; i < lengthB; i++){		
		env->DeleteLocalRef(jarrayB[i]);
	}
	for (int i = 0; i < lengthC; i++){		
		env->DeleteLocalRef(jarrayC[i]);
	}

	delete[] jarrayA;
	delete[] jarrayB;
	delete[] jarrayC;

	

	delete[] A;delete[] B; delete[] C;

	delete[] Aflag;
	delete[] Bflag;
	delete[] Cflag;


	return retVal;
}

JNIEXPORT void JNICALL Java_pmi_MI_dispose(JNIEnv *, jclass){
	destory();
}

JNIEXPORT jdouble JNICALL Java_pmi_MI_mutualInformation(JNIEnv *, jclass, jobjectArray, jobjectArray, jint){
	return (jdouble)0;
}


//int _tmain(){
//	double** a = new double*[1];
//	double** b = new double*[1];
//	double** c = new double*[2];
//	a[0] = new double[400];
//	b[0] = new double[400];
//	c[0] = new double[400];
//	c[1] = new double[400];
//	for (int i = 0; i < 400; i++){
//		a[0][i] = sin(i);
//		b[0][i] = cos(i);
//		c[0][i] = tan(i);
//		c[1][i] = 0;
//	}
//	cout << calculatePMI(a,b,c,64, 400,1,1,2) << endl;
//	return 0;
//
//}