#include <iostream>
#include <string>
#include "libqr.h"
#include "opencv2/objdetect.hpp"
#include "opencv2/wechat_qrcode.hpp"
#include "string_format.h"
using namespace std;
using namespace cv;

static
vector<Point> transform_image(Mat &in, vector<Point> qr_points, Mat &out)
{
    Mat src = (Mat_<float>(4, 2) <<
               qr_points[0].x, qr_points[0].y,
               qr_points[1].x, qr_points[1].y,
               qr_points[2].x, qr_points[2].y,
               qr_points[3].x, qr_points[3].y
               );

    int min_x = qr_points[0].x;
    int min_y = qr_points[0].y;
    int max_x = qr_points[0].x;
    int max_y = qr_points[0].y;
    for (auto p: qr_points) {
        min_x = min(p.x, min_x);
        min_y = min(p.y, min_y);
        max_x = max(p.x, max_x);
        max_y = max(p.y, max_y);
    }
    Mat dst = (Mat_<float>(4, 2) <<
               min_x, min_y,
               max_x, min_y,
               max_x, max_y,
               min_x, max_y);

    Mat m = getPerspectiveTransform(src, dst);
    warpPerspective(in, out, m, in.size());
    vector<Point> ret;
    ret.push_back(Point(min_x, min_y));
    ret.push_back(Point(max_x, min_y));
    ret.push_back(Point(max_x, max_y));
    ret.push_back(Point(min_x, max_y));
    return ret;
}

bool detect_qr(ProcessState &ps, float margin_ratio, bool warp, string &err)
{
#if WECHAT_QRCODE_USE_MODEL
    auto wr = wechat_qrcode::WeChatQRCode(
            "wechat_qrcode/detect.prototxt",
            "wechat_qrcode/detect.caffemodel",
            "wechat_qrcode/sr.prototxt",
            "wechat_qrcode/sr.caffemodel");
#else
    auto wr = wechat_qrcode::WeChatQRCode();
#endif
    vector<Mat> qrs;
    auto r = wr.detectAndDecode(ps.preprocessed, qrs);

    if (!r.size()) {
        err = "qr not detected";
        return false;
    }

    ps.qrcode = r[0];
    auto rect = qrs[0];
    vector<Point> qr_points;
    qr_points.push_back(Point(rect.at<float>(0, 0) / ps.scale, rect.at<float>(0, 1) / ps.scale));
    qr_points.push_back(Point(rect.at<float>(1, 0) / ps.scale, rect.at<float>(1, 1) / ps.scale));
    qr_points.push_back(Point(rect.at<float>(2, 0) / ps.scale, rect.at<float>(2, 1) / ps.scale));
    qr_points.push_back(Point(rect.at<float>(3, 0) / ps.scale, rect.at<float>(3, 1) / ps.scale));
    ps.qr_points = qr_points;
    Mat warped;
    vector<Point> warped_qr_points;
    if (warp) {
        warped_qr_points = transform_image(*ps.orig, qr_points, warped);
    } else {
        warped = *ps.orig;
        warped_qr_points = qr_points;
    }
    int min_x = warped_qr_points[0].x;
    int min_y = warped_qr_points[0].y;
    int max_x = min_x;
    int max_y = min_y;
    for (auto p: warped_qr_points) {
        min_x = min(p.x, min_x);
        min_y = min(p.y, min_y);
        max_x = max(p.x, max_x);
        max_y = max(p.y, max_y);
    }
    int margin = (max_x - min_x) * margin_ratio;
    if (min_y < margin || min_x < margin || max_x + margin >= warped.cols || max_y + margin >= warped.rows) {
        err = "qr margin too small";
        return false;
    }
    int qr_width = max_x - min_x;
    int qr_height = max_y - min_y;
    if (qr_width < 200 && qr_height < 200 && qr_width < ps.orig->cols * 0.5 && qr_height < ps.orig->rows * 0.5) {
        printf("(%d, %d) in (%d, %d)\n", qr_width, qr_height, ps.orig->cols, ps.orig->rows);
        err = "qr too small";
        return false;
    }

    Rect qr_rect(min_x, min_y, max_x - min_x, max_y - min_y);
    ps.qr_straighten = warped(qr_rect);
    Rect qr_with_margin_rect(min_x - margin, min_y - margin,
                             max_x - min_x + margin * 2,
                             max_y - min_y + margin * 2);
    ps.straighten = warped(qr_with_margin_rect);
    Mat g;
    cvtColor(ps.straighten, g, COLOR_BGR2GRAY);
    equalizeHist(g, g);

    Rect left_mid_rect;
    left_mid_rect.x = 0;
    left_mid_rect.y = ps.straighten.rows / 6;
    left_mid_rect.width = ps.straighten.cols / 3;
    left_mid_rect.height = ps.straighten.rows / 3;
    ps.left_mid_area = ps.straighten(left_mid_rect);
    return true;
}

bool preprocess(ProcessState &ps)
{
    Mat gray;
    cvtColor(*ps.orig, gray, COLOR_BGR2GRAY);
    ps.scale = 1.0;
    const float size_cap = 512;
    if (ps.orig->rows > size_cap) {
        ps.scale = size_cap / ps.orig->rows;
    }
    if (ps.orig->cols > ps.orig->rows && ps.orig->cols > size_cap) {
        ps.scale = size_cap / ps.orig->cols;
    }
    resize(gray, ps.preprocessed, Size(), ps.scale, ps.scale);
    return true;
}

static
bool check_blur_by_sobel(ProcessState &ps, Mat &gray, string &err)
{
    const float thres = ps.sobel_thres;
    Mat sobel_x, sobel_y;
    Sobel(gray, sobel_x, CV_64F, 1, 0, 3);
    Sobel(gray, sobel_y, CV_64F, 0, 1, 3);

    Mat magnitude;
    magnitude = sobel_x.mul(sobel_x) + sobel_y.mul(sobel_y); // Gx^2 + Gy^2
    sqrt(magnitude, magnitude); // sqrt(Gx^2 + Gy^2)

    Scalar meanMagnitude = mean(magnitude);
    double sharpness = meanMagnitude[0] / (gray.rows * gray.cols) * 1000;
    ps.clarity = sharpness;
    if (sharpness < thres) {
        err = string_format("image too blurry: %lf < %lf", sharpness, thres);
        return false;
    }
    return true;
}

static
bool check_sharpness(ProcessState &ps, Mat &gray, int method, string &err)
{
    return check_blur_by_sobel(ps, gray, err);
}

#define COUNT_COMPONENTS 0
#if COUNT_COMPONENTS
static bool is_valid_pattern(Mat &img)
{
    Mat labels;
    Mat stats;
    Mat centroids;
    connectedComponentsWithStats(img, labels, stats, centroids);
    int valid = 0;
    for (auto i = 0; i < stats.rows; i++) {
        int area = stats.at<int>(i, CC_STAT_AREA);
        if (area > 5) {
            valid++;
        }
    }

    return valid > 25;
}
#endif

bool emblem_dot_angle(ProcessState &ps, InputArray in, float &angle, string &qrcode, string &err)
{
    try {

        ps.orig = (Mat *)in.getObj();
        preprocess(ps);

        if (!detect_qr(ps, 0, true, err)) {
            qrcode = ps.qrcode;
            err = "detect_qr: " + err;
            return false;
        }

        qrcode = ps.qrcode;

        Mat lma_gray;
        cvtColor(ps.left_mid_area, lma_gray, COLOR_BGR2GRAY);
        if (!check_sharpness(ps, lma_gray, ps.sharpness_method, err)) {
            return false;
        }

        angle = 0;
        return true;
    } catch (const std::exception &exc) {
        std::cout << exc.what() << std::endl;
        err = "exception";
        return false;
    } catch (...) {
        err = "unknown error";
        return false;
    }
}

static
Mat adaptive_gray(Mat &img)
{
    Mat ret;
    Mat mean, stddev;
    Mat channels[3];
    Mat hsv_img;

    meanStdDev(img, mean, stddev);

    int bgr_max_std_channel = 0;
    float bgr_max_std = stddev.at<float>(0, 0);
    for (int i = 1; i < 3; i++) {
        auto nv = stddev.at<float>(0, i);
        if (nv > bgr_max_std_channel) {
            bgr_max_std_channel = i;
            bgr_max_std = nv;
        }
    }
    cvtColor(img, hsv_img, COLOR_BGR2HSV);
    meanStdDev(img, hsv_img, stddev);
    int hsv_max_std_channel = 0;
    float hsv_max_std = stddev.at<float>(0, 0);
    for (int i = 1; i < 3; i++) {
        auto nv = stddev.at<float>(0, i);
        if (nv > hsv_max_std_channel) {
            hsv_max_std_channel = i;
            hsv_max_std = nv;
        }
    }
    if (hsv_max_std > bgr_max_std) {
        split(hsv_img, channels);
        printf("using hsv channel %d\n", hsv_max_std_channel);
        ret = channels[hsv_max_std_channel];
    } else {
        split(img, channels);
        printf("using rgb channel %d\n", bgr_max_std_channel);
        ret = channels[bgr_max_std_channel];
    }
    return ret;
}

static
bool cell_in_bg(int cell_x, int cell_y)
{
    return
        (cell_x == 1 && (cell_y > 0 && cell_y < 6)) ||
        (cell_x == 2 && (cell_y == 1 || cell_y == 5)) ||
        (cell_x == 3 && (cell_y == 1 || cell_y == 5)) ||
        (cell_x == 4 && (cell_y == 1 || cell_y == 5)) ||
        (cell_x == 5 && cell_y > 0 && cell_y < 6)
        ;
}

static
bool roi_in_bg(int w, int h, Point p)
{
    int cell_x = p.x * 7 / w;
    int cell_y = p.y * 7 / h;
    return cell_in_bg(cell_x, cell_y);
}

static
void roi_mask(Mat &img, int margin_pct)
{
    int counts[256] = { 0 };
    for (int i = 0; i < img.cols; i++) {
        for (int j = 0; j < img.rows; j++) {
            uint8_t p = img.at<uint8_t>(Point(i, j));
            counts[p]++;
        }
    }
    int cut = 20;
    int seen = 0;
    int total = img.cols * img.rows;
    int p05, p95;
    for (p05 = 0; seen < total * cut / 100 && p05 < 256; p05++) {
        seen += counts[p05];
    }

    seen = 0;
    for (p95 = 0; seen < total * (100 - cut) / 100 && p95 < 256; p95++) {
        seen += counts[p95];
    }

    printf("p05: %d, p95: %d\n", p05, p95);
    int cap = (p95 - p05) * margin_pct / 100;
    int min_thres = p05 + cap;
    int max_thres = p95 - cap;

    for (int i = 0; i < img.cols; i++) {
        for (int j = 0; j < img.rows; j++) {
            auto pos = Point(i, j);
            uint8_t p = img.at<uint8_t>(pos);
            if (!roi_in_bg(img.cols, img.rows, pos)) {
                img.at<uint8_t>(pos) = 0;
            } else if (p < min_thres) {
                img.at<uint8_t>(pos) = 0;
            } else if (p > max_thres) {
                img.at<uint8_t>(pos) = 0;
            } else {
                img.at<uint8_t>(pos) = 255;
            }
        }
    }
}

static
vector<float> roi_extract_features(Mat &img)
{
    vector<int> ones(49, 0);
    vector<int> zeroes(49, 0);
    for (int i = 0; i < img.cols; i++) {
        for (int j = 0; j < img.rows; j++) {
            auto pos = Point(i, j);
            int cell_x = pos.x * 7 / img.cols;
            int cell_y = pos.y * 7 / img.rows;
            int idx = cell_y * 7 + cell_x;
            assert(idx < 49);

            uint8_t p = img.at<uint8_t>(pos);
            if (p) {
                ones[idx]++;
            } else {
                zeroes[idx]++;
            }
        }
    }
    printf("ones:\n");
    for (int i = 0; i < 49; i++) {
        printf("%d ", ones[i]);
    }
    printf("\n");
    vector<float> ret;
    for (int i = 0; i < 49; i++) {
        int cell_x = i % 7;
        int cell_y = i / 7;
        if (!cell_in_bg(cell_x, cell_y)) {
            continue;
        }
        if (ones[i] || zeroes[i]) {
            ret.push_back(ones[i] / (float)(ones[i] + zeroes[i]));
        } else {
            ret.push_back(0);
        }
    }
    return ret;
}

static
float mean(vector<float> &a)
{
    float sum = 0;

    if (!a.size()) {
        return 0;
    }

    for (auto x: a) {
        sum += x;
    }

    return sum / a.size();
}

static
float covariance(vector<float> &a, vector<float> &b)
{
    float mean_a = mean(a);
    float mean_b = mean(b);
    float ret = 0;

    if (a.size() != b.size()) return 0;

    for (size_t i = 0; i < a.size(); i++) {
        ret += (a[i] - mean_a) * (b[i] - mean_b);
    }
    return ret;
}

static inline
bool valid_point(Mat &a, Point p)
{
    return p.x > 0 && p.x < a.cols && p.y > 0 && p.y < a.rows;
}

static inline
bool fuzzy_pixel_match(Mat &a, Point pa, Mat &b, Point pb)
{
    if (!valid_point(a, pa) || !valid_point(b, pb)) return false;
    return a.at<uint8_t>(pa) == b.at<uint8_t>(pb);
}

static
int fuzzy_pixel_cmp(Mat &b, Mat &a)
{
    int ret = 0;
    int w = a.cols;
    int h = a.rows;
    assert(a.cols == b.cols);
    assert(a.rows == b.rows);
    for (int i = 0; i < w; i++) {
        for (int j = 0; j < h; j++) {
            Point p(i, j);
            if (!roi_in_bg(w, h, p)) {
                ret++;
                continue;
            }
            bool same = false;
            int fuzziness = 1;
            for (int ii = -fuzziness; ii <= fuzziness; ii++) {
                for (int jj = -fuzziness; jj <= fuzziness; jj++) {
                    if (fuzzy_pixel_match(a, p, b, Point(i + ii, j + jj))) {
                        same = true;
                        goto out;
                    }
                }
            }
out:
            ret += same ? 1 : 0;
        }
    }
    return ret;
}

double emblem_roi_similarity(SimilarityAlg alg, InputArray std_in, InputArray frame_roi_in, string &err)
{
    Mat stdm = *(Mat *)std_in.getObj();
    Mat frame_roi = *(Mat *)frame_roi_in.getObj();
    err = "";

    Mat frame_gray = adaptive_gray(frame_roi);
    Mat std_gray = adaptive_gray(stdm);

    resize(frame_gray, frame_gray, std_gray.size());

    double ret = 0;
    Mat frame = frame_gray.clone();
    Mat std = std_gray.clone();

    roi_mask(frame, 20);
    roi_mask(std, 30);

    double same = fuzzy_pixel_cmp(frame, std);
    double total = frame.rows * frame.cols;
    double sim = same / total;
    printf("same: %lf, total: %lf, sim: %lf\n", same, total, sim);

    auto std_feature = roi_extract_features(std);
    auto frame_feature = roi_extract_features(frame);

    printf("\nstd:");
    for (auto x: std_feature) {
        printf("%.2lf ", x * 100);
    }

    printf("\nfrm:");
    for (auto x: frame_feature) {
        printf("%.2lf ", x * 100);
    }
    printf("\n");

    double cov = covariance(std_feature, frame_feature);
    printf("cov: %lf\n", cov);
    double t = cov * sim;
    ret = ret > t ? ret : t;
    return ret;
}