// SPDX-License-Identifier: Apache-2.0
// kpre (tt/kpre.py, OPT_PLAN R4 device pre-processing), reader (NCRISC). One camera's uint8 node image rows
// (1600 x 3, RGB or BGR as uploaded) -> per H row r: the two lerp operands of the C14 x0.8 horizontal resize as fp32
// pseudo-tiles (1024 contiguous values; the compute kernel's unpack -> SFPU -> pack round trip keeps the order):
//   P0[j] = LUT[c][u8[r][(5k + p) * 3 + s]],  P1[j] = LUT[c][u8[r][(5k + p + 1) * 3 + s]],  j = (4k + p) * 3 + c
// (output pixel x = 4k + p, phase p; s = 2 - c for RGB input) where LUT[c][u] = f32(u + (-mean_c)) (the C14
// normalisation, computed on the host: exact by construction). Once per program: the lerp weight pattern (4 fp32
// pseudo-tiles, w[j] = (2p + 1) / 8) and the 4 vertical weight tiles ((2q + 1) / 8 everywhere) into resident CBs.
// Data movement only; no semaphores, no multicast; common runtime args only (P3): the 4 buffer addresses.
#include <cstdint>

#include "api/dataflow/dataflow_api.h"
#include "kpre_common.h"

constexpr uint32_t N = get_compile_time_arg_val(0);
constexpr uint32_t S = get_compile_time_arg_val(1);
constexpr uint32_t W = get_compile_time_arg_val(2);
constexpr uint32_t RH = get_compile_time_arg_val(3);
constexpr uint32_t NC = get_compile_time_arg_val(4);
constexpr uint32_t GX = get_compile_time_arg_val(5);
constexpr uint32_t RA = get_compile_time_arg_val(6);
constexpr uint32_t RB = get_compile_time_arg_val(7);
constexpr uint32_t RD = get_compile_time_arg_val(8);
constexpr uint32_t SWAP = get_compile_time_arg_val(9);      // 1: the image is RGB (network channel c reads 2 - c)
constexpr uint32_t ROW_BYTES = get_compile_time_arg_val(10);  // node row: 1600 x 3 bytes
constexpr uint32_t PERIODS = get_compile_time_arg_val(11);  // W / 4 (4 output pixels from 5 node pixels)
constexpr uint32_t SPLIT = get_compile_time_arg_val(12);   // 1: the writer RISC builds the odd rows (cb_p2)
constexpr uint32_t U16 = get_compile_time_arg_val(13);     // 1: raw uint16 planes (normalised on the SFPU)
constexpr auto img_args = TensorAccessorArgs<14>();
constexpr auto lut_args = TensorAccessorArgs<img_args.next_compile_time_args_offset()>();
constexpr auto wx_args = TensorAccessorArgs<lut_args.next_compile_time_args_offset()>();
constexpr auto wy_args = TensorAccessorArgs<wx_args.next_compile_time_args_offset()>();
constexpr auto beta_args = TensorAccessorArgs<wy_args.next_compile_time_args_offset()>();

constexpr uint32_t cb_p = 0, cb_wx = 1, cb_wy = 2, cb_scratch = 4, cb_beta = 7;
constexpr uint32_t TB32 = 4096, TILE_WORDS = 1024;
constexpr uint32_t PT = 5;   // pseudo-tiles per node row: its 5 phase planes
static_assert(ROW_BYTES % 64 == 0, "node rows must be 64-byte multiples (DRAM read alignment)");
static_assert(PERIODS * 3 <= TILE_WORDS, "a phase plane fits one pseudo-tile");

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    const KpreGeom g = kpre_geom<N, S, W, RH, NC, RA, RB, RD>(core);
    if (!g.rows) {
        return;
    }
    const auto img = TensorAccessor(img_args, get_common_arg_val<uint32_t>(0), ROW_BYTES);
    const auto lut_t = TensorAccessor(lut_args, get_common_arg_val<uint32_t>(1), 3 * 256 * 4);
    const auto wx_t = TensorAccessor(wx_args, get_common_arg_val<uint32_t>(2), 4 * TB32);
    const auto wy_t = TensorAccessor(wy_args, get_common_arg_val<uint32_t>(3), 4 * TB32);
    // scratch: LUT (3 KB) then two node rows (double buffer)
    const uint32_t base = (get_write_ptr(cb_scratch) + 63u) & ~63u;
    const uint32_t lut_l1 = base;
    const uint32_t row_l1[2] = {base + 3072, base + 3072 + ROW_BYTES};
    noc_async_read(lut_t.get_noc_addr(0), lut_l1, 3 * 256 * 4);
    cb_reserve_back(cb_wx, 4);
    noc_async_read(wx_t.get_noc_addr(0), get_write_ptr(cb_wx), 4 * TB32);
    cb_reserve_back(cb_wy, 4);
    noc_async_read(wy_t.get_noc_addr(0), get_write_ptr(cb_wy), 4 * TB32);
    if constexpr (U16) {   // the per-lane -mean_c tile (normalisation on the SFPU)
        const auto beta_t = TensorAccessor(beta_args, get_common_arg_val<uint32_t>(5), TB32);
        cb_reserve_back(cb_beta, 1);
        noc_async_read(beta_t.get_noc_addr(0), get_write_ptr(cb_beta), TB32);
    }
    noc_async_read(img.get_noc_addr(g.rlo), row_l1[0], ROW_BYTES);
    noc_async_read_barrier();
    cb_push_back(cb_wx, 4);
    cb_push_back(cb_wy, 4);
    if constexpr (U16) {
        cb_push_back(cb_beta, 1);
    }
    const uint32_t* lut = reinterpret_cast<const uint32_t*>(lut_l1);
    constexpr uint32_t STEP = SPLIT ? 2 : 1;   // SPLIT: rows rlo, rlo + 2, ... here; the writer builds the others
    uint32_t buf = 0;
    for (uint32_t r = g.rlo; r <= g.rhi; r += STEP) {
        if (r + STEP <= g.rhi) {   // prefetch the next node row
            noc_async_read(img.get_noc_addr(r + STEP), row_l1[buf ^ 1], ROW_BYTES);
        }
        cb_reserve_back(cb_p, PT);
        uint32_t* d0 = reinterpret_cast<uint32_t*>(get_write_ptr(cb_p));
        const uint8_t* src = reinterpret_cast<const uint8_t*>(row_l1[buf]);
        asm volatile("" ::: "memory");
        if constexpr (U16) {
            kpre_build_row_u16<PERIODS, SWAP>(src, reinterpret_cast<uint16_t*>(d0));
        } else {
            kpre_build_row<PERIODS, SWAP>(src, lut, d0);
        }
        asm volatile("" ::: "memory");
        cb_push_back(cb_p, PT);
        noc_async_read_barrier();
        buf ^= 1;
    }
}