// SPDX-License-Identifier: Apache-2.0
// EPI writer (sceneseg-p150 OPT_PLAN EPI, tt/epi.py): the output tiles of this core's contiguous range (work split as
// the reader) to the interleaved output(s), flushed per tile, one barrier at the end: the main output (CB 16, when
// HAS_MAIN) and, with SPLIT (OPT_REPORT R2.3), the bf16 hi / lo terms (CBs 17 / 18).
// CLASS (K8, OPT_REPORT R2.4): the main output is the fp32 class-margin tile [m0, m1, m2, pad...] of 32 pixels; it is
// not written: the class of each pixel is decided here, exactly as ttnn.argmax(dim=1) over (m0, m1, m2) on the SFPU
// (argmax_nc_compute.cpp: running max with gt_binary_tile, strict '>', lowest index on ties; gt(a, b) = a > b in the
// sign-magnitude total order, false when both are zero / subnormal: ckernel_sfpu_binary_comp.h), and the uint8 class
// of the UNIT * 32 pixels of one image row is written as one row page of the uint8 ROW_MAJOR class map.
// CT args: [0] T, [1] NC, [2] GX, [3] OTB (main output tile bytes), [4] HAS_MAIN, [5] SPLIT, [6] UNIT, [7] CLASS,
// then the TensorAccessorArgs of the main (CLASS: the class map), hi and lo outputs (repeats of the first present
// one when absent). Common RT args: [main_addr, hi_addr, lo_addr].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t T = get_compile_time_arg_val(0);
constexpr uint32_t NC = get_compile_time_arg_val(1);
constexpr uint32_t GX = get_compile_time_arg_val(2);
constexpr uint32_t OTB = get_compile_time_arg_val(3);
constexpr uint32_t HAS_MAIN = get_compile_time_arg_val(4);
constexpr uint32_t SPLIT = get_compile_time_arg_val(5);
constexpr uint32_t UNIT = get_compile_time_arg_val(6);
constexpr uint32_t CLASS = get_compile_time_arg_val(7);
constexpr auto out_args = TensorAccessorArgs<8>();
constexpr auto hi_args = TensorAccessorArgs<out_args.next_compile_time_args_offset()>();
constexpr auto lo_args = TensorAccessorArgs<hi_args.next_compile_time_args_offset()>();
constexpr uint32_t cb_out = 16, cb_hi = 17, cb_lo = 18, cb_cls = 19;
constexpr uint32_t TB16 = 2048;
constexpr uint32_t ROWB = UNIT * 32;                       // CLASS: bytes (pixels) per class-map row page

// SFPU gt(a, b) on fp32 bit patterns (no NaN in the margins)
FORCE_INLINE bool sfpu_gt(uint32_t a, uint32_t b) {
    if ((a & 0x7F800000u) == 0 && (b & 0x7F800000u) == 0) {
        return false;                                       // both zero / subnormal: equal
    }
    const int32_t ka = (a & 0x80000000u) ? -(int32_t)(a & 0x7FFFFFFFu) : (int32_t)a;
    const int32_t kb = (b & 0x80000000u) ? -(int32_t)(b & 0x7FFFFFFFu) : (int32_t)b;
    return ka > kb;
}

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = T / UNIT / NC, E = T / UNIT % NC;
    const uint32_t t0 = (core * B + (core < E ? core : E)) * UNIT;
    const uint32_t n = (B + (core < E ? 1 : 0)) * UNIT;
    if constexpr (CLASS) {
        const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(0), ROWB);
        const uint32_t row = get_write_ptr(cb_cls);
        volatile tt_l1_ptr uint8_t* cls = reinterpret_cast<volatile tt_l1_ptr uint8_t*>(row);
        for (uint32_t t = t0; t < t0 + n; ++t) {
            const uint32_t j = t % UNIT;
            if (j == 0 && t != t0) {
                noc_async_write_barrier();                  // the previous row page left the buffer
            }
            cb_wait_front(cb_out, 1);
            volatile tt_l1_ptr uint32_t* m = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_read_ptr(cb_out));
            for (uint32_t r = 0; r < 32; ++r) {
                const uint32_t w = (r >> 4) * 512 + (r & 15) * 16;  // face (r / 16) * 2, row r % 16, column 0
                uint32_t best = m[w];
                uint8_t c = 0;
                const uint32_t m1 = m[w + 1], m2 = m[w + 2];
                if (sfpu_gt(m1, best)) {
                    best = m1;
                    c = 1;
                }
                if (sfpu_gt(m2, best)) {
                    c = 2;
                }
                cls[j * 32 + r] = c;
            }
            cb_pop_front(cb_out, 1);
            if (j == UNIT - 1) {
                noc_async_write(row, out.get_noc_addr(t / UNIT), ROWB);
            }
        }
        noc_async_write_barrier();
        return;
    }
    const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(0), OTB);
    const auto hi = TensorAccessor(hi_args, get_common_arg_val<uint32_t>(1), TB16);
    const auto lo = TensorAccessor(lo_args, get_common_arg_val<uint32_t>(2), TB16);
    for (uint32_t t = t0; t < t0 + n; ++t) {
        if constexpr (HAS_MAIN) {
            cb_wait_front(cb_out, 1);
            noc_async_write(get_read_ptr(cb_out), out.get_noc_addr(t), OTB);
        }
        if constexpr (SPLIT) {
            cb_wait_front(cb_hi, 1);
            cb_wait_front(cb_lo, 1);
            noc_async_write(get_read_ptr(cb_hi), hi.get_noc_addr(t), TB16);
            noc_async_write(get_read_ptr(cb_lo), lo.get_noc_addr(t), TB16);
        }
        noc_async_writes_flushed();
        if constexpr (HAS_MAIN) {
            cb_pop_front(cb_out, 1);
        }
        if constexpr (SPLIT) {
            cb_pop_front(cb_hi, 1);
            cb_pop_front(cb_lo, 1);
        }
    }
    noc_async_write_barrier();
}