/*
 * vf_deblendlogo.c v0.1
 *
 * Removes alpha-blended TV station logos by reversing the alpha
 * blending process based on a template.
 *
 * Some artifacts are still left due to the colour resolution lost
 * in the alpha-blending process, but the watermark is much less
 * noticable.
 *
 * Usage: -vf deblendlogo=x:y:template.png
 *    Where template.png is a PNG file containing an RGBA image
 *   representing the logo template. x and y define where on the
 *   image the template is to be placed.
 *
 * There are several caveats with this filter:
 *    * It can sometimes be hard to get a good template image.
 *     For logos which are all white, you can use a frame from
 *     your video where the picture has faded to black. However,
 *     if the logo features black elements you need to find
 *     a frame which has solid white behind the logo. Fortunately,
 *     for most channels you will only have to find the template
 *     once as long as your video resolution does not change.
 *    * The filter makes no attempt to detect portions of the
 *     input video where the logo is not present. For these portions
 *     a negative watermark will be rendered onto the image.
 *      Personally, I prefer having a dark watermark for a short
 *     period to having a light watermark for the majority of a
 *     show, but your mileage may vary.
 *    * Compression artifacts as well as the loss of colour depth
 *     caused by the original alpha-blend can cause a small
 *     amount of distortion where the logo was. This is unavoidable.
 *     You can reduce the distortion by doing the logo removal
 *     before applying "clever" compression algorithms.
 *
 * If the logo you are trying to remove has opaque or near-opaque
 * elements, you may like to try the "delogo" filter instead.
 *
 */

/* This filter requires libpng */
#include "../config.h"
#ifdef HAVE_PNG

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <png.h>
#include <inttypes.h>
#include "mp_image.h"
#include "../mp_msg.h"
#include "vf.h"
#include "img_format.h"

#include "../postproc/swscale.h"
#include "../libvo/fastmemcpy.h"

#define DEBLENDLOGO_NAME "deblendlogo"
#define DEBLENDLOGO_MODNAME "vf_" DEBLENDLOGO_NAME

struct vf_priv_s {
    int x, y, tw, th;
    unsigned char* tmpl;
    struct SwsContext* yv12torgb;
    struct SwsContext* rgbtoyv12;

};

static int
query_format(struct vf_instance_s* vf, unsigned int fmt) {
    switch (fmt) {
        case IMGFMT_ABGR:
        case IMGFMT_BGRA:
        case IMGFMT_ARGB:
        case IMGFMT_RGBA:
        case IMGFMT_YV12:
            return vf_next_query_format(vf, fmt);
        default:
            return 0;
    }
}

static int
config(struct vf_instance_s* vf,
       int width, int height, int d_width, int d_height,
       unsigned int flags, unsigned int outfmt)
{
    if (vf->priv->x < 0 || vf->priv->y < 0
            || vf->priv->x + vf->priv->tw > width
            || vf->priv->y + vf->priv->th > height) {
        mp_msg(MSGT_VFILTER,MSGL_WARN,DEBLENDLOGO_MODNAME ": The template is hanging out of the video frame!\n");
    }

    /* If YV12 format is selected, we must set up the SwScaler conversion contexts */
    if (outfmt == IMGFMT_YV12) {
        vf->priv->yv12torgb = sws_getContext(width, height, IMGFMT_YV12,
                                             width, height, IMGFMT_BGRA,
                                             0, NULL, NULL, NULL);
        vf->priv->rgbtoyv12 = sws_getContext(width, height, IMGFMT_BGRA,
                                             width, height, IMGFMT_YV12,
                                             0, NULL, NULL, NULL);
    }

    return vf_next_config(vf, width, height, d_width, d_height, flags, outfmt);
}

static int
put_image(struct vf_instance_s* vf, mp_image_t* mpi) {
    mp_image_t* dmpi;
    unsigned int bpp = mpi->bpp / 8;
    unsigned int x, y, red, green, blue;
    struct vf_priv_s* priv = (struct vf_priv_s*)vf->priv;
    struct SwsContext *scalectx;

    if (mpi->imgfmt == IMGFMT_YV12) {
        dmpi = vf_get_image(vf->next, IMGFMT_BGRA, MP_IMGTYPE_EXPORT,
                MP_IMGFLAG_ACCEPT_STRIDE | MP_IMGFLAG_PREFER_ALIGNED_STRIDE,
                mpi->w, mpi->h);

        dmpi->stride[0] = mpi->w * 4;
        dmpi->planes[0] = malloc(dmpi->stride[0] * mpi->h);

        sws_scale(priv->yv12torgb, mpi->planes, mpi->stride, 0, mpi->h, dmpi->planes, dmpi->stride);
    }
    else {
        dmpi = vf_get_image(vf->next, mpi->imgfmt, MP_IMGTYPE_TEMP,
                MP_IMGFLAG_ACCEPT_STRIDE | MP_IMGFLAG_PREFER_ALIGNED_STRIDE,
                mpi->w, mpi->h);

        memcpy_pic(dmpi->planes[0],mpi->planes[0],mpi->w*bpp, mpi->h,
               dmpi->stride[0],mpi->stride[0]);
    }

    switch (dmpi->imgfmt) {
        case IMGFMT_ABGR: {
            red = 1;
            green = 2;
            blue = 3;
        } break;
        case IMGFMT_BGRA: {
            red = 0;
            green = 1;
            blue = 2;
        } break;
        case IMGFMT_ARGB: {
            red = 1;
            green = 2;
            blue = 3;
        } break;
        case IMGFMT_RGBA: {
            red = 0;
            green = 1;
            blue = 2;
        } break;
        default: {
            static int shownerror;
            if (shownerror != 1) {
                mp_msg(MSGT_VFILTER,MSGL_ERR,DEBLENDLOGO_MODNAME ": Unsupported image format (%s).\n", vo_format_name(dmpi->imgfmt));
                shownerror = 1;
            }
            return vf_next_put_image(vf, dmpi);
        }
    }

    /*
        The equation used here is just a re-arranged version of that
        used to do alpha blending in the first place.

        w: watermark r,g,b (from the template)
        s: source r,g,b (the result)
        b: blended r,g,b (the video data)
        a: alpha (from the template)

            s = b - (w * a)
                -----------     (for each component)
                   1 - a

        vpos: memory location in the video image
        tpos: memory location in the template image

     */

    for (x = 0; x < priv->tw; x++) {
        for (y = 0; y < priv->th; y++) {
            int vpos = ((y+priv->y) * dmpi->stride[0]) + ((x+priv->x) * 4);
            int tpos = ((y * priv->tw) + x) * 4;

            double alpha = ((double)priv->tmpl[tpos + 3]) / 255.0;

            double wr = priv->tmpl[tpos],
                wg = priv->tmpl[tpos + 1],
                wb = priv->tmpl[tpos + 2];
            double br = dmpi->planes[0][vpos+red],
                bg = dmpi->planes[0][vpos+green],
                bb = dmpi->planes[0][vpos+blue];
            double sr,sg,sb;

            if (alpha == 0.0 || alpha == 1.0) continue;

            sr = ((br - (wr * alpha)) / (1.0 - alpha));
            sg = ((bg - (wg * alpha)) / (1.0 - alpha));
            sb = ((bb - (wb * alpha)) / (1.0 - alpha));
            if (sr < 0.0) sr = 0.0;
            if (sg < 0.0) sg = 0.0;
            if (sb < 0.0) sb = 0.0;

            dmpi->planes[0][vpos+red] = sr;
            dmpi->planes[0][vpos+green] = sg;
            dmpi->planes[0][vpos+blue] = sb;
        }
    }

    if (mpi->imgfmt == IMGFMT_YV12) {
        /* Convert back to YV12 */
        mp_image_t* tdmpi = dmpi;

        dmpi = vf_get_image(vf->next, IMGFMT_YV12, MP_IMGTYPE_TEMP,
                MP_IMGFLAG_ACCEPT_STRIDE | MP_IMGFLAG_PREFER_ALIGNED_STRIDE,
                tdmpi->w, tdmpi->h);

        sws_scale(priv->rgbtoyv12, tdmpi->planes, tdmpi->stride, 0, tdmpi->h, dmpi->planes, dmpi->stride);

        free(tdmpi->planes[0]);
        tdmpi->planes[0] = 0;
    }

    return vf_next_put_image(vf, dmpi);
}

static void uninit(struct vf_instance_s* vf) {
    if (vf->priv->tmpl != 0)
        free(vf->priv->tmpl);
    vf->priv->tmpl = 0;

    if (vf->priv->yv12torgb != 0)
        sws_freeContext(vf->priv->yv12torgb);
    if (vf->priv->rgbtoyv12 != 0)
        sws_freeContext(vf->priv->rgbtoyv12);

    vf->priv->yv12torgb = 0;
    vf->priv->rgbtoyv12 = 0;

    if (vf->priv != 0)
        free(vf->priv);
    vf->priv = 0;
}

static int
open(vf_instance_t* vf, char* args) {
    vf->config = config;
    vf->uninit = uninit;
    vf->query_format = query_format;
    vf->put_image = put_image;
    vf->priv = malloc(sizeof(struct vf_priv_s));
    vf->priv->x = 0;
    vf->priv->y = 0;
    vf->priv->tw = -1;
    vf->priv->th = -1;
    vf->priv->tmpl = 0;
    vf->priv->yv12torgb = 0;
    vf->priv->rgbtoyv12 = 0;

    if (args) {
        char tmplname[1000];
        char png_cookie[8];
        png_structp png_ptr;
        png_infop info_ptr;
        png_infop end_info;
        png_bytepp img_rows;
        int filelen, i;
        FILE* tmplfile = 0;
	    int result = sscanf(args, "%d:%d:%s",
               &vf->priv->x, &vf->priv->y, tmplname);

        if (result != 3) {
            mp_msg(MSGT_VFILTER, MSGL_ERR, DEBLENDLOGO_MODNAME ": Invalid arguments. Usage: deblendlogo=x:y:templatefile\n");
            return 0;
        }

        if (! (tmplfile = fopen(tmplname, "rw"))) {
            mp_msg(MSGT_VFILTER, MSGL_ERR, DEBLENDLOGO_MODNAME ": Failed to load template file %s\n", tmplname);
            return 0;
        }

        fread(png_cookie, 1, sizeof(png_cookie), tmplfile);
        if (png_sig_cmp(png_cookie, 0, sizeof(png_cookie))) {
            mp_msg(MSGT_VFILTER, MSGL_ERR, DEBLENDLOGO_MODNAME ": Given template file is not a PNG image\n");
            fclose(tmplfile);
            return 0;
        }

        fseek(tmplfile, 0, 0);

        if (! (png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, 0, 0, 0))) {
            mp_msg(MSGT_VFILTER, MSGL_ERR, DEBLENDLOGO_MODNAME ": Failed to allocate libpng read struct\n");
            fclose(tmplfile);
            return 0;
        }

        if (! (info_ptr = png_create_info_struct(png_ptr))) {
            mp_msg(MSGT_VFILTER, MSGL_ERR, DEBLENDLOGO_MODNAME ": Failed to allocate libpng info struct\n");
            png_destroy_read_struct(&png_ptr, 0, 0);
            fclose(tmplfile);
            return 0;
        }

        if (! (end_info = png_create_info_struct(png_ptr))) {
            mp_msg(MSGT_VFILTER, MSGL_ERR, DEBLENDLOGO_MODNAME ": Failed to allocate libpng end info struct\n");
            png_destroy_read_struct(&png_ptr, &info_ptr, 0);
            fclose(tmplfile);
            return 0;
        }

        #define DEBLEND_CLEANUP { png_destroy_read_struct(&png_ptr, &info_ptr, &end_info); \
                                       fclose(tmplfile); return 0; }

        /* FIXME: All of this PNG reading stuff should really be ready to catch a longjmp,
             but that only matters if the template image is corrupt, which should happen
             rarely. Right now it'll just cause mplayer to exit with an error message. */
        png_init_io(png_ptr, tmplfile);
        png_read_info(png_ptr, info_ptr);

        if (png_get_color_type(png_ptr,info_ptr) != PNG_COLOR_TYPE_RGB_ALPHA) {
            mp_msg(MSGT_VFILTER, MSGL_ERR, DEBLENDLOGO_MODNAME ": Template image must be an RGB image with an alpha channel\n");
            DEBLEND_CLEANUP;
            return 0;
        }

        if (png_get_bit_depth(png_ptr,info_ptr) == 16) {
            png_set_strip_16(png_ptr);
        }

        vf->priv->tw = png_get_image_width(png_ptr, info_ptr);
        vf->priv->th = png_get_image_height(png_ptr, info_ptr);

        if (! (vf->priv->tmpl = (char*)malloc(4 * vf->priv->tw * vf->priv->th))) {
            mp_msg(MSGT_VFILTER, MSGL_ERR, DEBLENDLOGO_MODNAME ": Failed to allocate memory for template image\n");
            DEBLEND_CLEANUP;
            return 0;
        }

        if (! (img_rows = (png_bytepp)malloc(sizeof(png_bytep) * vf->priv->th))) {
            mp_msg(MSGT_VFILTER, MSGL_ERR, DEBLENDLOGO_MODNAME ": Failed to allocate memory\n");
            DEBLEND_CLEANUP;
            return 0;
        }

        for (i = 0; i < vf->priv->th; i++) {
            img_rows[i] = vf->priv->tmpl + (i * vf->priv->tw * 4);
        }

        png_read_image(png_ptr, img_rows);

        png_read_end(png_ptr, end_info);
        png_destroy_read_struct(&png_ptr, &info_ptr, &end_info);

    }
    else {
        mp_msg(MSGT_VFILTER, MSGL_ERR, DEBLENDLOGO_MODNAME ": Invalid arguments. Usage: deblendlogo=x:y:templatefile\n");
        return 0;
    }
    return 1;
}

vf_info_t vf_info_deblendlogo = {
    "remove alpha-blended station logos",
    DEBLENDLOGO_NAME,
    "Martin Atkins",
    "",
    open,
    NULL
};

#endif