/* GStreamer

 * Copyright (C) <1999> Erik Walthinsen <omega@cse.ogi.edu>

 * Copyright (C) <2007> Wim Taymans <wim.taymans@collabora.co.uk>

 * Copyright (C) <2007> Edward Hervey <edward.hervey@collabora.co.uk>

 * Copyright (C) <2007> Jan Schmidt <thaytan@noraisin.net>

 *

 * This library is free software; you can redistribute it and/or

 * modify it under the terms of the GNU Library General Public

 * License as published by the Free Software Foundation; either

 * version 2 of the License, or (at your option) any later version.

 *

 * This library 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

 * Library General Public License for more details.

 *

 * You should have received a copy of the GNU Library General Public

 * License along with this library; if not, write to the

 * Free Software Foundation, Inc., 59 Temple Place - Suite 330,

 * Boston, MA 02111-1307, USA.

 */

#ifdef HAVE_CONFIG_H

#include "config.h"

#endif

#include <gst/gst.h>

#include <gst/base/gstbasetransform.h>

#include <gst/video/video.h>

#include <gst/controller/gstcontroller.h>

#include <stdlib.h>

#include <string.h>

#include <math.h>

#ifndef M_PI

#define M_PI  3.14159265358979323846

#endif

#define GST_TYPE_ALPHA \

  (gst_alpha_get_type())

#define GST_ALPHA(obj) \

  (G_TYPE_CHECK_INSTANCE_CAST((obj),GST_TYPE_ALPHA,GstAlpha))

#define GST_ALPHA_CLASS(klass) \

  (G_TYPE_CHECK_CLASS_CAST((klass),GST_TYPE_ALPHA,GstAlphaClass))

#define GST_IS_ALPHA(obj) \

  (G_TYPE_CHECK_INSTANCE_TYPE((obj),GST_TYPE_ALPHA))

#define GST_IS_ALPHA_CLASS(klass) \

  (G_TYPE_CHECK_CLASS_TYPE((klass),GST_TYPE_ALPHA))

typedef struct _GstAlpha GstAlpha;

typedef struct _GstAlphaClass GstAlphaClass;

typedef enum

{

  ALPHA_METHOD_SET,

  ALPHA_METHOD_GREEN,

  ALPHA_METHOD_BLUE,

  ALPHA_METHOD_CUSTOM,

}

GstAlphaMethod;

GST_DEBUG_CATEGORY_STATIC (gst_alpha_debug);

#define GST_CAT_DEFAULT gst_alpha_debug

struct _GstAlpha

{

  GstBaseTransform parent;

  /* caps */

  GstVideoFormat format;

  gint width, height;

  gboolean ayuv;

  gdouble alpha;

  guint target_r;

  guint target_g;

  guint target_b;

  GstAlphaMethod method;

  gfloat angle;

  gfloat noise_level;

  guint black_sensitivity;

  guint white_sensitivity;

  gfloat y;                     /* chroma color */

  gint8 cb, cr;

  gint8 kg;

  gfloat accept_angle_cos;

  gfloat accept_angle_sin;

  guint8 accept_angle_tg;

  guint8 accept_angle_ctg;

  guint8 one_over_kc;

  guint8 kfgy_scale;

};

struct _GstAlphaClass

{

  GstBaseTransformClass parent_class;

};

/* elementfactory information */

static const GstElementDetails gst_alpha_details =

GST_ELEMENT_DETAILS ("Alpha filter",

    "Filter/Effect/Video",

    "Adds an alpha channel to video - uniform or via chroma-keying",

    "Wim Taymans <wim@fluendo.com>\n"

    "Edward Hervey <edward.hervey@collabora.co.uk>\n"

    "Jan Schmidt <thaytan@noraisin.net>");

/* Alpha signals and args */

enum

{

  /* FILL ME */

  LAST_SIGNAL

};

#define DEFAULT_METHOD ALPHA_METHOD_SET

#define DEFAULT_ALPHA 1.0

#define DEFAULT_TARGET_R 0

#define DEFAULT_TARGET_G 255

#define DEFAULT_TARGET_B 0

#define DEFAULT_ANGLE 20.0

#define DEFAULT_NOISE_LEVEL 2.0

#define DEFAULT_BLACK_SENSITIVITY 100

#define DEFAULT_WHITE_SENSITIVITY 100

enum

{

  PROP_0,

  PROP_METHOD,

  PROP_ALPHA,

  PROP_TARGET_R,

  PROP_TARGET_G,

  PROP_TARGET_B,

  PROP_ANGLE,

  PROP_NOISE_LEVEL,

  PROP_BLACK_SENSITIVITY,

  PROP_WHITE_SENSITIVITY,

  PROP_LAST

};

static GstStaticPadTemplate gst_alpha_src_template =

GST_STATIC_PAD_TEMPLATE ("src",

    GST_PAD_SRC,

    GST_PAD_ALWAYS,

    GST_STATIC_CAPS (GST_VIDEO_CAPS_YUV ("AYUV"))

    );

static GstStaticPadTemplate gst_alpha_sink_template =

    GST_STATIC_PAD_TEMPLATE ("sink",

    GST_PAD_SINK,

    GST_PAD_ALWAYS,

    GST_STATIC_CAPS (GST_VIDEO_CAPS_YUV ("AYUV") ";" GST_VIDEO_CAPS_YUV ("I420")

    )

    );

static gboolean gst_alpha_start (GstBaseTransform * trans);

static gboolean gst_alpha_get_unit_size (GstBaseTransform * btrans,

    GstCaps * caps, guint * size);

static GstCaps *gst_alpha_transform_caps (GstBaseTransform * btrans,

    GstPadDirection direction, GstCaps * caps);

static gboolean gst_alpha_set_caps (GstBaseTransform * btrans,

    GstCaps * incaps, GstCaps * outcaps);

static GstFlowReturn gst_alpha_transform (GstBaseTransform * btrans,

    GstBuffer * in, GstBuffer * out);

static void gst_alpha_init_params (GstAlpha * alpha);

static void gst_alpha_set_property (GObject * object, guint prop_id,

    const GValue * value, GParamSpec * pspec);

static void gst_alpha_get_property (GObject * object, guint prop_id,

    GValue * value, GParamSpec * pspec);

GST_BOILERPLATE (GstAlpha, gst_alpha, GstBaseTransform,

    GST_TYPE_BASE_TRANSFORM);

#define GST_TYPE_ALPHA_METHOD (gst_alpha_method_get_type())

static GType

gst_alpha_method_get_type (void)

{

  static GType alpha_method_type = 0;

  static const GEnumValue alpha_method[] = {

    {ALPHA_METHOD_SET, "Set/adjust alpha channel", "set"},

    {ALPHA_METHOD_GREEN, "Chroma Key green", "green"},

    {ALPHA_METHOD_BLUE, "Chroma Key blue", "blue"},

    {ALPHA_METHOD_CUSTOM, "Chroma Key on target_r/g/b", "custom"},

    {0, NULL, NULL},

  };

  if (!alpha_method_type) {

    alpha_method_type = g_enum_register_static ("GstAlphaMethod", alpha_method);

  }

  return alpha_method_type;

}

static void

gst_alpha_base_init (gpointer g_class)

{

  GstElementClass *element_class = GST_ELEMENT_CLASS (g_class);

  gst_element_class_set_details (element_class, &gst_alpha_details);

  gst_element_class_add_pad_template (element_class,

      gst_static_pad_template_get (&gst_alpha_sink_template));

  gst_element_class_add_pad_template (element_class,

      gst_static_pad_template_get (&gst_alpha_src_template));

  GST_DEBUG_CATEGORY_INIT (gst_alpha_debug, "alpha", 0,

      "alpha - Element for adding alpha channel to streams");

}

static void

gst_alpha_class_init (GstAlphaClass * klass)

{

  GObjectClass *gobject_class;

  GstBaseTransformClass *btrans_class;

  gobject_class = (GObjectClass *) klass;

  btrans_class = (GstBaseTransformClass *) klass;

  gobject_class->set_property = gst_alpha_set_property;

  gobject_class->get_property = gst_alpha_get_property;

  g_object_class_install_property (G_OBJECT_CLASS (klass), PROP_METHOD,

      g_param_spec_enum ("method", "Method",

          "How the alpha channels should be created", GST_TYPE_ALPHA_METHOD,

          DEFAULT_METHOD, (GParamFlags) G_PARAM_READWRITE));

  g_object_class_install_property (G_OBJECT_CLASS (klass), PROP_ALPHA,

      g_param_spec_double ("alpha", "Alpha", "The value for the alpha channel",

          0.0, 1.0, DEFAULT_ALPHA,

          (GParamFlags) G_PARAM_READWRITE | GST_PARAM_CONTROLLABLE));

  g_object_class_install_property (G_OBJECT_CLASS (klass), PROP_TARGET_R,

      g_param_spec_uint ("target_r", "Target Red", "The Red target", 0, 255,

          DEFAULT_TARGET_R,

          (GParamFlags) G_PARAM_READWRITE | GST_PARAM_CONTROLLABLE));

  g_object_class_install_property (G_OBJECT_CLASS (klass), PROP_TARGET_G,

      g_param_spec_uint ("target_g", "Target Green", "The Green target", 0, 255,

          DEFAULT_TARGET_G,

          (GParamFlags) G_PARAM_READWRITE | GST_PARAM_CONTROLLABLE));

  g_object_class_install_property (G_OBJECT_CLASS (klass), PROP_TARGET_B,

      g_param_spec_uint ("target_b", "Target Blue", "The Blue target", 0, 255,

          DEFAULT_TARGET_B,

          (GParamFlags) G_PARAM_READWRITE | GST_PARAM_CONTROLLABLE));

  g_object_class_install_property (G_OBJECT_CLASS (klass), PROP_ANGLE,

      g_param_spec_float ("angle", "Angle", "Size of the colorcube to change",

          0.0, 90.0, DEFAULT_ANGLE,

          (GParamFlags) G_PARAM_READWRITE | GST_PARAM_CONTROLLABLE));

  g_object_class_install_property (G_OBJECT_CLASS (klass), PROP_NOISE_LEVEL,

      g_param_spec_float ("noise_level", "Noise Level", "Size of noise radius",

          0.0, 64.0, DEFAULT_NOISE_LEVEL,

          (GParamFlags) G_PARAM_READWRITE | GST_PARAM_CONTROLLABLE));

  g_object_class_install_property (G_OBJECT_CLASS (klass),

      PROP_BLACK_SENSITIVITY, g_param_spec_uint ("black-sensitivity",

          "Black Sensitivity", "Sensitivity to dark colors", 0, 128,

          DEFAULT_BLACK_SENSITIVITY,

          (GParamFlags) G_PARAM_READWRITE | GST_PARAM_CONTROLLABLE));

  g_object_class_install_property (G_OBJECT_CLASS (klass),

      PROP_WHITE_SENSITIVITY, g_param_spec_uint ("white-sensitivity",

          "Sensitivity", "Sensitivity to bright colors", 0, 128,

          DEFAULT_WHITE_SENSITIVITY,

          (GParamFlags) G_PARAM_READWRITE | GST_PARAM_CONTROLLABLE));

  btrans_class->start = GST_DEBUG_FUNCPTR (gst_alpha_start);

  btrans_class->transform = GST_DEBUG_FUNCPTR (gst_alpha_transform);

  btrans_class->get_unit_size = GST_DEBUG_FUNCPTR (gst_alpha_get_unit_size);

  btrans_class->transform_caps = GST_DEBUG_FUNCPTR (gst_alpha_transform_caps);

  btrans_class->set_caps = GST_DEBUG_FUNCPTR (gst_alpha_set_caps);

}

static void

gst_alpha_init (GstAlpha * alpha, GstAlphaClass * klass)

{

  alpha->alpha = DEFAULT_ALPHA;

  alpha->method = DEFAULT_METHOD;

  alpha->target_r = DEFAULT_TARGET_R;

  alpha->target_g = DEFAULT_TARGET_G;

  alpha->target_b = DEFAULT_TARGET_B;

  alpha->angle = DEFAULT_ANGLE;

  alpha->noise_level = DEFAULT_NOISE_LEVEL;

  alpha->black_sensitivity = DEFAULT_BLACK_SENSITIVITY;

  alpha->white_sensitivity = DEFAULT_WHITE_SENSITIVITY;

}

/* do we need this function? */

static void

gst_alpha_set_property (GObject * object, guint prop_id,

    const GValue * value, GParamSpec * pspec)

{

  GstAlpha *alpha;

  g_return_if_fail (GST_IS_ALPHA (object));

  alpha = GST_ALPHA (object);

  switch (prop_id) {

    case PROP_METHOD:

      alpha->method = g_value_get_enum (value);

      switch (alpha->method) {

        case ALPHA_METHOD_GREEN:

          alpha->target_r = 0;

          alpha->target_g = 255;

          alpha->target_b = 0;

          break;

        case ALPHA_METHOD_BLUE:

          alpha->target_r = 0;

          alpha->target_g = 0;

          alpha->target_b = 255;

          break;

        default:

          break;

      }

      gst_alpha_init_params (alpha);

      break;

    case PROP_ALPHA:

      alpha->alpha = g_value_get_double (value);

      break;

    case PROP_TARGET_R:

      alpha->target_r = g_value_get_uint (value);

      gst_alpha_init_params (alpha);

      break;

    case PROP_TARGET_G:

      alpha->target_g = g_value_get_uint (value);

      gst_alpha_init_params (alpha);

      break;

    case PROP_TARGET_B:

      alpha->target_b = g_value_get_uint (value);

      gst_alpha_init_params (alpha);

      break;

    case PROP_ANGLE:

      alpha->angle = g_value_get_float (value);

      gst_alpha_init_params (alpha);

      break;

    case PROP_NOISE_LEVEL:

      alpha->noise_level = g_value_get_float (value);

      gst_alpha_init_params (alpha);

      break;

    case PROP_BLACK_SENSITIVITY:

      alpha->black_sensitivity = g_value_get_uint (value);

      break;

    case PROP_WHITE_SENSITIVITY:

      alpha->white_sensitivity = g_value_get_uint (value);

      break;

    default:

      G_OBJECT_WARN_INVALID_PROPERTY_ID (object, prop_id, pspec);

      break;

  }

}

static void

gst_alpha_get_property (GObject * object, guint prop_id, GValue * value,

    GParamSpec * pspec)

{

  GstAlpha *alpha;

  g_return_if_fail (GST_IS_ALPHA (object));

  alpha = GST_ALPHA (object);

  switch (prop_id) {

    case PROP_METHOD:

      g_value_set_enum (value, alpha->method);

      break;

    case PROP_ALPHA:

      g_value_set_double (value, alpha->alpha);

      break;

    case PROP_TARGET_R:

      g_value_set_uint (value, alpha->target_r);

      break;

    case PROP_TARGET_G:

      g_value_set_uint (value, alpha->target_g);

      break;

    case PROP_TARGET_B:

      g_value_set_uint (value, alpha->target_b);

      break;

    case PROP_ANGLE:

      g_value_set_float (value, alpha->angle);

      break;

    case PROP_NOISE_LEVEL:

      g_value_set_float (value, alpha->noise_level);

      break;

    case PROP_BLACK_SENSITIVITY:

      g_value_set_uint (value, alpha->black_sensitivity);

      break;

    case PROP_WHITE_SENSITIVITY:

      g_value_set_uint (value, alpha->white_sensitivity);

      break;

    default:

      G_OBJECT_WARN_INVALID_PROPERTY_ID (object, prop_id, pspec);

      break;

  }

}

static gboolean

gst_alpha_get_unit_size (GstBaseTransform * btrans,

    GstCaps * caps, guint * size)

{

  GstVideoFormat format;

  gint width, height;

  if (!gst_video_format_parse_caps (caps, &format, &width, &height))

    return FALSE;

  *size = gst_video_format_get_size (format, width, height);

  GST_DEBUG_OBJECT (btrans, "unit size = %d for format %d w %d height %d",

      *size, format, width, height);

  return TRUE;

}

static GstCaps *

gst_alpha_transform_caps (GstBaseTransform * btrans,

    GstPadDirection direction, GstCaps * caps)

{

  GstCaps *ret;

  GstStructure *structure;

  gint i;

  ret = gst_caps_copy (caps);

  /* When going from the SINK pad to the src, we just need to make sure the

   * format is AYUV */

  if (direction == GST_PAD_SINK) {

    for (i = 0; i < gst_caps_get_size (ret); i++) {

      structure = gst_caps_get_structure (ret, i);

      gst_structure_set (structure, "format",

          GST_TYPE_FOURCC, GST_MAKE_FOURCC ('A', 'Y', 'U', 'V'), NULL);

    }

  } else {

    GstCaps *ayuv_caps;

    /* In the other direction, prepend a copy of the caps with format AYUV, 

     * and set the first to I420 */

    ayuv_caps = gst_caps_copy (ret);

    for (i = 0; i < gst_caps_get_size (ret); i++) {

      structure = gst_caps_get_structure (ret, i);

      gst_structure_set (structure, "format",

          GST_TYPE_FOURCC, GST_MAKE_FOURCC ('I', '4', '2', '0'), NULL);

    }

    gst_caps_append (ret, ayuv_caps);

  }

  gst_caps_do_simplify (ret);

  return ret;

}

static gboolean

gst_alpha_set_caps (GstBaseTransform * btrans,

    GstCaps * incaps, GstCaps * outcaps)

{

  GstAlpha *alpha = GST_ALPHA (btrans);

  if (!gst_video_format_parse_caps (incaps, &alpha->format,

          &alpha->width, &alpha->height))

    return FALSE;

  if (alpha->format == GST_VIDEO_FORMAT_AYUV)

    alpha->ayuv = TRUE;

  else

    alpha->ayuv = FALSE;

  return TRUE;

}

static void

gst_alpha_set_ayuv (guint8 * src, guint8 * dest, gint width, gint height,

    gdouble alpha)

{

  gint b_alpha = (gint) (alpha * 255);

  gint y, x;

  for (y = 0; y < height; y++) {

    for (x = 0; x < width; x++) {

      *dest++ = (*src++ * b_alpha) >> 8;

      *dest++ = *src++;

      *dest++ = *src++;

      *dest++ = *src++;

    }

  }

}

static void

gst_alpha_set_i420 (guint8 * src, guint8 * dest, gint width, gint height,

    gdouble alpha)

{

  gint b_alpha = (gint) (alpha * 255);

  guint8 *srcY;

  guint8 *srcU;

  guint8 *srcV;

  gint i, j;

  gint src_wrap, src_uv_wrap;

  gint y_stride, uv_stride;

  gboolean odd_width;

  y_stride = gst_video_format_get_row_stride (GST_VIDEO_FORMAT_I420, 0, width);

  uv_stride = gst_video_format_get_row_stride (GST_VIDEO_FORMAT_I420, 1, width);

  src_wrap = y_stride - width;

  src_uv_wrap = uv_stride - (width / 2);

  srcY = src;

  srcU = src + gst_video_format_get_component_offset (GST_VIDEO_FORMAT_I420,

      1, width, height);

  srcV = src + gst_video_format_get_component_offset (GST_VIDEO_FORMAT_I420,

      2, width, height);

  odd_width = (width % 2 != 0);

  for (i = 0; i < height; i++) {

    for (j = 0; j < width / 2; j++) {

      *dest++ = b_alpha;

      *dest++ = *srcY++;

      *dest++ = *srcU;

      *dest++ = *srcV;

      *dest++ = b_alpha;

      *dest++ = *srcY++;

      *dest++ = *srcU++;

      *dest++ = *srcV++;

    }

    /* Might have one odd column left to do */

    if (odd_width) {

      *dest++ = b_alpha;

      *dest++ = *srcY++;

      *dest++ = *srcU;

      *dest++ = *srcV;