wavelets2d.h

#pragma once

#include <stdint.h>
#include <string.h>

#define WAVELET_DIM 512

extern void wavelet_forward_2d(uint8_t *mat, size_t N);
extern void wavelet_inverse_2d(uint8_t *mat, size_t N);

static inline uint8_t haar_newY(uint8_t y1, uint8_t y2) {
    return 2 * y2 - y1;
}

static inline void wavelet_predict_horizontal(uint8_t *vec, size_t n, bool invert) {
    size_t const half = n >> 1;
    uint32_t predict = 0;
    for (size_t i = 0; i < half; i++) {
        size_t j = i + half;
        if (i < half - 1) {
            predict = (vec[i] + vec[i + 1]) / 2;
        }
        else if (n == 2) {
            predict = vec[0];
        }
        else {
            predict = haar_newY(vec[i - 1], vec[i]);
        }
        if (!invert)
            vec[j] = vec[j] - predict;
        else
            vec[j] = vec[j] + predict;
    }
}

static inline void wavelet_predict_vertical(uint8_t *vec, size_t n, bool invert) {
    size_t const half = n >> 1;
    uint32_t predict = 0;
    for (size_t i = 0; i < half; i++) {
        size_t i0 = i * WAVELET_DIM;
        size_t i1 = (i * WAVELET_DIM) + WAVELET_DIM;
        size_t i2 = (i * WAVELET_DIM) + WAVELET_DIM;
        size_t j = (i * WAVELET_DIM) + (half * WAVELET_DIM);
        if (i < half - 1) {
            predict = (vec[i0] + vec[i1]) / 2;
        }
        else if (n == 2) {
            predict = vec[0];
        }
        else {
            predict = haar_newY(vec[i2], vec[i0]);
        }
        if (!invert)
            vec[j] = vec[j] - predict;
        else
            vec[j] = vec[j] + predict;
    }
}

static inline void wavelet_forward_horizontal(uint8_t *mat_a, uint8_t *mat_b, size_t n) {
    size_t half = n >> 1;
    for (size_t y = 0; y < n; y++) {
        size_t o = sizeof(uint8_t) * y * WAVELET_DIM;
        uint8_t *vec_a = mat_a + o;
        uint8_t *vec_b = mat_b + o;
        for (size_t i = 0, j = 0; i < n; i = i + 2, j++) {
            vec_b[j] = vec_a[i];
            vec_b[half + j] = vec_a[i + 1];
        }
        wavelet_predict_horizontal(vec_b, n, false);
    }
}

static inline void wavelet_forward_vertical(uint8_t *mat_a, uint8_t *mat_b, size_t n) {
    size_t half = n >> 1;
    for (size_t x = 0; x < n; x++) {
        size_t o = sizeof(uint8_t) * x;
        uint8_t *vec_a = mat_a + o;
        uint8_t *vec_b = mat_b + o;
        for (size_t i = 0, j = 0; i < n; i = i + 2, j++) {
            size_t i0 = i * WAVELET_DIM;
            size_t i1 = (i * WAVELET_DIM) + WAVELET_DIM;
            size_t j0 = j * WAVELET_DIM;
            size_t j1 = (j * WAVELET_DIM) + (half * WAVELET_DIM);
            vec_b[j0] = vec_a[i0];
            vec_b[j1] = vec_a[i1];
        }
        wavelet_predict_vertical(vec_b, n, false);
    }
}

static inline void wavelet_inverse_horizontal(uint8_t *mat_a, uint8_t *mat_b, size_t n) {
    size_t half = n >> 1;
    for (size_t y = 0; y < n; y++) {
        size_t o = sizeof(uint8_t) * y * WAVELET_DIM;
        uint8_t *vec_a = mat_a + o;
        uint8_t *vec_b = mat_b + o;
        wavelet_predict_horizontal(vec_a, n, true);
        for (size_t i = 0, j = 0; i < n; i = i + 2, j++) {
            vec_b[i] = vec_a[j];
            vec_b[i + 1] = vec_a[half + j];
        }
    }
}

static inline void wavelet_inverse_vertical(uint8_t *mat_a, uint8_t *mat_b, size_t n) {
    size_t half = n >> 1;
    for (size_t x = 0; x < n; x++) {
        size_t o = sizeof(uint8_t) * x;
        uint8_t *vec_a = mat_a + o;
        uint8_t *vec_b = mat_b + o;
        wavelet_predict_vertical(vec_a, n, true);
        for (size_t i = 0, j = 0; i < n; i = i + 2, j++) {
            size_t i0 = i * WAVELET_DIM;
            size_t i1 = (i * WAVELET_DIM) + WAVELET_DIM;
            size_t j0 = j * WAVELET_DIM;
            size_t j1 = (j * WAVELET_DIM) + (half * WAVELET_DIM);
            vec_b[i0] = vec_a[j0];
            vec_b[i1] = vec_a[j1];
        }
    }
}

static inline void swap(uint8_t **p1, uint8_t **p2) {
    uint8_t *temp = *p1;
    *p1 = *p2;
    *p2 = temp;
}

void wavelet_forward_2d(uint8_t *mat, size_t N) {
    uint8_t mat_temp[WAVELET_DIM * WAVELET_DIM];
    memset(mat_temp, 0, WAVELET_DIM * WAVELET_DIM);

    uint8_t *mat_a = mat;
    uint8_t *mat_b = &mat_temp[0];

    for (size_t n = N; n > 1; n = n >> 1) {
        wavelet_forward_horizontal(mat_a, mat_b, n);
        swap(&mat_a, &mat_b);
        wavelet_forward_vertical(mat_a, mat_b, n);
        swap(&mat_a, &mat_b);
    }
}

void wavelet_inverse_2d(uint8_t *mat, size_t N) {
    uint8_t mat_temp[WAVELET_DIM * WAVELET_DIM];
    memset(mat_temp, 0, WAVELET_DIM * WAVELET_DIM);

    uint8_t *mat_a = mat;
    uint8_t *mat_b = &mat_temp[0];

    for (size_t n = 2; n <= N; n = n << 1) {
        wavelet_inverse_vertical(mat_a, mat_b, n);
        swap(&mat_a, &mat_b);
        wavelet_inverse_horizontal(mat_a, mat_b, n);
        swap(&mat_a, &mat_b);
    }
}