/*
 * transition_vqm.c -- video quality measurement
 * Copyright (c) 2012-2024 Dan Dennedy <dan@dennedy.org>
 * Core psnr and ssim routines based on code from
 *   qsnr (C) 2010 E. Oriani, ema <AT> fastwebnet <DOT> it
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>
 */

#include "common.h"
#include <framework/mlt.h>
#include <math.h>
#include <stdio.h>
#include <string.h>
#include <QColor>
#include <QFont>
#include <QImage>
#include <QPainter>
#include <QPalette>
#include <QString>

static double calc_psnr(const uint8_t *a, const uint8_t *b, int size, int bpp)
{
    double mse = 0.0;
    int n = size + 1;

    while (--n) {
        int diff = *a - *b;
        mse += diff * diff;
        a += bpp;
        b += bpp;
    }

    return 10.0 * log10(255.0 * 255.0 / (mse == 0 ? 1e-10 : mse / size));
}

static double calc_ssim(
    const uint8_t *a, const uint8_t *b, int width, int height, int window_size, int bpp)
{
    int windows_x = width / window_size;
    int windows_y = height / window_size;
    double avg = 0.0;

    if (!windows_x || !windows_y)
        return 0.0;

    // for each window
    for (int y = 0; y < windows_y; ++y)
        for (int x = 0; x < windows_x; ++x) {
            int base_offset = x * window_size + y * window_size * width;
            double ref_acc = 0.0, ref_acc_2 = 0.0, cmp_acc = 0.0, cmp_acc_2 = 0.0,
                   ref_cmp_acc = 0.0;

            // accumulate the pixel values for this window
            for (int j = 0; j < window_size; ++j)
                for (int i = 0; i < window_size; ++i) {
                    uint8_t c_a = a[bpp * (base_offset + j * width + i)];
                    uint8_t c_b = b[bpp * (base_offset + j * width + i)];
                    ref_acc += c_a;
                    ref_acc_2 += c_a * c_a;
                    cmp_acc += c_b;
                    cmp_acc_2 += c_b * c_b;
                    ref_cmp_acc += c_a * c_b;
                }

            // compute the SSIM for this window
            // http://en.wikipedia.org/wiki/SSIM
            // http://en.wikipedia.org/wiki/Variance
            // http://en.wikipedia.org/wiki/Covariance
            double n_samples = window_size * window_size, ref_avg = ref_acc / n_samples,
                   ref_var = ref_acc_2 / n_samples - ref_avg * ref_avg,
                   cmp_avg = cmp_acc / n_samples,
                   cmp_var = cmp_acc_2 / n_samples - cmp_avg * cmp_avg,
                   ref_cmp_cov = ref_cmp_acc / n_samples - ref_avg * cmp_avg,
                   c1 = 6.5025, // (0.01*255.0)^2
                c2 = 58.5225,   // (0.03*255)^2
                ssim_num = (2.0 * ref_avg * cmp_avg + c1) * (2.0 * ref_cmp_cov + c2),
                   ssim_den = (ref_avg * ref_avg + cmp_avg * cmp_avg + c1)
                              * (ref_var + cmp_var + c2);

            // accumulate the SSIM
            avg += ssim_num / ssim_den;
        }

    // return the average SSIM
    return avg / windows_x / windows_y;
}

static int get_image(mlt_frame a_frame,
                     uint8_t **image,
                     mlt_image_format *format,
                     int *width,
                     int *height,
                     int writable)
{
    mlt_frame b_frame = mlt_frame_pop_frame(a_frame);
    mlt_properties properties = MLT_FRAME_PROPERTIES(a_frame);
    mlt_transition transition = MLT_TRANSITION(mlt_frame_pop_service(a_frame));
    uint8_t *b_image;
    int window_size = mlt_properties_get_int(MLT_TRANSITION_PROPERTIES(transition), "window_size");
    double psnr[3], ssim[3];

    *format = mlt_image_yuv422;
    mlt_frame_get_image(b_frame, &b_image, format, width, height, writable);
    mlt_frame_get_image(a_frame, image, format, width, height, writable);

    psnr[0] = calc_psnr(*image, b_image, *width * *height, 2);
    psnr[1] = calc_psnr(*image + 1, b_image + 1, *width * *height / 2, 4);
    psnr[2] = calc_psnr(*image + 3, b_image + 3, *width * *height / 2, 4);
    ssim[0] = calc_ssim(*image, b_image, *width, *height, window_size, 2);
    ssim[1] = calc_ssim(*image + 1, b_image + 1, *width / 2, *height, window_size, 4);
    ssim[2] = calc_ssim(*image + 3, b_image + 3, *width / 2, *height, window_size, 4);
    mlt_properties_set_double(properties, "meta.vqm.psnr.y", psnr[0]);
    mlt_properties_set_double(properties, "meta.vqm.psnr.cb", psnr[1]);
    mlt_properties_set_double(properties, "meta.vqm.psnr.cr", psnr[2]);
    mlt_properties_set_double(properties, "meta.vqm.ssim.y", ssim[0]);
    mlt_properties_set_double(properties, "meta.vqm.ssim.cb", ssim[1]);
    mlt_properties_set_double(properties, "meta.vqm.ssim.cr", ssim[2]);
    printf("%05d %05.2f %05.2f %05.2f %5.3f %5.3f %5.3f\n",
           mlt_frame_get_position(a_frame),
           psnr[0],
           psnr[1],
           psnr[2],
           ssim[0],
           ssim[1],
           ssim[2]);

    // copy the B frame to the bottom of the A frame for comparison
    window_size = mlt_image_format_size(*format, *width, *height, NULL) / 2;
    memcpy(*image + window_size, b_image + window_size, window_size);

    if (!mlt_properties_get_int(MLT_TRANSITION_PROPERTIES(transition), "render"))
        return 0;

    // get RGBA image for Qt drawing
    *format = mlt_image_rgba;
    mlt_frame_get_image(a_frame, image, format, width, height, 1);

    // convert mlt image to qimage
    QImage img;
    convert_mlt_to_qimage_rgba(*image, &img, *width, *height);

    // setup Qt drawing
    QPainter painter(&img);
    painter.setRenderHints(QPainter::Antialiasing | QPainter::TextAntialiasing
#if QT_VERSION < QT_VERSION_CHECK(6, 0, 0)
                           | QPainter::HighQualityAntialiasing
#endif
    );
    // draw some stuff with Qt
    QFont font;
    font.setBold(true);
    font.setPointSize(30 * *height / 1080);
    painter.setPen(Qt::black);
    painter.drawLine(0, *height / 2 + 1, *width, *height / 2);
    painter.setPen(Qt::white);
    painter.drawLine(0, *height / 2 - 1, *width, *height / 2);
    painter.setFont(font);
    auto s = QString::asprintf("Frame: %05d\nPSNR:   %05.2f (Y) %05.2f (Cb) %05.2f (Cr)\nSSIM:    "
                               "%5.3f (Y) %5.3f (Cb) %5.3f (Cr)",
                               mlt_frame_get_position(a_frame),
                               psnr[0],
                               psnr[1],
                               psnr[2],
                               ssim[0],
                               ssim[1],
                               ssim[2]);
    painter.setPen(Qt::black);
    painter.drawText(52, *height * 8 / 10 + 2, *width, *height, 0, s);
    painter.setPen(Qt::white);
    painter.drawText(50, *height * 8 / 10, *width, *height, 0, s);

    // finish Qt drawing
    painter.end();
    convert_qimage_to_mlt_rgba(&img, *image, *width, *height);

    return 0;
}

static mlt_frame process(mlt_transition transition, mlt_frame a_frame, mlt_frame b_frame)
{
    mlt_frame_push_service(a_frame, transition);
    mlt_frame_push_frame(a_frame, b_frame);
    mlt_frame_push_get_image(a_frame, get_image);

    return a_frame;
}

extern "C" {

mlt_transition transition_vqm_init(mlt_profile profile,
                                   mlt_service_type type,
                                   const char *id,
                                   void *arg)
{
    mlt_transition transition = mlt_transition_new();

    if (transition) {
        mlt_properties properties = MLT_TRANSITION_PROPERTIES(transition);

        if (!createQApplicationIfNeeded(MLT_TRANSITION_SERVICE(transition))) {
            mlt_transition_close(transition);
            return NULL;
        }
        transition->process = process;
        mlt_properties_set_int(properties, "_transition_type", 1); // video only
        mlt_properties_set_int(properties, "window_size", 8);
        printf("frame psnr[Y] psnr[Cb] psnr[Cr] ssim[Y] ssim[Cb] ssim[Cr]\n");
    }

    return transition;
}

} // extern "C"
