clear all;
clc;

% Tests all images in folder 'folder'. Must be in BMP format.
% For each call of main, Quality factor is second argument, and third argument is quantization table value in 2011 approach.

folder = "tid2008";
printResults(folder,25,64)
printResults(folder,25,32)
printResults(folder,50,16)
printResults(folder,80,8)

function results = testFunction(folder,qf,fixedfactor)
    files = dir(fullfile(folder,'*.BMP'));
    results = zeros([3,5]);
    for k = 1:length(files)
        fileName = fullfile(folder,files(k).name);
        results = results + main(fileName,qf,fixedfactor);
        % results(2,1) = results(1,1) - results(2,1)
    end
    results = results./length(files);
end

function tbl = printResults(folder,qf,fixedfactor)
    T = testFunction(folder,qf,fixedfactor);
     tbl = sprintf("| | | Size | MS-SSIM | FSIMc | MSE | PSNR | \n" + ...
        "|---+---+----+--+--+--+---|\n" + ...
        "| QF = %i | Unmodified | %i | %.3f | %.3f | %.3f | %.3f |\n" + ...
        "| | textbf{Our tables} | %i | %.3f | %.3f | %.3f | %.3f | \n" + ...
        "| | cite{DBLP:conf/fc/JohnsonWL11} (%i) | %i | %.3f | %.3f | %.3f | %.3f |", ...
        qf,round(T(1,1)), T(1,2) , T(1,3) , T(1,4) , T(1,5), ...
        round(T(2,1)), T(2,2) , T(2,3) , T(2,4) , T(2,5), ...
        fixedfactor,round(T(3,1)), T(3,2) , T(3,3) , T(3,4) , T(3,5));
end
    
%testimg = "I07.BMP";
%main(testimg,50,32)
% main(testimg,50,16)
% main(testimg,80,16)

function table = main(img,qf,fixedfactor)

    %%% Standard quantization tables Q = 50
    % https://www.w3.org/Graphics/JPEG/itu-t81.pdf
    std_quantizaton_luminace = [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];

    std_quantizaton_chrominance = [16,  18,  24,  47,  99,  99,  99,  99;
                                   18,  21,  26,  66,  99,  99,  99,  99;
                                   24,  26,  56,  99,  99,  99,  99,  99;
                                   47,  66,  99,  99,  99,  99,  99,  99;
                                   99,  99,  99,  99,  99,  99,  99,  99;
                                   99,  99,  99,  99,  99,  99,  99,  99;
                                   99,  99,  99,  99,  99,  99,  99,  99;
                                   99,  99,  99,  99,  99,  99,  99,  99];

    %%% Code for generating quantization tables for other quality factors
    % Follows approach from cjpeg (see jcparam.c)
    Quality_Factor = qf;
    if (Quality_Factor < 50)
        S = 5000/Quality_Factor;
    else
        S = 200 - 2*Quality_Factor;
    end
    quantization_luminance = floor((S*std_quantizaton_luminace + 50)/100);
    quantization_luminance(quantization_luminance == 0) = 1;
    quantization_chrominance = floor((S*std_quantizaton_chrominance + 50)/100);
    quantization_chrominance(quantization_chrominance == 0) = 1;

    %%% Classic tables
    classic_quantization_table_string = generate_quantization_table_string("# Standard quantification tables\n", quantization_luminance, quantization_chrominance);
    fid = fopen('classic_tables.txt', 'wt');
    fprintf(fid, classic_quantization_table_string);
    fclose(fid);
    system(sprintf("cjpeg -optimize -baseline -qtables classic_tables.txt -qslots 0,1,1 -sample 2x1 -outfile %s__classic__.jpg %s", img, img));

    %%% Our tables (only powers of two)
    % Find quantization table giving size closest to classic version
    out_quantization_table_string_match_size = closest_in_size(img, quantization_luminance, quantization_chrominance);
    fid = fopen('our_tables_size.txt', 'wt');
    fprintf(fid, out_quantization_table_string_match_size);
    fclose(fid);
    system(sprintf("cjpeg -optimize -baseline -qtables our_tables_size.txt -qslots 0,1,1 -sample 2x1 -outfile %s__our_size__.jpg %s", img, img));

    %%% "Bad" '11 tables
    % Note that there are so few options here that it doesn't make sense to automate testing. Very quickly one can find the two closest
    bad11_quantization_luminance = ones(8).*fixedfactor;
    bad11_quantization_chrominance = ones(8).*fixedfactor;
    bad11_quantization_table_string = generate_quantization_table_string("# Custom quantification tables using single power of two\n", bad11_quantization_luminance, bad11_quantization_chrominance);
    fid = fopen('11_tables.txt', 'wt');
    fprintf(fid, bad11_quantization_table_string);
    fclose(fid);
    system(sprintf("cjpeg -optimize -baseline -qtables 11_tables.txt -qslots 0,1,1 -sample 2x1 -outfile %s__11__.jpg %s", img, img));

    %%% Analysis
    I_full = imread(img);
    I_compressed_classic = imread(sprintf("%s__classic__.jpg", img));
    I_compressed_our = imread(sprintf("%s__our_size__.jpg", img));
    I_compressed_11 = imread(sprintf("%s__11__.jpg", img));

    % Mean Squared Error
    mse_results = [NaN, NaN, NaN];
    mse_results(1) = immse(I_full,I_compressed_classic);
    mse_results(2) = immse(I_full,I_compressed_our);
    mse_results(3) = immse(I_full,I_compressed_11);

    % peak noise to signal rate
    psnr_results = [NaN, NaN, NaN];
    psnr_results(1) = psnr(I_compressed_classic,I_full);
    psnr_results(2) = psnr(I_compressed_our,I_full);
    psnr_results(3) = psnr(I_compressed_11,I_full);

    % multiscale structural similarity
    msss_results = [NaN, NaN, NaN];
    msss_results(1) = mean(multissim(I_compressed_classic,I_full));
    msss_results(2) = mean(multissim(I_compressed_our,I_full));
    msss_results(3) = mean(multissim(I_compressed_11,I_full));

    % Feature Similarity
    % Uses the FSIM library
    % https://ieeexplore.ieee.org/document/5705575
    % http://www.comp.polyu.edu.hk/~cslzhang/IQA/FSIM/FSIM.htm
    fsim_results = [NaN, NaN, NaN];
    [~,fsimc] = FSIM(I_full,I_compressed_classic);
    fsim_results(1) = fsimc;
    [~,fsimz] = FSIM(I_full,I_compressed_our);
    fsim_results(2) = fsimc;
    [~,fsimc] = FSIM(I_full,I_compressed_11);
    fsim_results(3) = fsimc;

    % size
    size_classic = dir(sprintf("%s__classic__.jpg", img)).bytes;
    size_our = dir(sprintf("%s__our_size__.jpg", img)).bytes;
    size_11 = dir(sprintf("%s__11__.jpg", img)).bytes;

    % Pretty print - made to be formatted by emacs' orgmode
    table = [[size_classic,size_our,size_11].',msss_results.',fsim_results.',mse_results.',psnr_results.'];
    % table = sprintf("| QF = %i | Size | MS-SSIM | FSIMc | MSE | PSNR | \n |- | | | | | |\n | Unmodified | %i | %.3f | %.3f | %.3f | %.3f |\n| cite{DBLP:conf/fc/JohnsonWL11} (%i) | %i | %.3f | %.3f | %.3f | %.3f | \n | Our compression | %i | %.3f | %.3f | %.3f | %.3f |",qf,size_classic,msss_results(1), fsim_results(1), mse_results(1), psnr_results(1),fixedfactor,size_11,msss_results(3), fsim_results(3), mse_results(3), psnr_results(3),size_our,msss_results(2), fsim_results(2), mse_results(2), psnr_results(2));
end

% Round x to a near power of two (up or down), p specifying where to make the cut
function y = to_power2(x, p)
    threshold = (1+p)*2^(floor(log2(x)));
    if x < threshold
        y = 2^(floor(log2(x)));
    else
        y = 2^(ceil(log2(x)));
    end
end

% Convert M to a string of the format cjpeg understands
function tmp = table_to_string(M)
    tmp = '';
    for i = 1:8
        for j = 1:8
            tmp = sprintf("%s%s ", tmp, num2str(M(i,j)));
        end
        tmp = strcat(tmp, '\n');
    end
end

% Generate quantization table string for inputted tables
function tmp = generate_quantization_table_string(name, luminance_table, chrominance_table)
    standard_comment_string = sprintf('# Created by Simon Erfurth\n# %s\n# Example use:\n# cjpeg -optimize -baseline -qtables this_file.txt -qslots 0,1,1 -sample 2x1 -outfile out.jpg in.bmp\n',datetime("today"));
    quantization_luminance_string = sprintf("# Luma\n%s",table_to_string(luminance_table));
    quantization_chrominance_string = sprintf("# Chroma\n%s", table_to_string(chrominance_table));
    tmp = sprintf("%s%s%s%s", name, standard_comment_string, quantization_luminance_string, quantization_chrominance_string);
end

% Functions to find quantization table giving size closest to classic version
function table_string = closest_in_size(img, quantization_luminance, quantization_chrominance)
    target_size = dir(sprintf("%s__classic__.jpg", img)).bytes;
    
    p_max = 1;
    q_lum_max = generate_quantization_table(quantization_luminance,p_max);
    q_chr_max = generate_quantization_table(quantization_chrominance,p_max);
    
    p_min = 0;
    q_lum_min = generate_quantization_table(quantization_luminance,p_min);
    q_chr_min = generate_quantization_table(quantization_chrominance,p_min);

    % And handle special case where the following loop doesn't execute at all (e.g. if Q = 100)
    q_lum_new = q_lum_max;
    q_chr_new = q_chr_max;
    
    while (any(any(q_lum_min ~= q_lum_max)) || any(any(q_chr_min ~= q_chr_max))) && (p_max - p_min) > 0.000001
        p = (p_max + p_min)/2;
        q_lum_new = generate_quantization_table(quantization_luminance,p);
        q_chr_new = generate_quantization_table(quantization_chrominance,p);
        img_file_name = create_image(q_lum_new,q_chr_new,img);
        current_size = dir(img_file_name).bytes;
        if current_size > target_size
            p_max = p;
            q_lum_max = q_lum_new;
            q_chr_max = q_chr_new;
        else
            p_min = p;
            q_lum_min = q_lum_new;
            q_chr_min = q_chr_new;
        end
    end
    table_string = generate_quantization_table_string("# Custom quantification tables using powers of two\n # Optimized to give same size as standard variant\n", q_lum_new, q_chr_new);
end

function new_table = generate_quantization_table(start_table,p)
    new_table = ones(8);
    for i = 1:8
        for j = 1:8
            new_table(i,j) = to_power2(start_table(i,j),p);
        end
    end
end

function out_file_name = create_image(quantization_luminance, quantization_chrominance, img_file_name)
    tmp_quantization_table_string = generate_quantization_table_string("# Temporary quantification tables\n", quantization_luminance, quantization_chrominance);
    fid = fopen('tmp_tables.txt', 'wt');
    fprintf(fid, tmp_quantization_table_string);
    fclose(fid);
    system(sprintf("cjpeg -optimize -baseline -qtables tmp_tables.txt -qslots 0,1,1 -sample 2x1 -outfile %s__tmp__.jpg %s", img_file_name, img_file_name));
    out_file_name = sprintf("%s__tmp__.jpg", img_file_name);
end