#include <opencv2/core.hpp>
#include <opencv2/ml.hpp>

class MLP
{
public:
    MLP();                                          // 构造函数，mlp 的初始化
    ~MLP() { }

    void loadXML( const QString &xml  );            // 加载已经训练好的模型
    void saveXML( const QString &xml  );            // 保存训练好的模型
    void loadCSV( const QString &csv );             // 加载正负样本用于训练
    void train();                                   // 训练

    // 最后，对图片进行判别：是否为路牌
    //  输入为一串图片路径
    //  输出为图片路径和一个标志，true 说明图片为路牌，false 说明不是
    QList<QPair<QString, bool> > predictImages( const QStringList &images );

private:
    cv::Ptr<cv::ml::ANN_MLP> mlp;                   // 存储了参数
};

mlp = cv::ml::ANN_MLP::create();
// 层数和每层 neuron 个数对结果有一定影响，下文对此有论述
mlp->setLayerSizes( (cv::Mat)(cv::Mat_<int>(1,5)
                                << FEATURENUM, FEATURENUM / 2, FEATURENUM / 6, FEATURENUM / 24, 1) );

// 激活函数设置为 sigmoid
mlp->setActivationFunction( cv::ml::ANN_MLP::SIGMOID_SYM, 1.0, 1.0 );

// 训练方法为反向传播
mlp->setTrainMethod( cv::ml::ANN_MLP::BACKPROP, 0.1, 0.1 );
mlp->setTermCriteria( cv::TermCriteria(
                            cv::TermCriteria::COUNT + cv::TermCriteria::EPS
                            , 5000
                            , 0.01 ) );

// pos 和 neg 是图片路径列表，分别是正负样本集合
int np = pos.length();
int nn = neg.length();

// 生成数据集，这里的 features 和 flags 就是机器学习里
// 常说的 data 和 label
float *features = new float[FEATURENUM*(np + nn)];
float *flags    = new float[np + nn];

// 用图片初始化 features，并设置好 labels
for ( int i = 0; i < np; ++i ) {
    Utils::img2feature(qPrintable(pos.at(i)), features + FEATURENUM*i);
    *(flags + i)        =       1.0f;       // 正样本为 1
}
for ( int i = 0; i < nn; ++i ) {
    Utils::img2feature(qPrintable(neg.at(i)), features + FEATURENUM*np+FEATURENUM*i);
    *(flags + np + i)   =       -1.0f;      // 负样本为 -1
}

// 存到 OpenCV 的矩阵结构 Mat 里，并生成数据集
cv::Mat featureMat( np+nn, FEATURENUM,  CV_32F, features );
cv::Mat flagMat(    np+nn, 1,           CV_32F, flags );
cv::Ptr<cv::ml::TrainData> trainSet // mat 的每一行代表一条数据
    = cv::ml::TrainData::create( featureMat, cv::ml::ROW_SAMPLE, flagMat );

// 然后就可以训练了
mlp->train( trainSet );

// 不要忘了释放内存
delete[] features;
delete[] flags;

float feature[FEATURENUM];
Utils::img2feature( "path/to/image", feature );

cv::Mat featureMat  ( 1, FEATURENUM,    CV_32F, feature );
cv::Mat result      ( 1, 1,             CV_32FC1 );

mlp->predict( featureMat, result );
float *p = result.ptr<float>(0);

// 保存
mlp->save( "path/to/save/model(xml file)" );

// 加载
mlp = cv::Algorithm::load<cv::ml::ANN_MLP>( "path/to/saved-model.xml" );

// feature 已经分配好内存了，内存大小是 sizeof(float)*FEATURENUM
// FEATURENUM 是一个宏，定义为
//      #define FEATURENUM  ( 24*(3+3)+ 3 + 4*3 )
// 看了后面的数据处理你大概能懂这些数字每个代表什么

bool Utils::img2feature( const char *filePath, float *feature )
{
    Mat img = imread( filePath, cv::IMREAD_COLOR );
    if ( img.rows != IMGSIZE || img.cols != IMGSIZE ) {
        qDebug() << __FUNCDNAME__ << "failed, because of wrong dim," << "\tfilepath:" << filePath;
        return false;
    }

    int allCountR          =   0  , allCountG          =   0  , allCountB          =   0  ;
    int rowCountR[IMGSIZE] = { 0 }, rowCountG[IMGSIZE] = { 0 }, rowCountB[IMGSIZE] = { 0 };
    int colCountR[IMGSIZE] = { 0 }, colCountG[IMGSIZE] = { 0 }, colCountB[IMGSIZE] = { 0 };
    int tempR[4]           = { 0 }, tempG[4]           = { 0 }, tempB[4]           = { 0 };

    for( int i = 0; i < IMGSIZE; ++i ) {
        for( int j = 0; j < IMGSIZE; ++j ) {

            int b = img.at<cv::Vec3b>(i,j)[0];
            int g = img.at<cv::Vec3b>(i,j)[1];
            int r = img.at<cv::Vec3b>(i,j)[2];

            allCountR += r;     rowCountR[i] += r;     colCountR[j] += r;
            allCountG += g;     rowCountG[i] += g;     colCountG[j] += g;
            allCountB += b;     rowCountB[i] += b;     colCountB[j] += b;

            /*
             * 分成四个部分，统计各自区域内的 rgb 比例，四个区域低编号为：
             *
             *           0 | 1
             *           --+--
             *           2 | 3
            */
            int index = 2 * (i < IMGSIZE / 2 ? 0 : 1) + (j < IMGSIZE / 2 ? 0 : 1);
            tempR[index] += r;
            tempG[index] += g;
            tempB[index] += b;
        }
    }

    // 每行每列 rgb
    for ( int i=0; i < IMGSIZE; ++i ) {
        feature[i*6+0] = (float)rowCountR[i]/allCountR; // 第 i 行的红色占全图红色的比例
        feature[i*6+1] = (float)rowCountG[i]/allCountG; // 绿色
        feature[i*6+2] = (float)rowCountB[i]/allCountB; // 蓝色
        feature[i*6+3] = (float)colCountR[i]/allCountR; // 第 i 列
        feature[i*6+4] = (float)colCountG[i]/allCountG;
        feature[i*6+5] = (float)colCountB[i]/allCountB;
    }
    // rgb 三种颜色的比例
    feature[IMGSIZE*6+0] = (float)allCountR / (allCountR + allCountG +allCountB);
    feature[IMGSIZE*6+1] = (float)allCountG / (allCountR + allCountG +allCountB);
    feature[IMGSIZE*6+2] = (float)allCountB / (allCountR + allCountG +allCountB);
    // 四个区域
    for ( int i = 0; i < 4; ++i ) {
        // 区域内的 rgb 比例
        int total = tempR[i] + tempG[i] + tempB[i];
        feature[IMGSIZE*6+3 + 3*i+0] = (float)tempR[i] / total;
        feature[IMGSIZE*6+3 + 3*i+1] = (float)tempG[i] / total;
        feature[IMGSIZE*6+3 + 3*i+2] = (float)tempB[i] / total;
    }
    return true;
}