#include <iostream>
#include <opencv2/opencv.hpp>
#include <opencv2/highgui/highgui.hpp>

void convertToGrayScale(cv::Mat & in, cv::Mat & out) {
    for (int i = 0; i < in.rows; i++) {
        for (int j = 0; j < in.cols; j++) {
            cv::Vec3b pixel = in.at<cv::Vec3b>(i, j);
            uchar gsColor = (uchar)(0.0722 * pixel.val[0] + 0.7152 * pixel.val[1] + 0.2126 * pixel.val[2]);
            out.at<uchar>(i, j) = gsColor;
        }
    }
}

void getBinaryModel(cv::Mat & in, int * out) {
    uchar threshold = 0;
    int histogram[256] = {0};
    int sum = 0;
    double omega1 = 0, omega2, mu1, mu2, sigma = 0, max = 0;
    int pixelNum = in.rows * in.cols;
    for (int i = 0; i < in.rows; i++) {
        for (int j = 0; j < in.cols; j++) {
            ++histogram[in.at<uchar>(i, j)];
            sum += in.at<uchar>(i, j);
        }
    }
    for (int i = 0; i < 256; ++i) {
        omega1 += histogram[i];
        if (omega1 == 0)
            continue;
        omega2 = pixelNum - omega1;
        if (omega2 == 0)
            break;
        sigma += i * histogram[i];
        mu1 = sigma / omega1;
        mu2 = (sum - sigma) / omega2;
        double_t between = omega1 * omega2 * (mu1 - mu2) * (mu1 - mu2);
        if (between > max) {
            max = between;
            threshold = (uchar)i;
        }
    }

    for (int i = 0; i < in.rows; i++) {
        for (int j = 0; j < in.cols; j++) {
            if (in.at<uchar>(i, j) < threshold) {
                out[i * in.cols + j] = 1;
            } else {
                out[i * in.cols + j] = 0;
            }
        }
    }
}

int getLeft(int * model, int width, int i, int j) {
    if (j == 0) {
        return 0;
    }
    return model[i * width + j - 1];
}

int getUp(int * model, int width, int i, int j) {
    if (i == 0) {
        return 0;
    }
    return model[(i - 1) * width + j];
}

unsigned long countAreas(int * model, int width, int height) {
    int curComp = 0;
    int compNum = 0;
    int leftComp, upComp;
    std::map<int, int> comps;
    for (int i = 0; i < height; i++) {
        for (int j = 0; j < width; j++) {
            int pos = i * width + j;
            if (model[pos] == 1) {
                leftComp = getLeft(model, width, i, j);
                upComp = getUp(model, width, i, j);
                if (leftComp == 0 && upComp == 0) {
                    ++curComp;
                    ++compNum;
                    comps[curComp] = curComp;
                    model[pos] = curComp;
                } else {
                    if (leftComp == 0) {
                        model[pos] = comps[upComp];
                    } else if (upComp == 0 || comps[leftComp] == comps[upComp]) {
                        model[pos] = comps[leftComp];
                    }
                    else {
                        if (leftComp < upComp) {
                            model[pos] = comps[leftComp];
                            if (comps[upComp] != comps[leftComp]) {
                                comps[upComp] = comps[leftComp];
                            }
                        } else {
                            model[pos] = comps[upComp];
                            if (comps[leftComp] != comps[upComp]) {
                                comps[leftComp] = comps[upComp];
                            }
                        }
                    }
                }
            }
        }
    }
    std::set<int> difComps;
    for (int i = 0; i < height; i++) {
        for (int j = 0; j < width; j++) {
            int pos = i * width + j;
            if (model[pos] != 0) {
                model[pos] = comps[model[pos]];
                difComps.insert(model[pos]);
            }
        }
    }

    return difComps.size();
}

void dilation(int * modelSrc, int * modelDst, int width, int height, uchar * matrix, int size) {
    for (int i = 0; i < height; i++) {
        for (int j = 0; j < width; j++) {
            int pos = i * width + j;
            if (modelSrc[pos] == 1) {
                for (int k = 0; k < size; k++) {
                    for (int m = 0; m < size; m++) {
                        if (matrix[k * size + m] == 1) {
                            int x = j - size / 2 + m;
                            int y = i - size / 2 + k;
                            if (x >= 0 && x < width && y >= 0 && y < height)
                                modelDst[y * width + x] = matrix[k * size + m];
                        }
                    }
                }
            }
        }
    }
}

int main() {
    cv::Mat image;
    image = cv::imread("img/img2.jpg", cv::IMREAD_COLOR);
    if(!image.data )
    {
        std::cerr <<  "Could not open or find the image" << std::endl ;
        return -1;
    }
    cv::Mat grayScaleImage(image.size(), CV_8U);
    convertToGrayScale(image, grayScaleImage);
    int model[image.rows * image.cols];
    int tmp[image.rows * image.cols];
    getBinaryModel(grayScaleImage, model);

    /*for (int i = 0; i < image.rows * image.cols; i++)
    {
        std::cerr <<  model[i] << " ";
        if (i % image.rows == 0)
            std::cerr << std::endl;
    }*/


    uchar matrix[9][9] = {
            0,0,0,1,1,1,0,0,0,
            0,0,1,1,1,1,1,0,0,
            0,1,1,1,1,1,1,1,0,
            1,1,1,1,1,1,1,1,1,
            1,1,1,1,1,1,1,1,1,
            1,1,1,1,1,1,1,1,1,
            0,1,1,1,1,1,1,1,0,
            0,0,1,1,1,1,1,0,0,
            0,0,0,1,1,1,0,0,0,
    };

    dilation(model, tmp, image.cols, image.rows, &matrix[0][0], 9);

    cv::Mat modelImage(image.size(), CV_8U);
    for (int i = 0; i < modelImage.rows; i++) {
        for (int j = 0; j < modelImage.cols; j++) {
            if (model[i * modelImage.cols + j] == 1) {
                modelImage.at<uchar>(i, j) = 255;
            } else {
                modelImage.at<uchar>(i, j) = 0;
            }
        }
    }

    cv::Mat afterIm(image.size(), CV_8U);
    for (int i = 0; i < afterIm.rows; i++) {
        for (int j = 0; j < afterIm.cols; j++) {
            if (tmp[i * afterIm.cols + j] == 1) {
                afterIm.at<uchar>(i, j) = 255;
            } else {
                afterIm.at<uchar>(i, j) = 0;
            }
        }
    }

    std::cerr << countAreas(tmp, image.cols, image.rows);

    /*namedWindow("before", cv::WINDOW_AUTOSIZE);
    imshow("before", modelImage);*/

    namedWindow("after", cv::WINDOW_AUTOSIZE);
    imshow("after", afterIm);

    cv::waitKey(0);

    return 0;
}