/*

 * Recognizer.cpp

 *

 *  Created on: Mar 23, 2011

 *      Author: ams

 *

 *

 */

#include <cstdio>

#include <unistd.h>

#include <dirent.h>

#include <opencv/cv.h>

#include <opencv/highgui.h>

#include <opencv/ml.h>

#include "MoUT.hpp"

// Size of the sample images in pixels

#define SAMPLE_WIDTH 27

#define SAMPLE_HEIGHT 20

// Multiplied by 3 for the three color channels

#define SAMPLE_SIZE (SAMPLE_WIDTH * SAMPLE_HEIGHT * 3)

void trainNN()

{

        CvANN_MLP brains;

        int samples;

        FILE *fin;

        fin = fopen("./cup.dat", "r");

        //Get the number of samples.

        fscanf(fin, "%d", &samples);

        printf("Found training file with %d samples...\n", samples);

        /* Create a matrix to hold the training data

         * A sample is a 20x27 RGB image, so there are 1620 inputs in a sample.

         * There is also one output, so there are a total of 1621 values.

         * With 2304 samples, that means there are 3,734,784 values, and at

         * 4 bytes per value, over 14 megs of data. This should be trivial, but it

         * causes a segmentation fault. Using unit8_t appears to be ok...

         */

        uint8_t trainingData[samples][SAMPLE_SIZE + 1];

        /* OpenCV matrices for the sample data. One holds inputs, the other holds outputs

         * This has to be floating point, despite the fact that the input isn't */

        CvMat* trainInput = cvCreateMat(samples, SAMPLE_SIZE, CV_32FC1);

        CvMat* trainOutput = cvCreateMat(samples, 1, CV_32FC1);

        /* Weights of each sample, which should be all the same */

        CvMat* sampleWeights = cvCreateMat(samples, 1, CV_32FC1);

        /* Matrix representation of the network.

         * The first number is the count of layers.

         */

        CvMat* nnLayers = cvCreateMat(4, 1, CV_32SC1);

        CvMat trainInput1, trainOutput1, nnLayers1, sampleWeights1;

        cvGetRows(trainInput, &trainInput1, 0, samples);

        cvGetRows(trainOutput, &trainOutput1, 0, samples);

        cvGetRows(sampleWeights, &sampleWeights1, 0, samples);

        cvGetRows(nnLayers, &nnLayers1, 0, 4);

        /*Setting the number of neurons on each layer of the ANN

         * Layer 1: 1620 neurons (One for each channel of each input pixel)

         * Layer 2: Undecided, this parameter needs tweaking

         *  Layer 3: Undecided, this parameter also needs tweaking

         *  Layer 4: 1 neurons (1 output, for whether the input has a red cup in it)

         *

         *  Heuristic principles for the design of artificial neural networks, by

         *  Stephen Walczak and Narciso Cerpa, has some useful pointers.

         */

        cvSet1D(&nnLayers1, 0, cvScalar(SAMPLE_SIZE));

        cvSet1D(&nnLayers1, 1, cvScalar((int) (SAMPLE_SIZE * 0.75)));

        cvSet1D(&nnLayers1, 2, cvScalar((int) (SAMPLE_SIZE * 0.65)));

        cvSet1D(&nnLayers1, 3, cvScalar(1));

        /* Read the data from the training file. This might take

         * something like two forevers.

         */

        printf("Reading training data...");

        for (int sample = 0; sample < samples; sample++)

        {

                for (int value = 0; value < SAMPLE_SIZE + 1; value++)

                {

                        //%hhu is an integer promoted from a char that is unsigned

                        //This may be a problem, since trainingData is uint8_t (char)

                        fscanf(fin, "%hhu", &trainingData[sample][value]);

                }

        }

        //Done reading the file

        fclose(fin);

        printf("done.\n");

        //Assemble the ML training data.

        printf("Assembling training data...");

        for (int sample =0; sample < samples; sample++)

        {

                //Input values

                for (int value = 0; value < SAMPLE_SIZE; value++)

                {

                        //load each data value

                        cvSetReal2D(&trainInput1, sample, value, trainingData[sample][value]);

                }

           //Output values

           cvSet1D(&trainOutput1, sample, cvScalar(trainingData[sample][SAMPLE_SIZE + 1]));

           //Weight (setting everything to 1)

           cvSet1D(&sampleWeights1, sample, cvScalar(1));

        }

        printf("done.\n");

        //Create the ANN

        brains.create(&nnLayers1);

        //Train the ANN

        printf("Training ANN...");

        brains.train(

                        trainInput,

                        trainOutput,

                        sampleWeights,

                        0,                                                                                                                                                                                        //Take all samples into account

                        CvANN_MLP_TrainParams(

                                        cvTermCriteria(

                                                        CV_TERMCRIT_ITER + CV_TERMCRIT_EPS,                //Stop after either a number of iterations or a certain accuracy is achieved

                                                        10000,                                                                                                                                        //Maximum iterations

                                                        1.0),                                                                                                                                                //Accuracy TODO this needs finagling

                                        CvANN_MLP_TrainParams::BACKPROP,                                                //Use backpropagation

                                        0.01,                                                                                                                                                        // coefficeint to multiply weight gradient by. 0.1 is recommended by online sources

                                        0.05)                                                                                                                                                        // backpropagation "moment scale", couldn't find documentation

                        );

        printf("done.\n");

        //Save the results to a file

        CvFileStorage* annOutFile = cvOpenFileStorage("./annParams", 0, CV_STORAGE_WRITE_TEXT);

        brains.write(annOutFile, "annParams");

        printf("Wrote to file.\n");

}

/* Use the neural net to determine if the target occurs in an image

 */

bool recognize(string dirName)

{

        // Load the NN from the file

        CvFileStorage* annInFile = cvOpenFileStorage("./annParams", 0, CV_STORAGE_READ);

        CvFileNode* annInNode;

        CvANN_MLP ann;

        ann.read(annInFile, annInNode);

        //get all the images in the directory

        DIR *dpdf;

        struct dirent *epdf;

        dpdf = opendir(dirName.c_str());

        if (dpdf != NULL)

        {

                //Loop over all the files in the directory

                while (epdf = readdir(dpdf))

                {

                        //skip the dot files

                        if(strcmp(epdf->d_name, ".") == 0 || strcmp(epdf->d_name, "..") == 0)

                        {

                                continue;

                        }

                        string fullName = (dirName + "/" + epdf->d_name).c_str();

                        //Downsample and convert to an input vector

                        string nnVector = img2NNVect(fullName);

                        float strSample[SAMPLE_SIZE];

                        vector<string> tokens;

                        Tokenize(nnVector, tokens);

                        vector<string>::iterator tokIt;

                        int index = 0;

                        float newVal = 0.0;

                        for(tokIt = tokens.begin(); tokIt != tokens.end(); tokIt++)

                        {

                                if(fromString<float>(newVal, (*tokIt), std::dec))

                                {

                                        strSample[index] = newVal;

                                }

                                else

                                {

                                        strSample[index] = -1.0;

                                        printf("Error converting %s to a float\n", (*tokIt).c_str());

                                }

                                index++;

                        }

                        CvMat sample = cvMat(1, SAMPLE_SIZE, CV_32FC1, strSample);

                        //Matrix to store the output in

                        float output[1];

                        CvMat nnOutput = cvMat(1, 1, CV_32FC1, output);

                        // Feed to the NN

                        ann.predict(&sample, &nnOutput);

                        // Display the results

                        printf("File: %s\t\tOutput: %f", fullName.c_str(), nnOutput.data.fl[0]);

                }

        }

        return false;

}

void usage()

{

        printf("Please specify training (-t) or recognizing (-r)\n");

}

int main(int argc, char **argv)

{

        int c;

        while ((c = getopt(argc, argv, "tr")) != -1)

        {

                switch (c)

                {

                case 't':

                        trainNN();

                        break;

                case 'r':

                        recognize("/data/cup_untrained");

                        //recognize("/data/cup_untrained");

                        break;

                default:

                        usage();

                        return EXIT_FAILURE;

                }

        }

        return EXIT_SUCCESS;

}