#include <iostream>
#define _USE_MATH_DEFINES
#include <cmath>
#include <opencv2/highgui/highgui.hpp>

// Velikost bloku bude 8x8 pixelu
static const int N = 8;


// JPEG kvantizacni matice
static double coeffsQ[] = {
        16,  11,  10,  16,  24,  40,  51,  61,
        12,  12,  14,  19,  26,  58,  60,  55,
        14,  13,  16,  24,  40,  57,  69,  56,
        14,  17,  22,  29,  51,  87,  80,  62,
        18,  22,  37,  56,  68, 109, 103,  77,
        24,  35,  55,  64,  81, 104, 113,  92,
        49,  64,  78,  87, 103, 121, 120, 101,
        72,  92,  95,  98, 112, 100, 103,  99
    };

/**
 * Dopredna 2D DCT
 */
static void fdct2D(const cv::Mat &src, cv::Mat &dst) {
    for (unsigned short v = 0; v < N; v++) {
        for (unsigned short u = 0; u < N; u++) {
            double g = 0;
            //u dimenze
            double au = (u == 0) ? 1.0 / sqrt(N) : sqrt(2.0 / N);
            //v dimenze
            double av =  (v == 0) ? 1.0 / sqrt(N) : sqrt(2.0 / N);

            for (unsigned short y = 0; y < N; y++) {
                for (unsigned short x = 0; x < N; x++) {
                    g += src.at<double>(y, x) * au * av * cos((M_PI / N) * (x + 0.5) * u) * cos((M_PI / N) * (y + .5) * v);
                }
            dst.at<double>(v, u) = g;
            }
        }
    }

    // OpenCV DCT - ZAKOMENTOVAT!
    //cv::dct(src, dst);
}

/**
 * Zpetna 2D DCT
 */
static void idct2D(const cv::Mat &src, cv::Mat &dst) {
   for (unsigned short y = 0; y < N; y++) {
        for (unsigned short x = 0; x < N; x++) {
            double f = 0;
            for (unsigned short v = 0; v < N; v++) {
                for (unsigned short u = 0; u < N; u++) {
                    //u dimenze
                    double au = (u == 0) ? 1.0 / sqrt(N) : sqrt(2.0 / N);
                    //v dimenze
                    double av =  (v == 0) ? 1.0 / sqrt(N) : sqrt(2.0 / N);
                    f += au * av * src.at<double>(v, u) * cos((M_PI / N) * (x + .5) * u) * cos((M_PI / N) * (y + .5) * v);
                }
            dst.at<double>(y, x) = f;
            }
        }
    }

    // OpenCV inverzni DCT - ZAKOMENTOVAT!
    //cv::dct(src, dst, cv::DCT_INVERSE);
}

/**
 * Kvantizace koeficientu
 */
static void quantize(const cv::Mat &src, cv::Mat &dst) {

    // Zadna kvantizace
    //src.copyTo(dst);
    double quantize;
    for (unsigned short i = 0; i < N; i++) {
        for (unsigned short j = 0; j < N; j++) {
            quantize = round(src.at<double>(i, j) / coeffsQ[(i*N)+j]);
            //dekvantizace
            dst.at<double>(i, j) =  quantize * coeffsQ[(i*N)+j];
            std::cout << dst.at<double>(i, j) << std::endl;
        }
    }
}

/**
 * Demonstrace komprese obrazu pomoci DCT
 */
static void demo(const cv::Mat &input) {
    // Sem se ulozi koeficienty DCT
    cv::Mat coeff(input.rows, input.cols, cv::DataType<double>::type);

    // Sem se ulozi koeficienty po provedeni kvantizace
    cv::Mat quant(input.rows, input.cols, cv::DataType<double>::type);

    // Tady bude dekodovany obrazek
    cv::Mat decod(input.rows, input.cols, cv::DataType<double>::type);

    // Zpracujeme obraz po blocich
    for (int y = 0; y < input.rows; y += N) {
        for (int x = 0; x < input.cols; x += N) {
            // Aktualni zpracovavany blok bude tento vyrez celeho obrazu
            cv::Rect rect(x, y, N, N);

            // Pomocne promenne, ktere reprezentuji prislusny vyrez
            cv::Mat inputBlock(input, rect);
            cv::Mat coeffBlock(coeff, rect);
            cv::Mat quantBlock(quant, rect);
            cv::Mat decodBlock(decod, rect);

            // Nejprve transformujeme blok obrazku na koeficienty pomoci DCT
            fdct2D(inputBlock, coeffBlock);

            // Pak provedeme kvantizaci techto koeficientu (ztrata informace, komprese)
            quantize(coeffBlock, quantBlock);

            // Z kvantovanych koeficientu zrekonstruujeme puvodni blok obrazu
            idct2D(quantBlock, decodBlock);
        }
    }

    // Cely postup zobrazime
    imshow("input",      cv::Mat_<uchar>(input));
    imshow("transform",  cv::abs(coeff / 15));
    imshow("quantized",  cv::abs(quant / 15));
    imshow("decoded",    cv::Mat_<uchar>(decod));
    imshow("difference", cv::Mat_<uchar>(cv::Mat(cv::abs(input-decod))));
}

int main(int argc, char **argv) {
    const char *imagename;
    imagename = argc > 1 ? argv[1] : "lenna.png";

    std::cout << "Loading " << imagename << std::endl;

    // Nacteni obrazku
    cv::Mat img = cv::Mat_<double>(cv::imread(imagename, CV_LOAD_IMAGE_GRAYSCALE));
    if (!img.data) {
        std::cerr << "Failed do load image" << std::endl;
        return -1;
    }

    // Spusteni dema
    demo(img);

    cv::waitKey();

    return 0;
}