#include "CasicSegPostProcess.h"

#include <math.h>
#include <QDateTime>
#include <QString>
#include "utils/LogUtil.h"

namespace iristrt
{

std::vector<cv::Mat> CasicSegPostProcess::getConnections(cv::Mat input, int connectivity) {
	cv::Mat output;
    int nLabels = cv::connectedComponents(input, output, connectivity);
	output.convertTo(output, CV_8UC1, 1);

	std::vector<cv::Mat> conns;
	for (int i = 1; i <= nLabels; i++) {
		cv::Mat dst = cv::Mat::zeros(output.size(), output.type());
		for (int row = 0; row < output.rows; row++) {
			for (int col = 0; col < output.cols; col++) {
				uchar label = output.at<uchar>(row, col);
				if (label == i) {
					dst.at<uchar>(row, col) = 255;
				}
			}
		}
		conns.push_back(dst);
	}
	return conns;
}

void CasicSegPostProcess::showConnections(std::vector<cv::Mat> conns) {
	for (int i = 0; i < conns.size(); i++) {
		cv::imshow("con", conns[i]);
		cv::waitKey(0);
	}
	cv::destroyAllWindows();
}

std::vector<CasicTriplet>  CasicSegPostProcess::getTripleSet(std::vector<cv::Mat> maskConns,
															std::vector<cv::Mat> irisConns, 
															std::vector<cv::Mat> pupilConns,
															int chessboardDistance) {
    std::vector<CasicTriplet> triplet_set;
    for (auto maskConn : maskConns) {
        for (auto irisConn : irisConns) {
			int maskIrisDis = getChessboardDistance(maskConn, irisConn, chessboardDistance);
            for (auto pupilConn : pupilConns) {
				int maskPupilDis = getChessboardDistance(maskConn, pupilConn, chessboardDistance);
				if (maskIrisDis <= chessboardDistance && maskPupilDis <= chessboardDistance) {
                    CasicTriplet triplet;
                    triplet.maskConn = maskConn;
                    triplet.irisConn = irisConn;
                    triplet.pupilConn = pupilConn;
					triplet_set.push_back(triplet);
				}
			}
		}
	}
	return triplet_set;
}

void CasicSegPostProcess::showTriple(CasicTriplet triple) {
	/*cv::Mat tripleImage;
	cv::add(std::get<0>(triple), std::get<1>(triple), tripleImage);
	cv::add(std::get<2>(triple), tripleImage, tripleImage);
	cv::imshow("tripleImage", tripleImage);
	cv::waitKey(0);*/

    cv::Mat mask = triple.maskConn;
    cv::Mat iris = triple.irisConn;
    cv::Mat pupil = triple.pupilConn;

	cv::imshow("triple-mask", mask);
	cv::waitKey(0);

	cv::imshow("triple-iris", iris);
	cv::waitKey(0);

	cv::imshow("triple-pupil", pupil);
	cv::waitKey(0);
}

int CasicSegPostProcess::getChessboardDistance(cv::Mat con1, cv::Mat con2, int chessboardDistance) {
	// mask contours
	std::vector<std::vector<cv::Point>> c1;
    cv::findContours(con1, c1,cv:: RETR_TREE, cv::CHAIN_APPROX_SIMPLE, cv::Point());
	//cout << "c1.size() " << c1.size() << endl;
	std::vector<cv::Point> c1All;
	for (auto c : c1) {
		//cout << "c.size()" << c.size() << endl;
		c1All.insert(c1All.end(), c.begin(), c.end());
	}
	//cout << "c1All.size()	" << c1All.size() << endl;

	// other contours
	std::vector<std::vector<cv::Point>> c2;
    cv::findContours(con2, c2, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE, cv::Point());
	//cout << "c2.size() " << c2.size() << endl;
	std::vector<cv::Point> c2All;
	for (auto c : c2) {
		//cout << "c.size()" << c.size() << endl;
		c2All.insert(c2All.end(), c.begin(), c.end());
	}
	//cout << "c2All.size()	" << c2All.size() << endl;

	int distance = 1000;
	for (auto p1 : c1All) {
		for (auto p2 : c2All) {
			int tempDistance = std::max(abs(p1.x - p2.x), abs(p1.y - p2.y));
			if (tempDistance < distance) {
				distance = tempDistance;
			}
			if (distance <= chessboardDistance) {
				break;
			}
		}
	}
	return distance;


}

CasicTriplet CasicSegPostProcess::getMaxTriple(std::vector<CasicTriplet> tripleSet) {
	long maxCount = 0;
	int maxIndex = -1;
	for (int i = 0; i < tripleSet.size(); i++) {
        CasicTriplet triplet = tripleSet[i];
        long count = cv::countNonZero(triplet.maskConn)
            + cv::countNonZero(triplet.irisConn)
            + cv::countNonZero(triplet.pupilConn);
		if (count > maxCount) {
			maxCount = count;
			maxIndex = i;
		}
	}
	return tripleSet[maxIndex];
}

std::vector<int> CasicSegPostProcess::leastSquareCircleFitting(std::vector<cv::Point> counter) {

	std::vector<int> result;
	int center_x, center_y, radius = 0;

	if (counter.size() < 3) {
		result.push_back(center_x);
		result.push_back(center_y);
		result.push_back(radius);
		return result;
	}

	double sum_x, sum_y = 0.0;
	double sum_x2, sum_y2 = 0.0;
	double sum_x3, sum_y3 = 0.0;
	double sum_xy, sum_x1y2, sum_x2y1 = 0.0;

	for (auto point : counter) {
		double x = (double)point.x;
		double y = (double)point.y;
		double x2 = x * x;
		double y2 = y * y;
		sum_x += x;
		sum_y += y;
		sum_x2 += x2;
		sum_y2 += y2;
		sum_x3 += x2 * x;
		sum_y3 += y2 * y;
		sum_xy += x * y;
		sum_x1y2 += x * y2;
		sum_x2y1 += x2 * y;
	}

	double N = counter.size();
	double	C = N * sum_x2 - sum_x * sum_x;
	double	D = N * sum_xy - sum_x * sum_y;
	double	E = N * sum_x3 + N * sum_x1y2 - (sum_x2 + sum_y2) * sum_x;
	double	G = N * sum_y2 - sum_y * sum_y;
	double	H = N * sum_x2y1 + N * sum_y3 - (sum_x2 + sum_y2) * sum_y;
	double	a = (H * D - E * G) / (C * G - D * D);
	double	b = (H * C - E * D) / (D * D - G * C);
	double	c = -(a * sum_x + b * sum_y + sum_x2 + sum_y2) / N;

	center_x = (int)(a / (-2));
	center_y = (int)(b / (-2));
	radius = (int)(sqrt(a * a + b * b - 4 * c) / 2);

	result.push_back(center_x);
	result.push_back(center_y);
	result.push_back(radius);
	return result;
}

bool CasicSegPostProcess::postProcess(cv::Mat & mask, cv::Mat & iris, cv::Mat & pupil, std::vector<int> & irisCircle, std::vector<int> & pupilCircle) {
	int chessboard_distance = 15;

	// cv::Mat mask,iris,pupil;
    mask.convertTo(mask, CV_8UC1);
    iris.convertTo(iris, CV_8UC1);
    pupil.convertTo(pupil, CV_8UC1);

    cv::threshold(mask, mask, 127, 255, cv::THRESH_BINARY);
    cv::threshold(iris, iris, 90, 255, cv::THRESH_BINARY);
    cv::threshold(pupil, pupil, 0, 255, cv::THRESH_BINARY + cv::THRESH_OTSU);

    cv::Mat element = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(5, 5));
    cv::morphologyEx(iris, iris, cv::MORPH_CLOSE, element);
/*
    std::string pupilSeg = QString("%1/%2/%3-pupil-thre.bmp").arg("/home/casic/irisLogs").arg(QDate::currentDate().toString("yyyyMMdd")).arg(QTime::currentTime().toString("HHmmss")).toStdString();
    std::string maskSeg = QString("%1/%2/%3-mask-thre.bmp").arg("/home/casic/irisLogs").arg(QDate::currentDate().toString("yyyyMMdd")).arg(QTime::currentTime().toString("HHmmss")).toStdString();
    std::string irisSeg = QString("%1/%2/%3-iris-thre.bmp").arg("/home/casic/irisLogs").arg(QDate::currentDate().toString("yyyyMMdd")).arg(QTime::currentTime().toString("HHmmss")).toStdString();
    cv::imwrite(pupilSeg, pupil);
    cv::imwrite(maskSeg, mask);
    cv::imwrite(irisSeg, iris);
*/
//    cv::imshow("mask",mask);
//    cv::imshow("iris",iris);
//    cv::imshow("pupil",pupil);
//    cv::waitKey(0);

//    std::cout << "get connection ..."  << std::endl;
	
	std::vector<cv::Mat> maskCons = getConnections(mask, 8);
	std::vector<cv::Mat> irisCons = getConnections(iris, 8);
	std::vector<cv::Mat> pupilCons = getConnections(pupil, 8);

	if(maskCons.size() <= 0 || irisCons.size() <= 0 || pupilCons.size() <= 0)
	{
		return false;
	}

//    std::cout << "get triple set ..."  << std::endl;

    std::vector<CasicTriplet> tripleSet = getTripleSet(maskCons, irisCons, pupilCons, chessboard_distance);

	if (tripleSet.size() < 1) {
		return false;
	}

//    std::cout << "get max triple ..." << std::endl;

    CasicTriplet maxTriple = getMaxTriple(tripleSet);
	//showTriple(maxTriple);

//    std::cout << "least Square Circle Fitting ..."  << std::endl;

    cv::Mat bestIris = maxTriple.irisConn;
	std::vector<cv::Point> outPoints;
    cv::findNonZero(bestIris, outPoints);
	irisCircle = leastSquareCircleFitting(outPoints);
	int iris_x = irisCircle[0];
	int	iris_y = irisCircle[1];
	int	iris_r = irisCircle[2];

//    LOG_INFO(QString("outPoints size: %1").arg(outPoints.size()).toStdString());

    cv::Mat bestPupil = maxTriple.pupilConn;
	std::vector<std::vector<cv::Point>> innerCounter;
    cv::findContours(bestPupil, innerCounter, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
	std::vector<cv::Point> innerPoints;

//    LOG_INFO(QString("innerPoints size: %1, %2").arg(innerPoints.size()).arg(innerCounter.size()).toStdString());

    for (auto c : innerCounter) {
		innerPoints.insert(innerPoints.end(), c.begin(), c.end());
//        LOG_INFO(QString("innerPoints: %1, %2, %3, %4").arg(c.begin()->x).arg(c.begin()->y).arg(c.end()->x).arg(c.end()->y).toStdString());
	}
	pupilCircle = leastSquareCircleFitting(innerPoints);
	int pupil_x = pupilCircle[0];
	int pupil_y = pupilCircle[1];
	int	pupil_r = pupilCircle[2];

//    LOG_INFO(QString("innerPoints size: %1").arg(innerPoints.size()).toStdString());
/*
    std::string pupilbest = QString("%1/%2/%3-pupil-best.bmp").arg("/home/casic/irisLogs").arg(QDate::currentDate().toString("yyyyMMdd")).arg(QTime::currentTime().toString("HHmmss")).toStdString();
    std::string irisbest = QString("%1/%2/%3-iris-best.bmp").arg("/home/casic/irisLogs").arg(QDate::currentDate().toString("yyyyMMdd")).arg(QTime::currentTime().toString("HHmmss")).toStdString();
    cv::imwrite(pupilbest, bestPupil);
    cv::imwrite(irisbest, bestIris);
*/
//    std::cout << "iris: " << iris_x << " " << iris_y << " " << iris_r << std::endl;
//    std::cout << "pupil:" << pupil_x << " " << pupil_y << " " << pupil_r << std::endl;
//    LOG_INFO(QString("iris: %1, %2, %3").arg(iris_x).arg(iris_y).arg(iris_r).toStdString());
//    LOG_INFO(QString("pupil: %1, %2, %3").arg(pupil_x).arg(pupil_y).arg(pupil_r).toStdString());

//    std::cout << "bitwise_and ..."  << std::endl;

	cv::Mat circleMask = cv::Mat::zeros(mask.size(), CV_8UC1);
    cv::circle(circleMask,cv::Point(iris_x, iris_y), iris_r, 255, -1);
    cv::circle(circleMask,cv::Point(pupil_x, pupil_y), pupil_r, 0, -1);
	cv::bitwise_and(mask, circleMask, mask);

	return true;
}

} // end of namespace iristrt