/*
 * 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 "mp_image.h"
#include "../mp_msg.h"
#include "vf.h"
#include "img_format.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;
};

/* This filter is very picky about what formats it accepts, because
  it needs to be able to get at RGBA data to work on it with the RGBA template.
  In future, perhaps it should also be extended to work on YUV formats */
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:
            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");
    }
    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;

    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: {
            mp_msg(MSGT_VFILTER,MSGL_ERR,DEBLENDLOGO_MODNAME ": Unsupported image format.\n");
            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 += 1) {
        for (y = 0; y < priv->th; y += 1) {
            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) continue;
            if (alpha == 255) 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;
        }
    }

    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 != 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;

    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