/*

 * Copyright (c) Radzivon Bartoshyk, 10/2024. 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.

 *

 * 3.  Neither the name of the copyright holder nor the names of its

 * contributors may be used to endorse or promote products derived from

 * this software without specific prior written permission.

 *

 * 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 HOLDER 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.

 */

#[inline]

fn div_by_255(v: u16) -> u8 {

    ((((v + 0x80) >> 8) + v + 0x80) >> 8).min(255) as u8

}

const fn make_unpremultiplication_table() -> [u8; 65536] {

    let mut alpha = 0usize;

    let mut buf = [0u8; 65536];

    while alpha < 256 {

        let mut pixel = 0usize;

        while pixel < 256 {

            if alpha == 0 {

                buf[alpha * 255 + pixel] = 0;

            } else {

                let value = (pixel * 255 + alpha / 2) / alpha;

                buf[alpha * 255 + pixel] = if value > 255 { 255 } else { value as u8 };

            }

            pixel += 1;

        }

        alpha += 1;

    }

    buf

}

pub(crate) static UNPREMULTIPLICATION_TABLE: [u8; 65536] = make_unpremultiplication_table();

/// Associate alpha in place

///

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `in_place`: Slice to where premultiply

///

pub fn premultiply_rgba8(in_place: &mut [u8]) {

    // Almost all loops are not auto-vectorised without doing anything dirty.

    // So everywhere is just added something beautiful.

    for chunk in in_place.chunks_exact_mut(4) {

        let a = chunk[3] as u16;

        chunk[0] = div_by_255(chunk[0] as u16 * a);

        chunk[1] = div_by_255(chunk[1] as u16 * a);

        chunk[2] = div_by_255(chunk[2] as u16 * a);

        chunk[3] = div_by_255(255 * a);

    }

}

/// Associate alpha to new slice

///

/// Faster if you need to do a copy first.

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `source`: Source slice with RGBA data

///

pub fn premultiplied_rgba8(source: &[u8]) -> Vec<u8> {

    let mut target = vec![0u8; source.len()];

    // Almost all loops are not auto-vectorised without doing anything dirty.

    // So everywhere is just added something beautiful.

    for (dst, src) in target.chunks_exact_mut(4).zip(source.chunks_exact(4)) {

        let a = src[3] as u16;

        dst[0] = div_by_255(src[0] as u16 * a);

        dst[1] = div_by_255(src[1] as u16 * a);

        dst[2] = div_by_255(src[2] as u16 * a);

        dst[3] = div_by_255(255 * a);

    }

    target

}

/// Un premultiply alpha in place

///

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `in_place`: Slice to work on

///

///

pub fn unpremultiply_rgba8(in_place: &mut [u8]) {

    for chunk in in_place.chunks_exact_mut(4) {

        let a = chunk[3];

        let z = a as u16 * 255;

        chunk[0] = UNPREMULTIPLICATION_TABLE[(z + chunk[0] as u16) as usize];

        chunk[1] = UNPREMULTIPLICATION_TABLE[(z + chunk[1] as u16) as usize];

        chunk[2] = UNPREMULTIPLICATION_TABLE[(z + chunk[2] as u16) as usize];

    }

}

/// Associate alpha in place

///

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `in_place`: Slice to where premultiply

///

pub fn premultiply_la8(in_place: &mut [u8]) {

    // Almost all loops are not auto-vectorized without doing anything dirty.

    // So everywhere is just added something beautiful.

    for chunk in in_place.chunks_exact_mut(2) {

        let a = chunk[1] as u16;

        chunk[0] = div_by_255(chunk[0] as u16 * a);

        chunk[1] = div_by_255(255 * a);

    }

}

/// Associate alpha to a new destination

///

/// Faster if you need to do a copy first.

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `source`: Source slice with LA data

///

pub fn premultiplied_la8(source: &[u8]) -> Vec<u8> {

    let mut target = vec![0u8; source.len()];

    // Almost all loops are not auto-vectorised without doing anything dirty.

    // So everywhere is just added something beautiful.

    for (dst, src) in target.chunks_exact_mut(2).zip(source.chunks_exact(2)) {

        let a = src[1] as u16;

        dst[0] = div_by_255(src[0] as u16 * a);

        dst[1] = div_by_255(255 * a);

    }

    target

}

/// Un premultiply alpha in place

///

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `in_place`: Slice to work on

///

pub fn unpremultiply_la8(in_place: &mut [u8]) {

    // Almost all loops are not auto-vectorised without doing anything dirty.

    // So everywhere is just added something beautiful.

    for chunk in in_place.chunks_exact_mut(2) {

        let a = chunk[1];

        let z = a as u16 * 255;

        chunk[0] = UNPREMULTIPLICATION_TABLE[(z + chunk[0] as u16) as usize];

    }

}

/// Computes `round(v / (2**n - 1))`. The result is expected to fit in u16.

#[inline(always)]

fn div_by_2pn_m1(v: u32, n: u32) -> u16 {

    debug_assert!(n > 0 && n <= 16);

    let round = 1 << (n - 1);

    let v = v + round;

    (((v >> n) + v) >> n) as u16

}

#[inline]

fn div_by_1023(v: u32) -> u16 {

    div_by_2pn_m1(v, 10)

}

#[inline]

fn div_by_4095(v: u32) -> u16 {

    div_by_2pn_m1(v, 12)

}

#[inline]

fn div_by_65535(v: u32) -> u16 {

    div_by_2pn_m1(v, 16)

}

/// Associate alpha in place

///

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `in_place`: Slice to where premultiply

/// * `bit_depth`: Bit-depth of the image

///

pub fn premultiply_rgba16(in_place: &mut [u16], bit_depth: u32) {

    // Almost all loops are not auto-vectorised without doing anything dirty.

    // So everywhere is just added something beautiful.

    assert!(bit_depth > 0 && bit_depth <= 16);

    let max_colors = (1 << bit_depth) - 1;

    if bit_depth == 10 {

        for chunk in in_place.chunks_exact_mut(4) {

            let a = chunk[3] as u32;

            chunk[0] = div_by_1023(chunk[0] as u32 * a);

            chunk[1] = div_by_1023(chunk[1] as u32 * a);

            chunk[2] = div_by_1023(chunk[2] as u32 * a);

            chunk[3] = div_by_1023(1023 * a);

        }

    } else if bit_depth == 12 {

        for chunk in in_place.chunks_exact_mut(4) {

            let a = chunk[3] as u32;

            chunk[0] = div_by_4095(chunk[0] as u32 * a);

            chunk[1] = div_by_4095(chunk[1] as u32 * a);

            chunk[2] = div_by_4095(chunk[2] as u32 * a);

            chunk[3] = div_by_4095(4095 * a);

        }

    } else if bit_depth == 16 {

        for chunk in in_place.chunks_exact_mut(4) {

            let a = chunk[3] as u32;

            chunk[0] = div_by_65535(chunk[0] as u32 * a);

            chunk[1] = div_by_65535(chunk[1] as u32 * a);

            chunk[2] = div_by_65535(chunk[2] as u32 * a);

            chunk[3] = div_by_65535(65535 * a);

        }

    } else {

        for chunk in in_place.chunks_exact_mut(4) {

            let a = chunk[3] as u32;

            chunk[0] = div_by_2pn_m1(chunk[0] as u32 * a, bit_depth);

            chunk[1] = div_by_2pn_m1(chunk[1] as u32 * a, bit_depth);

            chunk[2] = div_by_2pn_m1(chunk[2] as u32 * a, bit_depth);

            chunk[3] = div_by_2pn_m1(max_colors * a, bit_depth);

        }

    }

}

/// Associate alpha to a new destination

///

/// Faster, if you need to copy data first.

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `source`: Source slice with RGBA16 data

/// * `bit_depth`: Bit-depth of the image

///

pub fn premultiplied_rgba16(source: &[u16], bit_depth: u32) -> Vec<u16> {

    let mut target = vec![0u16; source.len()];

    // Almost all loops are not auto-vectorised without doing anything dirty.

    // So everywhere is just added something beautiful.

    assert!(bit_depth > 0 && bit_depth <= 16);

    let max_colors = (1 << bit_depth) - 1;

    if bit_depth == 10 {

        for (dst, src) in target.chunks_exact_mut(4).zip(source.chunks_exact(4)) {

            let a = src[3] as u32;

            dst[0] = div_by_1023(src[0] as u32 * a);

            dst[1] = div_by_1023(src[1] as u32 * a);

            dst[2] = div_by_1023(src[2] as u32 * a);

            dst[3] = div_by_1023(1023 * a);

        }

    } else if bit_depth == 12 {

        for (dst, src) in target.chunks_exact_mut(4).zip(source.chunks_exact(4)) {

            let a = src[3] as u32;

            dst[0] = div_by_4095(src[0] as u32 * a);

            dst[1] = div_by_4095(src[1] as u32 * a);

            dst[2] = div_by_4095(src[2] as u32 * a);

            dst[3] = div_by_4095(4095 * a);

        }

    } else if bit_depth == 16 {

        for (dst, src) in target.chunks_exact_mut(4).zip(source.chunks_exact(4)) {

            let a = src[3] as u32;

            dst[0] = div_by_65535(src[0] as u32 * a);

            dst[1] = div_by_65535(src[1] as u32 * a);

            dst[2] = div_by_65535(src[2] as u32 * a);

            dst[3] = div_by_65535(65535 * a);

        }

    } else {

        for (dst, src) in target.chunks_exact_mut(4).zip(source.chunks_exact(4)) {

            let a = src[3] as u32;

            dst[0] = div_by_2pn_m1(src[0] as u32 * a, bit_depth);

            dst[1] = div_by_2pn_m1(src[1] as u32 * a, bit_depth);

            dst[2] = div_by_2pn_m1(src[2] as u32 * a, bit_depth);

            dst[3] = div_by_2pn_m1(max_colors * a, bit_depth);

        }

    }

    target

}

/// Associate alpha in place for up to 16 bit-depth image

///

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `in_place`: Slice to where premultiply

/// * `bit_depth`: Bit-depth of the image

///

pub fn premultiply_la16(in_place: &mut [u16], bit_depth: u32) {

    // Almost all loops are not auto-vectorised without doing anything dirty.

    // So everywhere is just added something beautiful.

    assert!(bit_depth > 0 && bit_depth <= 16);

    let max_colors = (1 << bit_depth) - 1;

    if bit_depth == 10 {

        for chunk in in_place.chunks_exact_mut(2) {

            let a = chunk[1] as u32;

            chunk[0] = div_by_1023(chunk[0] as u32 * a);

            chunk[1] = div_by_1023(1023 * a);

        }

    } else if bit_depth == 12 {

        for chunk in in_place.chunks_exact_mut(2) {

            let a = chunk[1] as u32;

            chunk[0] = div_by_4095(chunk[0] as u32 * a);

            chunk[1] = div_by_4095(4095 * a);

        }

    } else if bit_depth == 16 {

        for chunk in in_place.chunks_exact_mut(2) {

            let a = chunk[1] as u32;

            chunk[0] = div_by_65535(chunk[0] as u32 * a);

            chunk[1] = div_by_65535(65535 * a);

        }

    } else {

        for chunk in in_place.chunks_exact_mut(2) {

            let a = chunk[1] as u32;

            chunk[0] = div_by_2pn_m1(chunk[0] as u32 * a, bit_depth);

            chunk[1] = div_by_2pn_m1(max_colors * a, bit_depth);

        }

    }

}

/// Associate alpha for up to 16 bit-depth image to a new destination

///

/// Faster, if you need to copy data first.

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `source`: Slice with source LA16 data

/// * `bit_depth`: Bit-depth of the image

///

pub fn premultiplied_la16(source: &[u16], bit_depth: u32) -> Vec<u16> {

    let mut target = vec![0u16; source.len()];

    // Almost all loops are not auto-vectorised without doing anything dirty.

    // So everywhere is just added something beautiful.

    assert!(bit_depth > 0 && bit_depth <= 16);

    let max_colors = (1 << bit_depth) - 1;

    if bit_depth == 10 {

        for (dst, src) in target.chunks_exact_mut(2).zip(source.chunks_exact(2)) {

            let a = src[1] as u32;

            dst[0] = div_by_1023(src[0] as u32 * a);

            dst[1] = div_by_1023(1023 * a);

        }

    } else if bit_depth == 12 {

        for (dst, src) in target.chunks_exact_mut(2).zip(source.chunks_exact(2)) {

            let a = src[1] as u32;

            dst[0] = div_by_4095(src[0] as u32 * a);

            dst[1] = div_by_4095(4095 * a);

        }

    } else if bit_depth == 16 {

        for (dst, src) in target.chunks_exact_mut(2).zip(source.chunks_exact(2)) {

            let a = src[1] as u32;

            dst[0] = div_by_65535(src[0] as u32 * a);

            dst[1] = div_by_65535(65535 * a);

        }

    } else {

        for (dst, src) in target.chunks_exact_mut(2).zip(source.chunks_exact(2)) {

            let a = src[1] as u32;

            dst[0] = div_by_2pn_m1(src[0] as u32 * a, bit_depth);

            dst[1] = div_by_2pn_m1(max_colors * a, bit_depth);

        }

    }

    target

}

/// Un premultiply alpha in place for up to 16 bit-depth image

///

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `in_place`: Slice to work on

/// * `bit_depth`: Bit-depth of the image

///

///

pub fn unpremultiply_la16(in_place: &mut [u16], bit_depth: u32) {

    // Almost all loops are not auto-vectorised without doing anything dirty.

    // So everywhere is just added something beautiful.

    assert!(bit_depth > 0 && bit_depth <= 16);

    let max_colors = (1 << bit_depth) - 1;

    for chunk in in_place.chunks_exact_mut(2) {

        let a = chunk[1] as u32;

        if a != 0 {

            let a_recip = max_colors as f32 / a as f32;

            chunk[0] = (chunk[0] as f32 * a_recip) as u16;

        }

    }

}

/// Un premultiply alpha in place

///

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `in_place`: Slice to work on

/// * `bit_depth`: Bit-depth of the image

///

///

pub fn unpremultiply_rgba16(in_place: &mut [u16], bit_depth: u32) {

    // Almost all loops are not auto-vectorised without doing anything dirty.

    // So everywhere is just added something beautiful.

    assert!(bit_depth > 0 && bit_depth <= 16);

    let max_colors = (1 << bit_depth) - 1;

    for chunk in in_place.chunks_exact_mut(4) {

        let a = chunk[3] as u32;

        if a != 0 {

            let a_recip = max_colors as f32 / a as f32;

            chunk[0] = (chunk[0] as f32 * a_recip) as u16;

            chunk[1] = (chunk[1] as f32 * a_recip) as u16;

            chunk[2] = (chunk[2] as f32 * a_recip) as u16;

        }

    }

}

/// Associate alpha in place

///

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `in_place`: Slice to where premultiply

///

pub fn premultiply_rgba_f32(in_place: &mut [f32]) {

    // Almost all loops are not auto-vectorised without doing anything dirty.

    // So everywhere is just added something beautiful.

    for chunk in in_place.chunks_exact_mut(4) {

        let a = chunk[3];

        chunk[0] *= a;

        chunk[1] *= a;

        chunk[2] *= a;

        chunk[3] = a;

    }

}

/// Associate alpha in place

///

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `in_place`: Slice to where premultiply

///

pub fn premultiply_luma_alpha_f32(in_place: &mut [f32]) {

    // Almost all loops are not auto-vectorised without doing anything dirty.

    // So everywhere is just added something beautiful.

    for chunk in in_place.chunks_exact_mut(2) {

        let a = chunk[1];

        chunk[0] *= a;

        chunk[2] = a;

    }

}

/// Associate alpha to a new destination

///

/// Faster, if you need to do a copy first

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `source`: Source slice with luma alpha

///

pub fn premultiplied_luma_alpha_f32(source: &[f32]) -> Vec<f32> {

    let mut target = vec![0.; source.len()];

    // Almost all loops are not auto-vectorised without doing anything dirty.

    // So everywhere is just added something beautiful.

    for (dst, src) in target.chunks_exact_mut(2).zip(source.chunks_exact(2)) {

        let a = src[2];

        dst[0] = src[0] * a;

        dst[1] = a;

    }

    target

}

/// Associate alpha to a new destination

///

/// Faster, if you need to do a copy first

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `source`: Source rgba slice

///

pub fn premultiplied_rgba_f32(source: &[f32]) -> Vec<f32> {

    let mut target = vec![0.; source.len()];

    // Almost all loops are not auto-vectorised without doing anything dirty.

    // So everywhere is just added something beautiful.

    for (dst, src) in target.chunks_exact_mut(4).zip(source.chunks_exact(4)) {

        let a = src[3];

        dst[0] = src[0] * a;

        dst[1] = src[1] * a;

        dst[2] = src[2] * a;

        dst[3] = a;

    }

    target

}

/// Un-premultiply alpha in place

///

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `in_place`: Slice to work on

///

pub fn unpremultiply_rgba_f32(in_place: &mut [f32]) {

    for chunk in in_place.chunks_exact_mut(4) {

        let a = chunk[3];

        if a != 0. {

            let a_recip = 1. / a;

            chunk[0] *= a_recip;

            chunk[1] *= a_recip;

            chunk[2] *= a_recip;

            chunk[3] = a;

        }

    }

}

/// Un-premultiply alpha in place

///

/// Note, for scaling alpha must be *associated*

///

/// # Arguments

///

/// * `in_place`: Slice to work on

///

pub fn unpremultiply_luma_alpha_f32(in_place: &mut [f32]) {

    for chunk in in_place.chunks_exact_mut(2) {

        let a = chunk[1];

        if a != 0. {

            let a_recip = 1. / a;

            chunk[0] *= a_recip;

            chunk[1] = a;

        }

    }

}