#include "mycvdigitssegm.h"

mycvDigitsSegm::mycvDigitsSegm()
{

}

mycvDigitsSegm::~mycvDigitsSegm()
{

}

void mycvDigitsSegm::segmentNum(const Mat &aFrame, Mat &aDst, vector<Mat> &aDigits, int aMode)
{
    switch(aMode)
    {
    case MODE_DETECT_CONTOURS: segmentNumContours(aFrame.clone(), aDst, aDigits); break;
    case MODE_DETECT_VERT_PROJ: segmentNumVerticalProj(aFrame.clone(), aDst, aDigits); break;
    }
}

void mycvDigitsSegm::loadPatterns(Mat &aSrc)
{
//	mcDigitFeatures.loadPattern(aSrc);
}

void mycvDigitsSegm::mycvPreprocCount(const Mat &aSrc, Mat &aDst)
{
    mycv mcMycv;
    Mat imgGray;
    cvtColor(aSrc, imgGray, CV_BGR2GRAY);

    int blockSz = mcMycv.blockSize(imgGray.cols, 120);

    /// Threshold
    Mat imgFrameThresh;

    blockSz = mcMycv.blockSize(imgGray.rows, 10);
    bilateralFilter(imgGray.clone(), imgGray, blockSz, 150, blockSz);
    // Threshold
    adaptiveThreshold(imgGray, imgFrameThresh, 255, CV_ADAPTIVE_THRESH_MEAN_C, CV_THRESH_BINARY, blockSz, 0);
    mcMycv.mycvRemoveHighLowObjects(imgFrameThresh, imgFrameThresh, 0, 0, 0, imgFrameThresh.rows/9, imgFrameThresh.rows/8,
                             (imgFrameThresh.rows/8)*(imgFrameThresh.rows/8));
    medianBlur(imgFrameThresh, imgFrameThresh, 3);

    // Remove filled collumns
    for(int i = 0 ; i < imgFrameThresh.cols ; i++)
    {
        int colSum = 0;
        for(int j = 0 ; j < imgFrameThresh.rows ; j++)
        {
            if(imgFrameThresh.at<uchar>(Point(i, j)) > 0)colSum++;
        }
        if(colSum == imgFrameThresh.rows)
        {
            for(int j = 0 ; j < imgFrameThresh.rows ; j++)
            {
                imgFrameThresh.at<uchar>(Point(i, j)) = 0;
            }
        }
    }

    // red mask
    Mat imgRed;
    mcMycv.mycvFindRed(aSrc, imgRed);
    blockSz = mcMycv.blockSize(imgFrameThresh.rows, 10);
    Mat element = getStructuringElement(MORPH_ELLIPSE, Size(blockSz, blockSz));
    morphologyEx(imgRed, imgRed, CV_MOP_DILATE, element);

//    bitwise_not ( imgRed, imgRed );
    imgFrameThresh = imgFrameThresh & imgRed;

    // Remove upper & lower borders
    Mat imgRowSum;
    vector<Point> charSep;
    mcMycv.mycvVerHorProjection(imgFrameThresh, imgRowSum, PROJ_HRIZONTAL, charSep);

    Point cutP = mcMycv.mycvCheckHorizBorders(imgFrameThresh, charSep);
    cvtColor(imgRowSum, imgRowSum, CV_GRAY2BGR);

    uchar *p1, *p2;
    p1 = (uchar *) imgFrameThresh.ptr<uchar>(cutP.x);
    p2 = (uchar *) imgFrameThresh.ptr<uchar>(cutP.y);
    for(int j = 0 ; j < imgFrameThresh.cols ; j++)
    {
        p1[j] = 0;
        p2[j] = 0;
    }

    mcMycv.mycvRemoveHighLowObjects(imgFrameThresh, imgFrameThresh, 0, imgFrameThresh.cols/5, 0);
    mcMycv.mycvRemoveHighLowObjects(imgFrameThresh, imgFrameThresh, 0, 0, 0, 10, 10);

    aDst = imgFrameThresh;
}

void mycvDigitsSegm::segmentNumVerticalProj(const Mat &aFrame, Mat &aDst, vector<Mat> &aDigits)
{
    mycv mcMycv;
    Mat imgFrameOrig = aFrame, imgFrameThresh;
    aDst = aFrame;

    // General frame preprocessing
    mycvPreprocCount(imgFrameOrig, imgFrameThresh);

    // Morphology
    int blockSz = mcMycv.blockSize(imgFrameThresh.rows, 20);
    Mat element = getStructuringElement(MORPH_ELLIPSE, Size(blockSz, blockSz));
    morphologyEx(imgFrameThresh, imgFrameThresh, CV_MOP_DILATE, element);

    // Fill up holes
    Mat imgThreshInv;
    bitwise_not(imgFrameThresh.clone(), imgThreshInv);

    // Remove small objects in inverted img -> small holes gone
    mcMycv.mycvRemoveHighLowObjects(imgThreshInv.clone(), imgThreshInv, 0, 0, 0, imgThreshInv.cols/7, imgThreshInv.rows-imgThreshInv.rows/5);
    bitwise_not(imgThreshInv.clone(), imgFrameThresh);

    // Remove too large objects
    mcMycv.mycvRemoveHighLowObjects(imgFrameThresh.clone(), imgFrameThresh, 0, 0, 0, imgFrameThresh.rows/10, imgFrameThresh.rows/6);

    // Morphology connect object in vertical direction
    blockSz = mcMycv.blockSize(imgFrameThresh.rows, 10);
    element = getStructuringElement(MORPH_ELLIPSE, Size(1+blockSz/10, blockSz*5));
    morphologyEx(imgFrameThresh, imgFrameThresh, CV_MOP_CLOSE, element);

    // Cut upper and lower frame 1/3
    Mat imgCut;
    int perc = imgFrameThresh.rows/3;
    Rect roi(0, perc, imgFrameThresh.cols-1, imgFrameThresh.rows-2*perc);
    imgCut = imgFrameThresh(roi).clone();

//        mcMycv.mycvRemoveBorderSmalls(imgCut, imgCut, imgCut.rows/4);

    /// Distinguish each digit using vertical projection
    Mat imgColSum;
    vector<Point> charSepars;   // x - start col, y - end col
    int maxSum = mcMycv.mycvVerHorProjection(imgCut, imgColSum, PROJ_VERTICAL, charSepars);

    if(maxSum == 0 || charSepars.size() == 0)return; // numbers not found

    Mat imgLines;
//        cvtColor(imgFrameOrig, imgLines, CV_GRAY2BGR);
    imgFrameOrig.copyTo(imgLines);
    // reduse redundant separators - leave 8 biggest
    while(charSepars.size() > 8)
    {
        uint minSep = numeric_limits<uint>::max();
        uint minPos = 0;
        for(int j = 0 ; j < (int)charSepars.size() ; j++)
        {
            // find smallest
            Point point = charSepars[j];
            if(minSep > point.y - point.x)
            {
                minPos = j;
                minSep = point.y - point.x;
            }
        }
        charSepars.erase(charSepars.begin()+minPos);
    }
    if(charSepars.size() == 0)return;

    // Evaluate mean width
    uint meanCharWidth = 0;
    for(int j = 0 ; j < (int)charSepars.size() ; j++)
    {
        meanCharWidth += charSepars[j].y - charSepars[j].x;
        line(imgLines, Point(charSepars[j].x, 0), Point(charSepars[j].x, imgColSum.cols-1), Scalar(255, 0, 0));
        line(imgLines, Point(charSepars[j].y, 0), Point(charSepars[j].y, imgColSum.cols-1), Scalar(0, 255, 0));
        line(imgLines, Point(charSepars[j].x, 0), Point(charSepars[j].y, 0), Scalar(0, 255, 0), 3);
    }
    meanCharWidth /= charSepars.size();

    // normalize digits width
    imgFrameOrig.copyTo(imgLines);
    for(int j = 0 ; j < (int)charSepars.size() ; j++)
    {
        Point act = charSepars[j];
        int width = act.y - act.x;
        int dif = meanCharWidth - width;
        charSepars[j].x -= dif/2;
        charSepars[j].y += dif/2;

        if(charSepars[j].x < 0)charSepars[j].x = 0;
        if(charSepars[j].y >= imgCut.cols)charSepars[j].y = imgCut.cols-1;

        line(imgLines, Point(charSepars[j].x, 0), Point(charSepars[j].x, imgColSum.cols-1), Scalar(255, 0, 0));
        line(imgLines, Point(charSepars[j].y, 0), Point(charSepars[j].y, imgColSum.cols-1), Scalar(0, 255, 0));
        line(imgLines, Point(charSepars[j].x, 0), Point(charSepars[j].y, 0), Scalar(0, 255, 0), 3);
    }



    // Create rectangles
    vector<Rect> charRects;
    for(int j = 0 ; j < charSepars.size() ; j++)
    {
        Rect newRect(charSepars[j].x, 0, charSepars[j].y-charSepars[j].x, imgLines.rows);
        charRects.push_back(newRect);
    }

    // Process digit
    mcDigitPreproc.loadFrame(aFrame);
    vector<Mat> digits;
    for(int j = 0 ; j < (int)charRects.size() ; j++)
    {
        Rect rect = charRects[j];
        Mat imgActDigit;
        if(mcDigitPreproc.preproc(imgActDigit, rect, PROC_MODE_NORMALIZE))
        {
            digits.push_back(imgActDigit);
//            float num = mcDigitFeatures.match(imgActDigit, 0);

//                float num = mcDigitGeometric.match(imgActDigit);
            // If digit recognized
//            if(num >= 0)
//            {
//                std::ostringstream ostream1;
//                ostream1 << num;
//                putText(imgLines, ostream1.str(), Point(charRects[j].x + 2, 30), CV_FONT_HERSHEY_COMPLEX, 1, Scalar(0,0,255));
//            }
        }
        rectangle(aDst, rect, Scalar(0,255,0));
    }

    aDigits = digits;
}

void mycvDigitsSegm::segmentNumContours(const Mat &aFrame, Mat &aDst, vector<Mat> &aDigits)
{
    mycv mcMycv;
    Mat imgFrameOrig = aFrame, imgFrameThresh;
//    aDst = aFrame;

    // General digit preprocessing
    mycvPreprocCount(imgFrameOrig, imgFrameThresh);


    int blockSz = mcMycv.blockSize(imgFrameOrig.rows, 10);

    Mat element = getStructuringElement(MORPH_ELLIPSE, Size(3, 1));
    morphologyEx(imgFrameThresh, imgFrameThresh, CV_MOP_CLOSE, element);
    mcMycv.mycvRemoveHighLowObjects(imgFrameThresh, imgFrameThresh, 0, 0, 0, imgFrameThresh.rows/10, imgFrameThresh.rows/5);
    mcMycv.mycvRemoveHighLowObjects(imgFrameThresh, imgFrameThresh, 0, imgFrameThresh.cols/6, imgFrameThresh.rows*7/8);

    // Morphology
    blockSz = mcMycv.blockSize(imgFrameThresh.cols, 80);
    element = getStructuringElement(MORPH_RECT, Size(1+blockSz/2, blockSz*2));
    morphologyEx(imgFrameThresh, imgFrameThresh, CV_MOP_CLOSE, element);

    blockSz = mcMycv.blockSize(imgFrameThresh.cols, 90);
    element = getStructuringElement(MORPH_ELLIPSE, Size(blockSz, blockSz));
    morphologyEx(imgFrameThresh, imgFrameThresh, CV_MOP_DILATE, element);

    /// Character segmentation
    vector<vector<Point> > contours;
    findContours(imgFrameThresh.clone(), contours, CV_RETR_EXTERNAL, CV_CHAIN_APPROX_SIMPLE);
    // The array for storing the approximation curve
    std::vector<cv::Point> approx;

    Mat imgSegChars;
    imgFrameOrig.copyTo(imgSegChars);
    vector<Rect> rects;
    for(int j = 0 ; j < (int)contours.size() ; j++)
    {
        approxPolyDP(
            cv::Mat(contours[j]),
            approx,
            cv::arcLength(cv::Mat(contours[j]), true) * 0.03,
            true
        );
        Rect rect;
        rect = boundingRect(approx);

        // Add more space up and down
        int increase = 10;
        rect.y -= increase;
        if(rect.y < 0)rect.y = 0;

        rect.height += increase;
        if(rect.y + rect.height >= imgFrameOrig.rows)rect.height = imgFrameOrig.rows - rect.y;

        // Filter inconvenient rects
        rects.push_back(rect);
        rectangle(imgSegChars, rect, Scalar(0, 0, 255));
    }

    // Check digits
    mcMycv.mycvCheckDigit(imgFrameThresh, rects);

    Mat imgAdjRects;
    imgFrameOrig.copyTo(imgAdjRects);
    mcDigitPreproc.loadFrame(aFrame);
    aDst = imgAdjRects.clone();
//    vector<Mat> digits;
    for(int j = 0 ; j < (int)rects.size() ; j++)
//    int j = 0;
    {
//            Rect rect(centers[j].x - meanRectSize.x/2 - meanRectSize.x/10, minY, meanRectSize.x + meanRectSize.x/10, meanRectSize.y);
        Rect rect = rects[j];
//        rectangle(imgAdjRects, rect, Scalar(0, 255, 0));
        Point circP = Point(rect.x+rect.width/2, rect.y+rect.height/2);
        circle(imgAdjRects, circP, 3, Scalar(0,0,255));
        Mat imgActDigit;
//            mycvRemoveHigherObjects(imgActDigit, imgActDigit, 0, 0, 0, imgActDigit.cols/15, imgActDigit.rows/4, imgActDigit.cols/2);

        if(mcDigitPreproc.preproc(imgActDigit, rect, PROC_MODE_NORMALIZE))
        {
        	aDigits.push_back(imgActDigit);
//                float num = mcDigitFeatures.match(imgActDigit, 0);
//                float num = mcDigitGeometric.match(imgActDigit);
//            float num = mcANN.test(imgActDigit);
//            if(num >= 0)
//            {
//                std::ostringstream ostream1;
//                ostream1 << num;
//                putText(aDst, ostream1.str(), Point(rect.x + 2, 30), CV_FONT_HERSHEY_COMPLEX, 1, Scalar(0,0,255));
//            }
            rectangle(aDst, rect, Scalar(0,255,0));
        }
    }

//    std::ostringstream ostream1;
//	ostream1 << "Num Of rect.: " << rects.size();
//	putText(aDst, ostream1.str(), Point(10, 30), CV_FONT_HERSHEY_COMPLEX, 1, Scalar(125));
    return;
//    aDigits = digits;
}
