/*
    PYGAME EXTENSION - pgext
    many parts of this project are taken from opensource projects:
        libSDL
        pygame

    Copyright (c) 2012-2013, Josef Vanzura <gindar@zamraky.cz>
    All rights reserved.

    Redistribution and use in source and binary forms, with or without
    modification, are permitted provided that the following conditions are met:

    1. Redistributions of source code must retain the above copyright notice, this
       list of conditions and the following disclaimer.
    2. Redistributions in binary form must reproduce the above copyright notice,
       this list of conditions and the following disclaimer in the documentation
       and/or other materials provided with the distribution.

    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
    ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
    WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
    DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
    ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
    (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
    LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
    ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
    SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

    The views and conclusions contained in the software and documentation are those
    of the authors and should not be interpreted as representing official policies,
    either expressed or implied, of the FreeBSD Project.
*/

#include "pgext.h"
#include "hsvlrgb.h"
#include "getset.h"

/*
    Set all opaque (alpha>0) pixels to specified color (R,G,B).
*/
static PyObject*
setColor( PyObject* self, PyObject* args )
{
    PyObject* py_surface = NULL;
    SDL_Surface *surface;
    int w,h,x,y = 0;
    char blend = 0;
    SDL_PixelFormat *format = NULL;
    PyObject *py_color = NULL;
    Uint8 color[4];
    Uint8 r=0,g=0,b=0,a=0;
    float sa=0, da=0, oa=0;
    Uint8 sr=0, sg=0, sb=0;

    PY_PARSEARGS_BEGIN
        "OO", &py_surface, &py_color
    PY_PARSEARGS_END

    getSurfaceAttrs(py_surface, &w, &h, &format, &surface);

    if(!RGBAFromColorObj (py_color, color)){
        Py_INCREF(Py_None);return Py_None;
    }

    if(color[3] == 0){
        Py_INCREF(Py_None);return Py_None;
    }

    if(color[3] != 255){
        blend = 1;
        sa = (float) (((float) color[3]) / 255.0f);
        sr = (Uint8) (((float) color[0]) * sa);
        sg = (Uint8) (((float) color[1]) * sa);
        sb = (Uint8) (((float) color[2]) * sa);
    }

    for( y = 0; y < h; y ++ ){
        for( x = 0; x < w; x ++ ){
            getPixel_RGBA(surface, x, y, &r, &g, &b, &a);
            if(a==0) continue;
            if(blend){
                da = (float) (((float) a) / 255.0f);
                da = da * (1 - sa);
                oa = sa + da;
                r = clampUint8((((float)r * da) + sr) / oa );
                g = clampUint8((((float)g * da) + sg) / oa );
                b = clampUint8((((float)b * da) + sb) / oa );
                setPixel_RGBA(surface, x, y, r, g, b, a );
            } else {
                setPixel_RGBA(surface, x, y, color[0], color[1], color[2], a );
            }
        }
    }

    Py_INCREF(Py_None);
    return Py_None;
}

/*
    Change alpha of surface
    method:
        0 alpha = alpha - shift
        1 alpha = alpha
        2 alpha = alpha * (shift/100)
*/
static PyObject*
setAlpha( PyObject* self, PyObject* args )
{
    PyObject* py_surface = NULL;
    SDL_Surface *surface;
    int w,h,x,y = 0;
    SDL_PixelFormat *format = NULL;
    int shift;
    int method = 0;
    double percent = 0;
    Uint8 r,g,b,a;

    PY_PARSEARGS_BEGIN
        "Oi|i", &py_surface, &shift, &method
    PY_PARSEARGS_END

    getSurfaceAttrs(py_surface, &w, &h, &format, &surface);
    if( method == 1 ){
        shift = clampUint8( shift );
    }
    if( method == 2 ){
        percent = ((double)(shift)) / 100.0f;
    }

    for( y = 0; y < h; y ++ ){
        for( x = 0; x < w; x ++ ){
            getPixel_RGBA(surface, x, y, &r, &g, &b, &a );
            if(method == 0){
                a = clampUint8(a + shift);
            }
            if(method == 1){
                a = shift;
            }
            if(method == 2){
                a = (Uint8) (a * percent);
            }
            setPixel_RGBA(surface, x, y, r, g, b, a );
        }
    }

    Py_INCREF(Py_None);return Py_None;

}

/*
    Greyscale
    method:
        0 - avg RGB
        1 - HSV (saturation=0)
        2 - HSL (saturation=0)
*/
static PyObject*
greyscale( PyObject* self, PyObject* args )
{
    PyObject* py_surface = NULL;
    SDL_PixelFormat *format = NULL;
    SDL_Surface *surface;
    int w,h = 0;
    Uint8 r,g,b,a;
    int x,y;
    double hue,sat,val;
    int v;
    int method = 0;

    PY_PARSEARGS_BEGIN
        "O|i", &py_surface, &method
    PY_PARSEARGS_END

    getSurfaceAttrs(py_surface, &w, &h, &format, &surface);

    if(method == 0){
        GETPIXELS_RGBA(surface);
            v = (int) ((float) (r + g + b) / 3.0f);
            setPixel_RGBA(surface, x, y, v, v, v, a);
        }}
    } else if(method == 1){
        GETPIXELS_RGBA(surface);
            RGB2HSV( r, g, b, &hue, &sat, &val );
            HSV2RGB( hue, 0, val, &r, &g, &b );
        SETPIXELS_RGBA(surface, a);
    } else if(method == 2){
        GETPIXELS_RGBA(surface);
            RGB2HSL( r, g, b, &hue, &sat, &val );
            HSL2RGB( hue, 0, val, &r, &g, &b );
        SETPIXELS_RGBA(surface, a);
    }
    Py_INCREF(Py_None);return Py_None;
}

/*
    Desaturate
*/
static PyObject*
desaturate( PyObject* self, PyObject* args )
{
    PyObject* py_surface = NULL;
    SDL_Surface *surface;
    int w,h,pxcnt = 0;
    SDL_PixelFormat *format = NULL;
    Uint8 r,g,b,a;
    int eq = 0;
    float percent = 0;

    PY_PARSEARGS_BEGIN
        "Of", &py_surface, &percent
    PY_PARSEARGS_END

    getSurfaceAttrs(py_surface, &w, &h, &format, &surface);
    pxcnt = w * h;
    if(percent >= 1.0){
        Py_INCREF(Py_None);return Py_None;
    }

    while( pxcnt-- ){
        getPixelN_RGBA(surface, pxcnt, &r, &g, &b, &a );
        eq = (int) ((float) (r + g + b) / 3.0f);
        r = (Uint8)(r + (((float) (eq - r)) * percent));
        g = (Uint8)(g + (((float) (eq - g)) * percent));
        b = (Uint8)(b + (((float) (eq - b)) * percent));
        setPixelN_RGBA(surface, pxcnt, r, g, b, a );
    }
    Py_INCREF(Py_None);return Py_None;
}



/*
    Invert surface colors
*/
static PyObject*
invert( PyObject* self, PyObject* args )
{
    PyObject* py_surface = NULL;
    SDL_Surface *surface;
    int w,h,pxcnt = 0;
    SDL_PixelFormat *format = NULL;
    Uint8 r,g,b,a;

    PY_PARSEARGS_BEGIN
        "O", &py_surface
    PY_PARSEARGS_END

    getSurfaceAttrs(py_surface, &w, &h, &format, &surface);
    pxcnt = w * h;

    while( pxcnt-- ){
        getPixelN_RGBA(surface, pxcnt, &r, &g, &b, &a);
        setPixelN_RGBA(surface, pxcnt, 255-r, 255-g, 255-b, a);
    }
    Py_INCREF(Py_None);return Py_None;
}


/*
    Adjust surface contrast
*/
static PyObject*
contrast( PyObject* self, PyObject* args )
{
    PyObject* py_surface = NULL;
    SDL_Surface *surface;
    int w,h,pxcnt = 0,cvalue = 0;
    float factor = 0;
    SDL_PixelFormat *format = NULL;
    Uint8 r,g,b,a;

    PY_PARSEARGS_BEGIN
        "Oi", &py_surface, &cvalue
    PY_PARSEARGS_END

    getSurfaceAttrs(py_surface, &w, &h, &format, &surface);
    pxcnt = w * h;

    factor = (259 * ((float)cvalue + 255)) / (255 * (259 - cvalue));

    while( pxcnt-- ){
        getPixelN_RGBA(surface, pxcnt, &r, &g, &b, &a);
        r = clampUint8(factor * (r - 128) + 128);
        g = clampUint8(factor * (g - 128) + 128);
        b = clampUint8(factor * (b - 128) + 128);
        setPixelN_RGBA(surface, pxcnt, r, g, b, a);
    }
    Py_INCREF(Py_None);return Py_None;
}

/*
    Adjust surface brightness
*/
static PyObject*
brightness( PyObject* self, PyObject* args )
{
    PyObject* py_surface = NULL;
    SDL_Surface *surface;
    int w,h,pxcnt = 0,bvalue = 0;
    SDL_PixelFormat *format = NULL;
    Uint8 r,g,b,a;

    PY_PARSEARGS_BEGIN
        "Oi", &py_surface, &bvalue
    PY_PARSEARGS_END

    getSurfaceAttrs(py_surface, &w, &h, &format, &surface);
    pxcnt = w * h;

    while( pxcnt-- ){
        getPixelN_RGBA(surface, pxcnt, &r, &g, &b, &a);
        r = clampUint8(r + bvalue);
        g = clampUint8(g + bvalue);
        b = clampUint8(b + bvalue);
        setPixelN_RGBA(surface, pxcnt, r, g, b, a);
    }
    Py_INCREF(Py_None);return Py_None;
}

/*
    Change hue of surface
    method: (only 0)
        0 - hue += shift
*/
static PyObject*
hue( PyObject* self, PyObject* args )
{
    PyObject* py_surface = NULL;
    SDL_Surface *surface;
    int w,h,pxcnt = 0;
    Uint32 pixel;
    SDL_PixelFormat *format = NULL;
    int shift = 0;
    int method = 0;
    Uint8 r,g,b,a;
    double hue,sat,val;

    PY_PARSEARGS_BEGIN
        "Oi|i", &py_surface, &shift, &method
    PY_PARSEARGS_END

    getSurfaceAttrs(py_surface, &w, &h, &format, &surface);
    pxcnt = w * h;

    while( pxcnt-- ){
        pixel = getPixelN(surface, pxcnt);
        SDL_GetRGBA( pixel, format, &r, &g, &b, &a );

        RGB2HSV( r,g,b, &hue, &sat, &val );

        if(method == 0){
            hue += (double)(shift);
        } else if(method == 1){
            hue = (double)(shift);
        }

        /* clamp hue 0-100 */
        if( hue < 0 ){
            hue = 360.0+hue;
        }
        if( hue > 360 ){
            hue = hue-360.0;
        }
        hue = clampValue( hue, 0.0, 360.0 );
        HSV2RGB( hue, sat, val, &r, &g, &b );
        pixel = SDL_MapRGBA( format, r, g, b, a );
        setPixelN(surface, pxcnt, pixel);
    }
    Py_INCREF(Py_None);return Py_None;
}

/*
    Change saturation of surface
    method:
        0 - sat += s_shift
        1 - sat = s_shift
        2 - sat =  sat * (s_shift/100.0)
*/
static PyObject*
saturation( PyObject* self, PyObject* args )
{
    PyObject* py_surface = NULL;
    SDL_Surface *surface;
    int w,h,pxcnt = 0;
    Uint32 pixel;
    SDL_PixelFormat *format = NULL;
    int shift = 0;
    int method = 0;
    Uint8 r,g,b,a;
    double hue,sat,val;
    double percent = 0;

    PY_PARSEARGS_BEGIN
        "Oi|i", &py_surface, &shift, &method
    PY_PARSEARGS_END

    getSurfaceAttrs(py_surface, &w, &h, &format, &surface);
    pxcnt = w * h;

    if( method == 2 ){
        percent = ((double)(shift))/100.0;
    }

    while( pxcnt-- ){
        pixel = getPixelN(surface, pxcnt);
        SDL_GetRGBA( pixel, format, &r, &g, &b, &a );

        RGB2HSV( r,g,b, &hue, &sat, &val );
        switch(method){
            case 0: sat += (double)(shift);break;
            case 1: sat = (double)(shift);break;
            case 2: sat = sat * percent;break;
        }

        /* clamp sat 0-100 */
        sat = clampValue( sat, 0.0, 100.0 );
        HSV2RGB( hue,sat,val, &r, &g, &b );
        pixel = SDL_MapRGBA( format, r, g, b, a );
        setPixelN(surface, pxcnt, pixel);
    }
    Py_INCREF(Py_None);return Py_None;
}

/*
    Change value of surface
    method:
        0 - val += v_shift
        1 - val = v_shift
        2 - val =  val * (v_shift/100.0)
*/
static PyObject*
value( PyObject* self, PyObject* args )
{
    PyObject* py_surface = NULL;
    SDL_Surface *surface;
    int w,h,pxcnt = 0;
    Uint32 pixel;
    SDL_PixelFormat *format = NULL;

    int shift = 0;
    int method = 0;
    Uint8 r,g,b,a;
    double hue,sat,val;
    double percent = 0;

    PY_PARSEARGS_BEGIN
        "Oi|i", &py_surface, &shift, &method
    PY_PARSEARGS_END

    getSurfaceAttrs(py_surface, &w, &h, &format, &surface);
    pxcnt = w * h;

    if( method == 2 ){
        percent = ((double)(shift))/100.0;
    }
    while( pxcnt-- ){
        pixel = getPixelN(surface, pxcnt);
        SDL_GetRGBA( pixel, format, &r, &g, &b, &a );

        RGB2HSV( r,g,b, &hue, &sat, &val );
        switch(method){
            case 0: val += (double)(shift);break;
            case 1: val = (double)(shift);break;
            case 2: val = val * percent;break;
        }

        /* clamp sat 0-100 */
        val = clampValue( val, 0, 100 );
        HSV2RGB( hue,sat,val, &r, &g, &b );
        pixel = SDL_MapRGBA( format, r, g, b, a );
        setPixelN(surface, pxcnt, pixel);
    }
    Py_INCREF(Py_None);return Py_None;
}

static PyObject*
lightness( PyObject* self, PyObject* args )
{
    PyObject* py_surface = NULL;
    SDL_Surface *surface;
    int w,h,pxcnt = 0;
    Uint32 pixel;
    SDL_PixelFormat *format = NULL;
    int shift = 0;
    int method = 0;
    Uint8 r,g,b,a;
    double hue,sat,val;
    double percent = 0;

    PY_PARSEARGS_BEGIN
        "Oi|i", &py_surface, &shift, &method
    PY_PARSEARGS_END

    getSurfaceAttrs(py_surface, &w, &h, &format, &surface);
    pxcnt = w * h;

    if( method == 2 ){
        percent = ((double)(shift))/100.0;
    }
    while( pxcnt-- ){
        pixel = getPixelN(surface, pxcnt);
        SDL_GetRGBA( pixel, format, &r, &g, &b, &a );

        RGB2HSL( r,g,b, &hue, &sat, &val );
        switch(method){
            case 0: val += (double)(shift);break;
            case 1: val = (double)(shift);break;
            case 2: val = val * percent;break;
        }

        /* clamp sat 0-100 */
        val = clampValue( val, 0.0, 100.0 );
        HSL2RGB( hue,sat,val, &r, &g, &b );
        pixel = SDL_MapRGBA( format, r, g, b, a );
        setPixelN(surface, pxcnt, pixel);
    }

    Py_INCREF(Py_None);return Py_None;
}

/* Color multiply */
static PyObject*
multiply( PyObject* self, PyObject* args )
{
    PyObject* py_surface = NULL;
    SDL_Surface *surface;
    int w,h,pxcnt = 0;
    SDL_PixelFormat *format = NULL;
    float shift = 1.0f;
    float rp = 1.0f, gp = 1.0f, bp = 1.0f;
    char rb = 1, gb = 1, bb = 1;
    Uint8 r,g,b,a;

    PY_PARSEARGS_BEGIN
        "Of|bbb", &py_surface, &shift, &rb, &gb, &bb
    PY_PARSEARGS_END

    getSurfaceAttrs(py_surface, &w, &h, &format, &surface);
    pxcnt = w * h;

    if(rb){rp = shift;}
    if(gb){gp = shift;}
    if(bb){bp = shift;}

    while( pxcnt-- ){
        getPixelN_RGBA(surface, pxcnt, &r, &g, &b, &a);

        r = clampUint8(r * rp);
        g = clampUint8(g * gp);
        b = clampUint8(b * bp);

        setPixelN_RGBA(surface, pxcnt, r, g, b, a);
    }

    Py_INCREF(Py_None);return Py_None;
}


/*
    Replace pixels (hue == check_hue) with new_hue.
*/
static PyObject*
colorize( PyObject* self, PyObject* args )
{
    PyObject* py_surface = NULL;
    SDL_Surface *surface;
    int w,h,x,y = 0;
    SDL_PixelFormat *format = NULL;
    int check_hue,new_hue;
    int s_shift = 0;
    int v_shift = 0;
    Uint8 r,g,b,a;
    double hue,sat,val;
    float sq;
    long count = 0;

    PY_PARSEARGS_BEGIN
        "Oii|ii",
        &py_surface, &check_hue, &new_hue, &v_shift, &s_shift
    PY_PARSEARGS_END

    getSurfaceAttrs(py_surface, &w, &h, &format, &surface);

    GETPIXELS_RGBA(surface);
        RGB2HSV( r,g,b, &hue, &sat, &val );
        if( (int)hue == check_hue ){
            hue = new_hue;
            if( s_shift != 0 || v_shift != 0 ){
                sq = (float) sat / 100.0f;
                if( s_shift != 0 ){
                    sat += ( ((double)(s_shift)) *sq);
                    sat = clampValue( sat, 0.0f, 100.0f );
                }
                if( v_shift != 0 ){
                    val += ( ((double)(v_shift)) *sq);
                    val = clampValue( val, 0.0f, 100.0f );
                }
            }
            count++;
        }
        HSV2RGB( hue,sat,val, &r, &g, &b );
    SETPIXELS_RGBA(surface, a);
    return Py_BuildValue("i", count );
}


/*
    Apply per-pixel alpha mask on surface
*/
static PyObject*
alphaMask( PyObject* self, PyObject* args )
{
    PyObject* py_surface = NULL;
    SDL_Surface *surface;
    SDL_Surface *target_surface;
    int w,h,x,y,mw,mh = 0;
    SDL_PixelFormat *format = NULL;
    PyObject* py_surface_mask = NULL;
    SDL_Surface *surface_mask;
    SDL_PixelFormat *format_mask = NULL;

    Uint8 r,g,b,a,tmp_a;
    double hue,sat,val;
    int method = 0;
    float alpha,alpha2;

    PY_PARSEARGS_BEGIN
        "OO|i",
        &py_surface, &py_surface_mask, &method
    PY_PARSEARGS_END

    getSurfaceAttrs(py_surface_mask, &mw, &mh, &format_mask, &surface_mask);
    getSurfaceAttrs(py_surface, &w, &h, &format, &surface);
    target_surface = createSurface( mw, mh );

    if( mw != w || mh != h ){
        return RAISE (PyExc_ValueError, "mask must be same size as surface");
    }

    if( method == 0 ){
        GETPIXELS_RGBA(surface_mask);
            RGB2HSV( r,g,b, &hue, &sat, &val );
            alpha = (((float)val) / 100.0f);
            /* set pixel */
            getPixel_RGBA(surface, x, y, &r, &g, &b, &a );
            tmp_a = (Uint8)( ((float)a) * alpha );
        SETPIXELS_RGBA(target_surface, tmp_a);
    } else if( method == 1 ){
        GETPIXELS_RGBA(surface_mask);
            tmp_a = a;
            getPixel_RGBA(surface, x, y, &r, &g, &b, &a );
        SETPIXELS_RGBA(target_surface, tmp_a);
    } else if( method == 2 ){
        GETPIXELS_RGBA(surface_mask);
            tmp_a = a;
            getPixel_RGBA(surface, x, y, &r, &g, &b, &a );
            /* alpha blend */
            alpha = ((float) a) / 255.0f;
            alpha2 = ((float) tmp_a) / 255.0f;
            tmp_a = (Uint8) (alpha * alpha2) * 255;
        SETPIXELS_RGBA(target_surface, tmp_a);
    } else if( method == 3 ){
        GETPIXELS_RGBA(surface_mask);
            tmp_a = a;/*mask alpha*/
            getPixel_RGBA(surface, x, y, &r, &g, &b, &a );
            if(tmp_a > 0){/* if mask alpha > 0 then set original alpha */
                tmp_a = a;
            }
        SETPIXELS_RGBA(target_surface, tmp_a);
    } else {
        GETPIXELS_RGBA(surface_mask);
        SETPIXELS_RGBA(target_surface, 255);
    }

    return PySurface_New( target_surface );
}


static PyMethodDef ColorMethods[] = {
    {"setColor", setColor, METH_VARARGS, "Set opaque pixels to color"},
    {"setAlpha", setAlpha, METH_VARARGS, "Set alpha channel of surface"},

    {"hue", hue, METH_VARARGS, "Hue"},
    {"saturation", saturation, METH_VARARGS, "Change saturation of surface"},
    {"value", value, METH_VARARGS, "Change value of surface"},
    {"lightness", lightness, METH_VARARGS, "Lightness"},
    {"multiply", multiply, METH_VARARGS, "Color multiply"},

    {"brightness", brightness, METH_VARARGS, "Surface brightness"},
    {"contrast", contrast, METH_VARARGS, "Surface contrast"},
    {"desaturate", desaturate, METH_VARARGS, "Desaturate surface"},
    {"invert", invert, METH_VARARGS, "Invert colors of surface"},
    {"greyscale", greyscale, METH_VARARGS, "Create greyscale surface"},
    {"colorize", colorize, METH_VARARGS, "Change hue of surface"},

    {"alphaMask", alphaMask, METH_VARARGS, ""},
    {NULL, NULL, NULL, NULL}
};

PyMODINIT_FUNC initcolor(void)
{
    Py_InitModule("color", ColorMethods);
    import_pygame_surface();
    import_pygame_color();
}
