#include <iostream>
#include <dirent.h>
#include <opencv2/opencv.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv/cvaux.hpp>
#include <math.h>
#include <opencv2/nonfree/features2d.hpp>
#include <unordered_map>

using namespace std;
using namespace cv;

vector<string> bikesFiles;
vector<string> carsFiles;
vector<string> cowsFiles;
vector<unordered_map<int,float>> histograms;
CvNormalBayesClassifier classifier_motor;
CvNormalBayesClassifier classifier_car;
CvNormalBayesClassifier classifier_cow;

unordered_map<int,float> initHashTable(){
    unordered_map<int,float> hist;
    for(int i=0;i<100;i++){
        hist.emplace(i,0);
    }
return hist;
}

Mat subMat(Mat hist, int st_index, int length){
    Mat histogram_mat(length, 100, CV_32FC1);
    for(int i=0;i<length;i++){
        for(int j=0;j<100;j++){
            histogram_mat.at<float>(i,j)=hist.at<float>(st_index,j);
        }
        st_index++;
    }
    return histogram_mat;
}

Mat histogramToMat(){
    Mat histogram_mat(326, 100, CV_32FC1);
    for(int i=0;i<histograms.size();i++){
        for(int j=0;j<histograms[i].size();j++){
            histogram_mat.at<float>(i,j)=histograms[i].at(j);
        }
    }
    return histogram_mat;
}

Mat initResponse(int st_index,int length){
    Mat response = Mat::zeros(261, 1, CV_32S);
    for(int i=st_index;i<st_index+length;i++){
        response.at<int>(i,0)=1;
    }
    return response;
}

vector<Mat> readImages(char path[], int type){
    vector<Mat> images;
    DIR*     dir;
    dirent*  pdir;
    char directory[100];
    strcpy(directory, path);
    dir = opendir(directory);
    readdir(dir);
    readdir(dir);
    while ((pdir = readdir(dir))) {
        char name[100];
        strcpy(name, directory);
        strcat(name,pdir->d_name);
        Mat src = imread(name,0);
        images.push_back(src);
        string fileName = pdir->d_name;
        if(type == 1)
            bikesFiles.push_back(fileName.substr(0,fileName.length()-4));
        if(type == 2)
            carsFiles.push_back(fileName.substr(0,fileName.length()-4));
        if(type == 3)
            cowsFiles.push_back(fileName.substr(0,fileName.length()-4));
    }
    closedir(dir);
    return images;
}

void applyHistEqualization(vector<Mat>& original){
for(int i=0;i<original.size();i++){
    Mat copy;
    equalizeHist( original[i], copy);
    original[i]=copy;
}
}

void applyResizing(vector<Mat>& original){
for(int i=0;i<original.size();i++){
    Mat copy;
    resize(original[i], copy, Size(300, 250));
    original[i]=copy;
}
}

vector<vector<KeyPoint>> denseDetector(vector<Mat> original){
vector<vector<KeyPoint>> keypoints;
Ptr<FeatureDetector> detector(new DenseFeatureDetector(1.f,1,0.1f,5,0,true,false));
for(int i=0;i<original.size();i++){
    vector<KeyPoint> image_keypoints;
    detector->detect(original[i],image_keypoints);
    keypoints.push_back(image_keypoints);
}
return keypoints;
}

vector<Mat> extractor(vector<Mat> images, vector<vector<KeyPoint>> keypoints){
vector<Mat> extracted_keypoints;
SiftDescriptorExtractor sde;
for(int i=0;i<images.size();i++){
    Mat descriptor;
    sde.compute(images[i], keypoints[i], descriptor);
    extracted_keypoints.push_back(descriptor);
}
return extracted_keypoints;
}

vector<Mat> extractFromSiftFile(vector<Mat> images, const char *type, int type1, const char *extension){
vector<Mat> extracted_keypoints;
for(int i=0;i<images.size();i++){
    string path= "";
    if(type1 ==1){
        path =  "/Users/msalah93/Desktop/Sem9/ComputerVision/CVProject/" + (string)type + "/"+ bikesFiles[i]+ (string)extension;
    }
    if(type1 ==2){
        path =  "/Users/msalah93/Desktop/Sem9/ComputerVision/CVProject/" + (string)type + "/"+ carsFiles[i]+(string)extension;
    }
    if(type1 ==3){
        path =  "/Users/msalah93/Desktop/Sem9/ComputerVision/CVProject/" + (string)type + "/"+ cowsFiles[i]+(string)extension;
    }
    Mat descriptor;
    FileStorage fs2(path, FileStorage::READ);
    fs2["SIFT"] >> descriptor;
    extracted_keypoints.push_back(descriptor);

}
return extracted_keypoints;
}

void saveDescriptors(vector<Mat>& images, const char *type, int type1, const char *extension){
for(int i=0;i<images.size();i++){
    string path= "";
    if(type1 ==1){
        path =  "/Users/msalah93/Desktop/Sem9/ComputerVision/CVProject/" + (string)type + "/"+ bikesFiles[i]+(string)extension;
    }
    if(type1 ==2){
        path =  "/Users/msalah93/Desktop/Sem9/ComputerVision/CVProject/" + (string)type + "/"+ carsFiles[i]+(string)extension;
    }
    if(type1 ==3){
        path =  "/Users/msalah93/Desktop/Sem9/ComputerVision/CVProject/" + (string)type + "/"+ cowsFiles[i]+(string)extension;
    }
    FileStorage fs(path, FileStorage::WRITE);
    fs << "SIFT" << images[i];
}
}

vector<Mat> getRandomImages(vector<Mat> motorbikes,vector<Mat> cars,vector<Mat> cows){
vector<Mat> randomised;
vector<int> indeces;
for(int i = 0;i<75;i++){
    int index = 0;
    while(find(indeces.begin(), indeces.end(), index)!=indeces.end())
        index = rand()%326;
    indeces.push_back(index);
    if(index<115)
        randomised.push_back(motorbikes[index]);
    else{
        if(index<215)
            randomised.push_back(cars[index-115]);
        else{
            if(index<326)
                randomised.push_back(cows[index-215]);
        }
    }
}
return randomised;
}

Mat applyKmeans(Mat images){
Mat labels;
Mat centroids;
kmeans(images,100,labels,TermCriteria(TermCriteria::COUNT, 100, 2.0),1,KMEANS_RANDOM_CENTERS,centroids);
return centroids;
}

float getEuclideanDistance(vector<float> v1, vector<float> v2){
float dist = 0;
for(int i=0;i<v1.size();i++){
    dist+= pow((v1[i]-v2[i]),2);
}
return sqrt(dist);
}

int adjustCentroid(vector<float> desc, Mat centroids){
int index = 0;
float min_dist = numeric_limits<float>::max();
for(int i=0;i<centroids.rows;i++){
//        vector<float> v;
//        centroids.row(i).copyTo(v);
    float dist_tmp = getEuclideanDistance(desc, centroids.row(i));
    if(dist_tmp<min_dist){
        min_dist = dist_tmp;
        index = i;
    }
}
return index;
}

void initialiseBoW(vector<Mat> images, Mat centroids){
for(int i=0;i<images.size();i++){
    unordered_map<int,float> hist = initHashTable();
    cout << hist.size() << endl;
    for(int j=0;j<images[i].rows;j++){
//            vector<float> v;
//            images[i].row(j).copyTo(v);
        int cent_index = adjustCentroid(images[i].row(j), centroids);
        hist.at(cent_index)= hist.at(cent_index)+1;
    }
    cout << "histogram list: " << histograms.size() << endl;
    histograms.push_back(hist);
}
}

void trainDataSet(Mat trainData){
//    Mat trainData = histogramToMat();
Mat motor_train = subMat(trainData, 0, 92);
Mat car_train = subMat(trainData, 115, 80);
Mat cow_train = subMat(trainData, 215, 89);
vector<Mat> trainTmp= {motor_train,car_train,cow_train};
Mat trainDataFinal;
vconcat(trainTmp, trainDataFinal);
Mat motor_bikes = initResponse(0, 92);
Mat cars = initResponse(92, 80);
Mat cows = initResponse(172, 89);
classifier_motor.train(trainDataFinal, motor_bikes);
cout << "trained motosekl" << endl;
classifier_car.train(trainDataFinal, cars);
cout << "trained cars" << endl;
classifier_cow.train(trainDataFinal, cows);
cout << "trained cows" << endl;
}

void testDataSet(Mat testData){
//    Mat testData = histogramToMat();
Mat motor_test = subMat(testData, 92, 23);
Mat car_test = subMat(testData, 195, 20);
Mat cow_test = subMat(testData, 304, 22);
vector<Mat> testTmp= {motor_test,car_test,cow_test};
Mat testDataFinal;
vconcat(testTmp, testDataFinal);
Mat motor_bikes,cars,cows;
classifier_motor.predict(testDataFinal, &motor_bikes);
cout << motor_bikes << endl;
classifier_car.predict(testDataFinal, &cars);
cout << cars << endl;
classifier_cow.predict(testDataFinal, &cows);
cout << cows << endl;
}

void saveBoWFiles(int start, int end,const char *type, int type1){
Mat imageHists= histogramToMat();
vector<Mat> matrix;
for(int i=start; i<end; i++){
    Mat currImage = Mat::zeros(1, 100, CV_32FC1);
    for(int j=0; j<100; j++){
        currImage.at<float>(0,j)=imageHists.at<float>(i,j);
    }
    matrix.push_back(currImage);
}
saveDescriptors(matrix, type,type1, ".bow");
}

void descriptorExtractor(){
char motorchar[] = "/Users/msalah93/Desktop/Sem9/PNGImages/motorbike-testset/";
char carchar[] = "/Users/msalah93/Desktop/Sem9/PNGImages/sideviews-cars/";
char cowchar[] = "/Users/msalah93/Desktop/Sem9/PNGImages/sideviews-cows2/";
vector<Mat> motorbikes = readImages(motorchar,1);
vector<Mat> cars = readImages(carchar,2);
vector<Mat> cows = readImages(cowchar,3);
applyHistEqualization(motorbikes);
applyHistEqualization(cars);
applyHistEqualization(cows);
applyResizing(motorbikes);
applyResizing(cars);
applyResizing(cows);

// Only done 1st time running
vector<vector<KeyPoint>> motorbikes_keypoints = denseDetector(motorbikes);
vector<vector<KeyPoint>> cars_keypoints = denseDetector(cars);
vector<vector<KeyPoint>> cows_keypoints = denseDetector(cows);
vector<Mat> motorbikes_mat = extractor(motorbikes,motorbikes_keypoints);
vector<Mat> cars_mat = extractor(cars,cars_keypoints);
vector<Mat> cows_mat = extractor(cows,cows_keypoints);
saveDescriptors(motorbikes_mat, "motorbikes",1);
saveDescriptors(cars_mat, "cars",2);
saveDescriptors(cows_mat, "cows",3);

vector<Mat> motorbikes_mat = extractFromSiftFile(motorbikes,"motorbikes",1,".sift");
vector<Mat> cars_mat = extractFromSiftFile(cars, "cars",2, ".sift");
vector<Mat> cows_mat = extractFromSiftFile(cows, "cows",3, ".sift");
vector<Mat> randomised = getRandomImages(motorbikes_mat,cars_mat,cows_mat);
cout << "finished randomised" << endl;

//Only Done 1st time running
Mat randomDesc;
vconcat(randomised, randomDesc);
Mat randomCentroids= applyKmeans(randomDesc);
cout << "finished kmeans" << endl;
initialiseBoW(motorbikes_mat, randomCentroids);
cout << "finished BOW motor" << endl;
initialiseBoW(cars_mat, randomCentroids);
cout << "finished BOW cars" << endl;
initialiseBoW(cows_mat, randomCentroids);
cout << "finished BOW cows" << endl;
saveBoWFiles(0, 115, "motorbikes", 1);
saveBoWFiles(115, 215, "cars", 2);
saveBoWFiles(215, 326, "cows", 3);

vector<Mat> motorbikes_mat_bow = extractFromSiftFile(motorbikes,"motorbikes",1,".bow");
vector<Mat> cars_mat_bow = extractFromSiftFile(cars, "cars",2, ".bow");
vector<Mat> cows_mat_bow = extractFromSiftFile(cows, "cows",3, ".bow");
Mat motor_bow, cars_bow, cows_bow;
vconcat(motorbikes_mat_bow,motor_bow);
vconcat(cars_mat_bow,cars_bow);
vconcat(cows_mat_bow,cows_bow);
vector<Mat> total_bow= {motor_bow, cars_bow, cows_bow};
Mat trainTestData;
vconcat(total_bow, trainTestData);
trainDataSet(trainTestData);
cout << "Trained." << endl;
testDataSet(trainTestData);
cout << "Tested";
}

int main(int argc, char** argv)
{
descriptorExtractor();
return 0;
}