/*

 * Copyright (C) 2024 Niklas Haas

 *

 * This file is part of FFmpeg.

 *

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

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

 * License as published by the Free Software Foundation; either

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

 *

 * FFmpeg 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

 * Lesser General Public License for more details.

 *

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

 * License along with FFmpeg; if not, write to the Free Software

 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

 */

#include <assert.h>

#include <string.h>

#include "libavutil/attributes.h"

#include "libavutil/avassert.h"

#include "libavutil/mem.h"

#include "cms.h"

#include "csputils.h"

#include "lut3d.h"

SwsLut3D *ff_sws_lut3d_alloc(void)

{

    SwsLut3D *lut3d = av_malloc(sizeof(*lut3d));

    if (!lut3d)

        return NULL;

    lut3d->dynamic = false;

    return lut3d;

}

void ff_sws_lut3d_free(SwsLut3D **plut3d)

{

    av_freep(plut3d);

}

bool ff_sws_lut3d_test_fmt(enum AVPixelFormat fmt, int output)

{

    return fmt == AV_PIX_FMT_RGBA64;

}

enum AVPixelFormat ff_sws_lut3d_pick_pixfmt(SwsFormat fmt, int output)

{

    return AV_PIX_FMT_RGBA64;

}

/**

 * v0 and v1 are 'black' and 'white'

 * v2 and v3 are closest RGB/CMY vertices

 * x >= y >= z are relative weights

 */

static av_always_inline

v3u16_t barycentric(int shift, int x, int y, int z,

                    v3u16_t v0, v3u16_t v1, v3u16_t v2, v3u16_t v3)

{

    const int a = (1 << shift) - x;

    const int b = x - y;

    const int c = y - z;

    const int d = z;

    av_assert2(x >= y);

    av_assert2(y >= z);

    return (v3u16_t) {

        (a * v0.x + b * v1.x + c * v2.x + d * v3.x) >> shift,

        (a * v0.y + b * v1.y + c * v2.y + d * v3.y) >> shift,

        (a * v0.z + b * v1.z + c * v2.z + d * v3.z) >> shift,

    };

}

static av_always_inline

v3u16_t tetrahedral(const SwsLut3D *lut3d, int Rx, int Gx, int Bx,

                    int Rf, int Gf, int Bf)

{

    const int shift = 16 - INPUT_LUT_BITS;

    const int Rn = FFMIN(Rx + 1, INPUT_LUT_SIZE - 1);

    const int Gn = FFMIN(Gx + 1, INPUT_LUT_SIZE - 1);

    const int Bn = FFMIN(Bx + 1, INPUT_LUT_SIZE - 1);

    const v3u16_t c000 = lut3d->input[Bx][Gx][Rx];

    const v3u16_t c111 = lut3d->input[Bn][Gn][Rn];

    if (Rf > Gf) {

        if (Gf > Bf) {

            const v3u16_t c100 = lut3d->input[Bx][Gx][Rn];

            const v3u16_t c110 = lut3d->input[Bx][Gn][Rn];

            return barycentric(shift, Rf, Gf, Bf, c000, c100, c110, c111);

        } else if (Rf > Bf) {

            const v3u16_t c100 = lut3d->input[Bx][Gx][Rn];

            const v3u16_t c101 = lut3d->input[Bn][Gx][Rn];

            return barycentric(shift, Rf, Bf, Gf, c000, c100, c101, c111);

        } else {

            const v3u16_t c001 = lut3d->input[Bn][Gx][Rx];

            const v3u16_t c101 = lut3d->input[Bn][Gx][Rn];

            return barycentric(shift, Bf, Rf, Gf, c000, c001, c101, c111);

        }

    } else {

        if (Bf > Gf) {

            const v3u16_t c001 = lut3d->input[Bn][Gx][Rx];

            const v3u16_t c011 = lut3d->input[Bn][Gn][Rx];

            return barycentric(shift, Bf, Gf, Rf, c000, c001, c011, c111);

        } else if (Bf > Rf) {

            const v3u16_t c010 = lut3d->input[Bx][Gn][Rx];

            const v3u16_t c011 = lut3d->input[Bn][Gn][Rx];

            return barycentric(shift, Gf, Bf, Rf, c000, c010, c011, c111);

        } else {

            const v3u16_t c010 = lut3d->input[Bx][Gn][Rx];

            const v3u16_t c110 = lut3d->input[Bx][Gn][Rn];

            return barycentric(shift, Gf, Rf, Bf, c000, c010, c110, c111);

        }

    }

}

static av_always_inline v3u16_t lookup_input16(const SwsLut3D *lut3d, v3u16_t rgb)

{

    const int shift = 16 - INPUT_LUT_BITS;

    const int Rx = rgb.x >> shift;

    const int Gx = rgb.y >> shift;

    const int Bx = rgb.z >> shift;

    const int Rf = rgb.x & ((1 << shift) - 1);

    const int Gf = rgb.y & ((1 << shift) - 1);

    const int Bf = rgb.z & ((1 << shift) - 1);

    return tetrahedral(lut3d, Rx, Gx, Bx, Rf, Gf, Bf);

}

/**

 * Note: These functions are scaled such that x == (1 << shift) corresponds to

 * a value of 1.0. This makes them suitable for use when interpolation LUT

 * entries with a fractional part that is just masked away from the index,

 * since a fractional coordinate of e.g. 0xFFFF corresponds to a mix weight of

 * just slightly *less* than 1.0.

 */

static av_always_inline v2u16_t lerp2u16(v2u16_t a, v2u16_t b, int x, int shift)

{

    const int xi = (1 << shift) - x;

    return (v2u16_t) {

        (a.x * xi + b.x * x) >> shift,

        (a.y * xi + b.y * x) >> shift,

    };

}

static av_always_inline v3u16_t lerp3u16(v3u16_t a, v3u16_t b, int x, int shift)

{

    const int xi = (1 << shift) - x;

    return (v3u16_t) {

        (a.x * xi + b.x * x) >> shift,

        (a.y * xi + b.y * x) >> shift,

        (a.z * xi + b.z * x) >> shift,

    };

}

static av_always_inline v3u16_t lookup_output(const SwsLut3D *lut3d, v3u16_t ipt)

{

    const int Ishift = 16 - OUTPUT_LUT_BITS_I;

    const int Cshift = 16 - OUTPUT_LUT_BITS_PT;

    const int Ix = ipt.x >> Ishift;

    const int Px = ipt.y >> Cshift;

    const int Tx = ipt.z >> Cshift;

    const int If = ipt.x & ((1 << Ishift) - 1);

    const int Pf = ipt.y & ((1 << Cshift) - 1);

    const int Tf = ipt.z & ((1 << Cshift) - 1);

    const int In = FFMIN(Ix + 1, OUTPUT_LUT_SIZE_I  - 1);

    const int Pn = FFMIN(Px + 1, OUTPUT_LUT_SIZE_PT - 1);

    const int Tn = FFMIN(Tx + 1, OUTPUT_LUT_SIZE_PT - 1);

    /* Trilinear interpolation */

    const v3u16_t c000 = lut3d->output[Tx][Px][Ix];

    const v3u16_t c001 = lut3d->output[Tx][Px][In];

    const v3u16_t c010 = lut3d->output[Tx][Pn][Ix];

    const v3u16_t c011 = lut3d->output[Tx][Pn][In];

    const v3u16_t c100 = lut3d->output[Tn][Px][Ix];

    const v3u16_t c101 = lut3d->output[Tn][Px][In];

    const v3u16_t c110 = lut3d->output[Tn][Pn][Ix];

    const v3u16_t c111 = lut3d->output[Tn][Pn][In];

    const v3u16_t c00  = lerp3u16(c000, c100, Tf, Cshift);

    const v3u16_t c10  = lerp3u16(c010, c110, Tf, Cshift);

    const v3u16_t c01  = lerp3u16(c001, c101, Tf, Cshift);

    const v3u16_t c11  = lerp3u16(c011, c111, Tf, Cshift);

    const v3u16_t c0   = lerp3u16(c00,  c10,  Pf, Cshift);

    const v3u16_t c1   = lerp3u16(c01,  c11,  Pf, Cshift);

    const v3u16_t c    = lerp3u16(c0,   c1,   If, Ishift);

    return c;

}

static av_always_inline v3u16_t apply_tone_map(const SwsLut3D *lut3d, v3u16_t ipt)

{

    const int shift = 16 - TONE_LUT_BITS;

    const int Ix = ipt.x >> shift;

    const int If = ipt.x & ((1 << shift) - 1);

    const int In = FFMIN(Ix + 1, TONE_LUT_SIZE - 1);

    const v2u16_t w0 = lut3d->tone_map[Ix];

    const v2u16_t w1 = lut3d->tone_map[In];

    const v2u16_t w  = lerp2u16(w0, w1, If, shift);

    const int base   = (1 << 15) - w.y;

    ipt.x = w.x;

    ipt.y = base + (ipt.y * w.y >> 15);

    ipt.z = base + (ipt.z * w.y >> 15);

    return ipt;

}

int ff_sws_lut3d_generate(SwsLut3D *lut3d, enum AVPixelFormat fmt_in,

                          enum AVPixelFormat fmt_out, const SwsColorMap *map)

{

    int ret;

    if (!ff_sws_lut3d_test_fmt(fmt_in, 0) || !ff_sws_lut3d_test_fmt(fmt_out, 1))

        return AVERROR(EINVAL);

    lut3d->dynamic = map->src.frame_peak.num > 0;

    lut3d->map = *map;

    if (lut3d->dynamic) {

        ret = ff_sws_color_map_generate_dynamic(&lut3d->input[0][0][0],

                                             &lut3d->output[0][0][0],

                                             INPUT_LUT_SIZE, OUTPUT_LUT_SIZE_I,

                                             OUTPUT_LUT_SIZE_PT, map);

        if (ret < 0)

            return ret;

        /* Make sure initial state is valid */

        ff_sws_lut3d_update(lut3d, &map->src);

        return 0;

    } else {

        return ff_sws_color_map_generate_static(&lut3d->input[0][0][0],

                                             INPUT_LUT_SIZE, map);

    }

}

void ff_sws_lut3d_update(SwsLut3D *lut3d, const SwsColor *new_src)

{

    if (!new_src || !lut3d->dynamic)

        return;

    lut3d->map.src.frame_peak = new_src->frame_peak;

    lut3d->map.src.frame_avg  = new_src->frame_avg;

    ff_sws_tone_map_generate(lut3d->tone_map, TONE_LUT_SIZE, &lut3d->map);

}

void ff_sws_lut3d_apply(const SwsLut3D *lut3d, const uint8_t *in, int in_stride,

                        uint8_t *out, int out_stride, int w, int h)

{

    while (h--) {

        const uint16_t *in16 = (const uint16_t *) in;

        uint16_t *out16 = (uint16_t *) out;

        for (int x = 0; x < w; x++) {

            v3u16_t c = { in16[0], in16[1], in16[2] };

            c = lookup_input16(lut3d, c);

            if (lut3d->dynamic) {

                c = apply_tone_map(lut3d, c);

                c = lookup_output(lut3d, c);

            }

            out16[0] = c.x;

            out16[1] = c.y;

            out16[2] = c.z;

            out16[3] = in16[3];

            in16  += 4;

            out16 += 4;

        }

        in  += in_stride;

        out += out_stride;

    }

}