#include "mycvdigitprocann.h"

#ifdef DEBUG_DIGIT
    #include <QDebug>
#endif

mycvDigitProcANN::mycvDigitProcANN()
{
    // numbers 0-9 and same number
    // of gaps between them -> 0.5-9.5
    mNumOfChars = 20;
    // Indicates trained ANN
    mbTrained = false;
}

mycvDigitProcANN::~mycvDigitProcANN()
{
    mcANN.clear();
    clearData();
}

int mycvDigitProcANN::addData(const Mat &aDat, int aLabel)
{
    Mat img;
    // Convert to gray
    if(aDat.dims > 1)
        cvtColor(aDat, img, CV_BGR2GRAY);
    else img = aDat.clone();
    // Reshape to one row
    Mat row = img.reshape(0, 1);

    // If some data already present
    if(!mDataSet.empty())
    {
        // Check data size
        if(mDataSet.cols != row.cols)return 0;
    }

    // Convert to float
    row.convertTo(row, CV_32FC1);

    // new data adding
    mDataSet.push_back(row);

    // add label
    mLabels.push_back(aLabel);

    return 1;
}

void mycvDigitProcANN::clearData()
{
    mDataSet.release();
    mLabels.release();
}

int mycvDigitProcANN::train(int aHiddenNeurons)
{
    // Check dataset
    if(mDataSet.empty() || mLabels.empty() ||
       mDataSet.rows != mLabels.rows)return 0;

    // Define ANN layers
//    cv::Mat layers = cv::Mat(4, 1, CV_32SC1);
//    layers.row(0) = cv::Scalar(mDataSet.cols);    // input layer
//    layers.row(1) = cv::Scalar(aHiddenNeurons);     // hidden layer 1
//    layers.row(2) = cv::Scalar(aHiddenNeurons);     // hidden layer 2
//    layers.row(3) = cv::Scalar(mNumOfChars);        // output layer

    cv::Mat layers = cv::Mat(3, 1, CV_32SC1);
    layers.row(0) = cv::Scalar(mDataSet.cols);    // input layer
    layers.row(1) = cv::Scalar(aHiddenNeurons);     // hidden layer 1
    layers.row(2) = cv::Scalar(mNumOfChars);        // output layer

    // Set ANN parameters
    CvANN_MLP_TrainParams params(

                                // terminate the training after either 1000
                                // iterations or a very small change in the
                                // network wieghts below the specified value
                                cvTermCriteria(CV_TERMCRIT_ITER+CV_TERMCRIT_EPS,     3000, 0.000001),
                                // use backpropogation for training
                                CvANN_MLP_TrainParams::BACKPROP,
                                // co-efficents for backpropogation training

                                0.1,
                                0.1);


    mcANN.create(layers, CvANN_MLP::SIGMOID_SYM, 1, 1);

    // Create training class of outputs
    Mat trainClasses;
    trainClasses.create( mDataSet.rows, mNumOfChars, CV_32FC1 );
    for( int i = 0; i < trainClasses.rows; i++ )
    {
        for( int k = 0; k < mNumOfChars; k++ )
        {
            //If class of data i is same than a k class
            if( k == mLabels.at<int>(i, 0) )
                trainClasses.at<float>(i,k) = 1;
            else
                trainClasses.at<float>(i,k) = 0;
        }
    }

#ifdef DEBUG_DIGIT
    imshow("data set", mDataSet);
    imshow("train classes", trainClasses);
#endif

    // Train ANN
    int iters = mcANN.train(mDataSet, trainClasses, cv::Mat(),cv::Mat(),params);

#ifdef DEBUG_DIGIT
    qDebug() << "ANN iterations = " << iters << endl;
#endif
    mbTrained = true;
    return 1;
}

float mycvDigitProcANN::test(Mat aTest)
{
    if(!mbTrained)return -1;
    // Reshape
    // Convert to gray
    Mat img;
    if(aTest.dims > 2)
        cvtColor(aTest, img, CV_BGR2GRAY);
    else img = aTest.clone();
    // Reshape to one row
    Mat row = img.reshape(0, 1);
#ifdef DEBUG_DIGIT
    imshow("test data", row);
#endif

    // Convert to float
    row.convertTo(row, CV_32FC1);

    // Check size
    /// TODO

    // Output vector
    Mat output(1, mNumOfChars, CV_32FC1);
    // Classify test data
    mcANN.predict(row, output);

    #ifdef DEBUG_DIGIT
        qDebug() << "ANN out:";
        for(int i = 0 ; i < output.cols ; i++)
            qDebug() << output.at<float>(0, i) << ", ";
        qDebug() << endl;
    #endif

    // Find max position
    Point maxLoc;
    double maxVal;
    minMaxLoc(output, 0, &maxVal, 0, &maxLoc);

    // Decode number
    if(maxLoc.x < 10)return (float)maxLoc.x;
    else return (float)(maxLoc.x-10)+0.5f;
}

void mycvDigitProcANN::saveXML(string aStr)
{
    cv::FileStorage fs(aStr, cv::FileStorage::WRITE);
    mcANN.write(*fs, "mlp");
}

void mycvDigitProcANN::loadXML(string aStr)
{
    mcANN.load(aStr.c_str(), "mlp");
    mbTrained = true;
}
