#include <iostream>
#include <fstream>
#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>

using namespace std;
using namespace cv;

/** classification header **/
#define NUM_ITERATIONS 10000
#define STEP_SIZE 0.4

/*
Comment on the effect of the stepsize on the convergence of optimization
----
The STEP_SIZE has a great effect on the speed of convergence. To make a fair comparison, we will keep the number of iterations NUM_ITERATIONS constant at 1e4.
Starting with small stepsize of 0.1 yields a loss of 0.4139, which decreases as we increment the stepsize by 0.1 (0.3663, 0.3470, 0.3432). For stepsizes of 0.35 and above the loss increases again, and as it turns out quite drastically so: (0.5 ~ 0.362542, 0.6 ~ 0.390917, 0.7 ~ 0.427088, 0.8 ~ 0.4685, 0.9 ~ 0.468582, 1.0 ~ 0.544636).
This behavior was somewhat expected and is most likely caused by overfitting to the training set.
*/

struct ClassificationParam{
    string posTrain, negTrain;
    string posTest, negTest;
};

// regression class for various regression methods
class LogisticRegression {

private:
    Mat train, test;                // each column is a feature vector
    Mat gtLabelTrain, gtLabelTest;  // row vector
    Mat phi;

    int loadFeatures(std::string& trainFile, std::string& testFile, Mat& feat, Mat& gtLabel);
	Mat calculateGradient(Mat &X, Mat &Y, Mat &phi);
	double calculateLoss(Mat &X, Mat &Y, Mat &phi);

public:
    LogisticRegression(ClassificationParam& param);
    int learnClassifier();
    int testClassifier();
    float sigmoid(float a);
	Mat sigmoid(Mat a);
    ~LogisticRegression(){}
};

/** regression header **/
#define FIN_RBF_NUM_CLUST 300
#define RBF_SIGMA 1e-3

// reading input parameters
struct RegressionParam{
    std::string regressionTrain;
    std::string regressionTest;
};

// models for regression
class Model{
public:
    Mat phi;        // each row models wi
    Mat sigma_sq;   // column vector
    Mat codeBook;   // codebook for finite kernel reg.
};

// regression class for various regression methods
class Regression{
private:
    Mat   trainx, trainw;
    Mat   testx, testw;
    Model linear_reg, fin_rbf_reg, dual_reg;

    int loadFeatures(std::string& fileName, Mat& vecx, Mat& vecw);

public:
    Regression(RegressionParam& param);
    ~Regression(){}
    int trainLinearRegression(); // TODO implement
    int trainFinite_RBF_KernelRegression(); // TODO implement
    int trainDualRegression(); // TODO implement

    int testLinearRegresssion(); // TODO implement
    int testFinite_RBF_KernelRegression(); // TODO implement
    int testDualRegression(); // TODO implement

};

int main() {
    RegressionParam rparam;
    rparam.regressionTrain = "../../../data/regression_train.txt";
    rparam.regressionTest  = "../../../data/regression_test.txt";

    Regression reg(rparam);

    // linear regression
    //reg.trainLinearRegression();
    //reg.testLinearRegresssion();
    //reg.trainFinite_RBF_KernelRegression();
    //reg.testFinite_RBF_KernelRegression();
    //reg.trainDualRegression();
    //reg.testDualRegression();

    ClassificationParam cparam;
    cparam.posTrain = "../../../data/bottle_train.txt";
    cparam.negTrain = "../../../data/horse_train.txt";
    cparam.posTest  = "../../../data/bottle_test.txt";
    cparam.negTest  = "../../../data/horse_test.txt";

    LogisticRegression cls(cparam);
    cls.learnClassifier();
    cls.testClassifier();

    return 0;
}

/** classification functions **/
LogisticRegression::LogisticRegression(ClassificationParam& param) {
    loadFeatures(param.posTrain,param.negTrain,train,gtLabelTrain);
    loadFeatures(param.posTest,param.negTest,test,gtLabelTest);
}

int LogisticRegression::loadFeatures(string& trainPos, string& trainNeg, Mat& feat, Mat& gtL){

    ifstream iPos(trainPos.c_str());
    if(!iPos) {
        cout<<"error reading train file: "<<trainPos<<endl;
        exit(-1);
    }
    ifstream iNeg(trainNeg.c_str());
    if(!iNeg) {
        cout<<"error reading test file: "<<trainNeg<<endl;
        exit(-1);
    }

    int rPos, rNeg, cPos, cNeg;
    iPos >> rPos;
    iPos >> cPos;
    iNeg >> rNeg;
    iNeg  >> cNeg;

    if(cPos != cNeg){
        cout<<"Number of features in pos and neg classes unequal"<<endl;
        exit(-1);
    }
    feat.create(cPos+1,rPos+rNeg,CV_32F); // each column is a feat vect
    gtL.create(1,rPos+rNeg,CV_32F);       // row vector


    // load positive examples
    for(int idr=0; idr<rPos; ++idr){
        gtL.at<float>(0,idr) = 1;
        feat.at<float>(0,idr) = 1;
        for(int idc=0; idc<cPos; ++idc){
            iPos >> feat.at<float>(idc+1,idr);
        }
    }

    // load negative examples
    for(int idr=0; idr<rNeg; ++idr){
        gtL.at<float>(0,rPos+idr) = 0;
        feat.at<float>(0,rPos+idr) = 1;
        for(int idc=0; idc<cNeg; ++idc){
            iNeg >> feat.at<float>(idc+1,rPos+idr);
        }
    }

    iPos.close();
    iNeg.close();

    return 0;
}

float LogisticRegression::sigmoid(float a){
    return 1.0f/(1+exp(-a));
}

Mat LogisticRegression::sigmoid(Mat a) {
	// Apply Sigmoid elementwise to all elements of matrix a
	for (auto y = 0u; y < a.rows; y++) {
		for (auto x = 0u; x < a.cols; x++) {
			a.at<float>(y, x) = sigmoid(a.at<float>(y, x));
		}
	}
	return a;
}


double LogisticRegression::calculateLoss(Mat &X, Mat &Y, Mat &phi) {

	Mat h, h_log, OneMinusTerm_log;
	// Sigmoid of the weighted sum
	h = sigmoid(phi * X);

	// Apply logarithm elementwise to h and 1-h respectively
	log(h, h_log);
	log(1 - h, OneMinusTerm_log);

	// Calculate Loss according to the function from the lecture
	double loss = (-1.0 / X.cols) * sum(Y.mul(h_log) + (1 - Y).mul(OneMinusTerm_log))[0];
	return loss;
}

Mat LogisticRegression::calculateGradient(Mat &X, Mat &Y, Mat &phi) {

	Mat h, gradient;
	// Sigmoid of the weighted sum
	h = sigmoid(phi * X);
	// Calulate gradient
	gradient = X * (h - Y).t();
	return gradient;
}

int LogisticRegression::learnClassifier() {

	// Placeholder for the weights
	phi = Mat(1, train.rows, CV_32F, 0.0);

	cout << "Starting logistic regression training with a learning rate of " << STEP_SIZE << " ..." << endl;
	cout << "-----" << endl;

	// Initialize the loss
	double initialLoss = FLT_MAX;
	// For every iteration...
	for (auto s = 0u; s < NUM_ITERATIONS; s++) {
		// ... Take a step in the direction of the negative gradient where the stepsize with which to take that step was determined earlier and stays constant over all iterations
		phi = phi - STEP_SIZE * calculateGradient(train, gtLabelTrain, phi).t();
		// Calculate loss to evaluate the training performance, which could be visualized somehow
		double loss = calculateLoss(train, gtLabelTrain, phi);
		if (s % 1000 == 0) {
			cout << "Loss on training set after iteration " << s << ": " << loss << endl;
		}
	}
	return 0;
}

int LogisticRegression::testClassifier() {

	Mat h;
	// Sigmoid of the weighted sum
	h = sigmoid(phi * test);

	cout << "Starting logistic regression testing ..." << endl;
	cout << "-----" << endl;
	int testSize = gtLabelTest.cols;
	float falsePredictions = 0;
	// For every test sample ...
	for (auto i = 0u; i < gtLabelTest.cols; i++) {
		// Binarize the results by rounding (values above 0.5 become 1, and values below become 0)
		float prediction = cvRound(h.at<float>(0, i));
		float teacher = gtLabelTest.at<float>(0, i);

		if (prediction != teacher)
			falsePredictions++;
	}

	float accuracy = (testSize - falsePredictions) / testSize;
	cout << "Loss on test set: " << calculateLoss(test, gtLabelTest, phi) << endl;
	cout << "Accuracy on test set: " << accuracy << endl;

	return 0;
}


/** regression functions **/
Regression::Regression(RegressionParam& param){
    // load features
    loadFeatures(param.regressionTrain,trainx,trainw);
    loadFeatures(param.regressionTest,testx,testw);
//    cout<<"features loaded successfully"<<endl;

    // model memory
    linear_reg.phi.create(trainx.rows,trainw.rows,CV_32F); linear_reg.phi.setTo(0);
    linear_reg.sigma_sq.create(trainw.rows,1,CV_32F); linear_reg.sigma_sq.setTo(0);
    fin_rbf_reg.phi.create(FIN_RBF_NUM_CLUST,trainw.rows,CV_32F);
    fin_rbf_reg.sigma_sq.create(trainw.rows,1,CV_32F);
    dual_reg.phi.create(trainx.cols,trainw.rows,CV_32F);
    dual_reg.sigma_sq.create(trainw.rows,1,CV_32F);

}
int Regression::loadFeatures(string& fileName, Mat& matx, Mat& matw){

    // init dimensions and file
    int numR, numC, dimW;
    ifstream iStream(fileName.c_str());
    if(!iStream){
        cout<<"cannot read feature file: "<<fileName<<endl;
        exit(-1);
    }

    // read file contents
    iStream >> numR;
    iStream >> numC;
    iStream >> dimW;
    matx.create(numC-dimW+1,numR,CV_32F); // each column is a feature
    matw.create(dimW,numR,CV_32F);        // each column is a vector to be regressed

    for(int r=0; r<numR; ++r){
        // read world data
        for(int c=0; c<dimW; ++c)
            iStream >> matw.at<float>(c,r);
        // read feature data
        matx.at<float>(0,r)=1;
        for(int c=0; c<numC-dimW; ++c)
            iStream >> matx.at<float>(c+1,r);
    }
    iStream.close();

    return 0;
}