/*
 * STDC Segmentation Post-Processing for KL630
 * Matches Python stdc630inference.py logic
 */

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

#include "kp_struct.h"
#include "ncpu_gen_struct.h"
#include "demo_post_utils.h"
#include "stdc_post_process.h"

typedef struct {
    uint8_t *prev_roi_luma;
    uint32_t prev_roi_capacity;
    uint32_t prev_roi_count;
    uint32_t frame_count;
    uint32_t consecutive_grass_frames;
    float prev_tree_roi_ratio;
} stdc_runtime_state_t;

static stdc_runtime_state_t g_stdc_state = {0};

static int stdc_ensure_prev_roi_capacity(uint32_t required)
{
    if (required <= g_stdc_state.prev_roi_capacity)
        return 0;

    uint8_t *new_buf = (uint8_t *)realloc(g_stdc_state.prev_roi_luma, required);
    if (NULL == new_buf) {
        printf("[stdc_post] realloc %u bytes failed\n", required);
        return -1;
    }

    g_stdc_state.prev_roi_luma = new_buf;
    g_stdc_state.prev_roi_capacity = required;
    return 0;
}

static int stdc_is_in_forward_roi(uint32_t row, uint32_t col, uint32_t rows, uint32_t cols)
{
    float y = ((float)row + 0.5f) / (float)rows;
    float x = ((float)col + 0.5f) / (float)cols;
    float t;
    float left;
    float right;

    if (y < FWD_ROI_TOP || y > FWD_ROI_BOTTOM)
        return 0;

    t = (y - FWD_ROI_TOP) / (FWD_ROI_BOTTOM - FWD_ROI_TOP);
    left = FWD_ROI_TOP_LEFT + (FWD_ROI_BOTTOM_LEFT - FWD_ROI_TOP_LEFT) * t;
    right = FWD_ROI_TOP_RIGHT + (FWD_ROI_BOTTOM_RIGHT - FWD_ROI_TOP_RIGHT) * t;

    return (x >= left && x <= right);
}

static float stdc_compute_motion_diff(struct kdp_image_s *image_p, uint8_t *is_moving)
{
    uint32_t raw_height = RAW_INPUT_ROW(image_p, 0);
    uint32_t raw_width = RAW_INPUT_COL(image_p, 0);
    uint32_t max_roi_pixels = raw_width * raw_height;
    uint8_t *y_plane = (uint8_t *)RAW_IMAGE_MEM_ADDR(image_p, 0);
    uint64_t diff_sum = 0;
    uint32_t roi_count = 0;
    float motion_diff = 0.0f;

    *is_moving = 1;

    if (0 == raw_width || 0 == raw_height || NULL == y_plane)
        return 0.0f;

    if (0 != stdc_ensure_prev_roi_capacity(max_roi_pixels))
        return 0.0f;

    for (uint32_t row = 0; row < raw_height; row++) {
        for (uint32_t col = 0; col < raw_width; col++) {
            if (!stdc_is_in_forward_roi(row, col, raw_height, raw_width))
                continue;

            uint8_t cur = y_plane[row * raw_width + col];
            if (g_stdc_state.prev_roi_count == roi_count + 1 || g_stdc_state.prev_roi_count > roi_count)
                diff_sum += (uint64_t)abs((int)cur - (int)g_stdc_state.prev_roi_luma[roi_count]);

            g_stdc_state.prev_roi_luma[roi_count++] = cur;
        }
    }

    if (roi_count > 0 && g_stdc_state.prev_roi_count == roi_count) {
        motion_diff = (float)diff_sum / (float)roi_count;
        *is_moving = (motion_diff >= STDC_MOTION_THRESHOLD) ? 1 : 0;
    }

    g_stdc_state.prev_roi_count = roi_count;
    return motion_diff;
}

/*
 * STDC NPU output: 8 channels x ~46 rows x ~96 cols (int8 CHW)
 * Argmax across channel axis gives the predicted class per pixel.
 * Scale/radix are the same for all channels, so argmax of int8
 * values is identical to argmax of dequantized floats.
 */
int stdc_post_process(int model_id, struct kdp_image_s *image_p)
{
    stdc_post_proc_params_t *params = (stdc_post_proc_params_t *)POSTPROC_PARAMS_P(image_p);
    float fps = 1.0f;

    struct output_node seg_node;
    get_output_node(&seg_node, image_p, 0);

    uint32_t num_ch  = seg_node.ch;   /* 8 */
    uint32_t num_row = seg_node.row;  /* ~46 */
    uint32_t num_col = seg_node.col;  /* ~96 */
    uint32_t total_pixels = num_row * num_col;

    if (params && params->fps > 0.0f)
        fps = params->fps;

    if (num_ch == 0 || num_row == 0 || num_col == 0 || total_pixels == 0) {
        printf("[stdc_post] invalid output node: ch=%u row=%u col=%u\n",
               num_ch, num_row, num_col);
        return KP_FW_INFERENCE_ERROR_101;
    }

    /* -------------------------------------------------------
     * 0. Initialise result buffer early so seg_map can be filled in the loop
     * ------------------------------------------------------- */
    stdc_result_t *result = (stdc_result_t *)POSTPROC_RESULT_MEM_ADDR(image_p);
    memset(result, 0, sizeof(stdc_result_t));

    /* -------------------------------------------------------
     * 1. Argmax: find winning class for every spatial pixel
     * ------------------------------------------------------- */
    uint32_t class_count[STDC_NUM_CLASSES] = {0};

    /* ROI pixel counters */
    uint32_t fwd_total = 0,  fwd_grass = 0, fwd_tree = 0;
    uint32_t col_total = 0;
    uint32_t col_count[STDC_NUM_CLASSES] = {0};
    float motion_diff;
    uint8_t is_moving = 1;

    /* Collision ROI boundaries in pixel coords */
    uint32_t col_r0 = (uint32_t)(COL_ROI_LEFT   * num_col);
    uint32_t col_r1 = (uint32_t)(COL_ROI_RIGHT   * num_col);
    uint32_t col_c0 = (uint32_t)(COL_ROI_TOP     * num_row);
    uint32_t col_c1 = (uint32_t)(COL_ROI_BOTTOM  * num_row);

    motion_diff = stdc_compute_motion_diff(image_p, &is_moving);

    for (uint32_t r = 0; r < num_row; r++) {
        for (uint32_t c = 0; c < num_col; c++) {
            /* Argmax across channels */
            int8_t best_val = *get_data(seg_node, 0, r, c);
            uint32_t best_cls = 0;
            for (uint32_t ch = 1; ch < num_ch && ch < STDC_NUM_CLASSES; ch++) {
                int8_t v = *get_data(seg_node, ch, r, c);
                if (v > best_val) {
                    best_val = v;
                    best_cls = ch;
                }
            }

            class_count[best_cls]++;

            /* Store argmax label for seg overlay */
            uint32_t seg_idx = r * num_col + c;
            if (seg_idx < STDC_SEG_MAP_MAX)
                result->seg_map[seg_idx] = (uint8_t)best_cls;

            /* Forward ROI */
            if (stdc_is_in_forward_roi(r, c, num_row, num_col)) {
                fwd_total++;
                if (best_cls == STDC_CLASS_GRASS)
                    fwd_grass++;
                if (best_cls == STDC_CLASS_TREE)
                    fwd_tree++;
            }

            /* Collision ROI */
            if (r >= col_c0 && r < col_c1 && c >= col_r0 && c < col_r1) {
                col_total++;
                col_count[best_cls]++;
            }
        }
    }

    /* -------------------------------------------------------
     * 2. Compute ratios and fill result struct
     * ------------------------------------------------------- */
    stdc_analysis_t *ana = &result->analysis;
    ana->seg_width  = num_col;
    ana->seg_height = num_row;
    ana->frame_count = ++g_stdc_state.frame_count;
    ana->motion_diff = motion_diff;
    ana->is_moving = is_moving;

    float inv_total = (total_pixels > 0) ? 1.0f / (float)total_pixels : 0.0f;
    for (uint32_t i = 0; i < STDC_NUM_CLASSES; i++)
        ana->class_ratio[i] = (float)class_count[i] * inv_total;

    /* Forward ROI grass ratio */
    ana->grass_roi_ratio = (fwd_total > 0)
        ? (float)fwd_grass / (float)fwd_total : 0.0f;
    ana->tree_roi_ratio = (fwd_total > 0)
        ? (float)fwd_tree / (float)fwd_total : 0.0f;
    ana->tree_roi_growth = ana->tree_roi_ratio - g_stdc_state.prev_tree_roi_ratio;
    g_stdc_state.prev_tree_roi_ratio = ana->tree_roi_ratio;

    /* Collision ROI ratios */
    float inv_col = (col_total > 0) ? 1.0f / (float)col_total : 0.0f;
    ana->col_person_ratio = (float)col_count[STDC_CLASS_PERSON] * inv_col;
    ana->col_car_ratio    = (float)col_count[STDC_CLASS_CAR]    * inv_col;
    ana->col_tree_ratio   = (float)col_count[STDC_CLASS_TREE]   * inv_col;
    ana->col_bunker_ratio = (float)col_count[STDC_CLASS_BUNKER] * inv_col;
    ana->col_pond_ratio   = (float)col_count[STDC_CLASS_POND]   * inv_col;

    /* -------------------------------------------------------
     * 3. Warning flags
     * ------------------------------------------------------- */
    ana->on_grass = (ana->grass_roi_ratio > THR_GRASS_ROI) ? 1 : 0;
    if (ana->on_grass && ana->is_moving)
        g_stdc_state.consecutive_grass_frames++;
    else
        g_stdc_state.consecutive_grass_frames = 0;

    ana->grass_duration_sec = (float)g_stdc_state.consecutive_grass_frames / fps;
    ana->grass_warning = (ana->grass_duration_sec >= STDC_WARNING_SECONDS) ? 1 : 0;
    ana->person_warning = (ana->class_ratio[STDC_CLASS_PERSON] > THR_PERSON_GLOBAL) ? 1 : 0;
    ana->car_warning    = (ana->class_ratio[STDC_CLASS_CAR]    > THR_CAR_GLOBAL)    ? 1 : 0;
    ana->greenery_warning = (ana->class_ratio[STDC_CLASS_GREENERY] > THR_GREENERY_GLOBAL) ? 1 : 0;
    ana->bunker_warning   = (ana->class_ratio[STDC_CLASS_BUNKER]   > THR_BUNKER_GLOBAL)   ? 1 : 0;
    ana->pond_warning     = (ana->class_ratio[STDC_CLASS_POND]     > THR_POND_GLOBAL)     ? 1 : 0;
    ana->tree_dense     = (ana->class_ratio[STDC_CLASS_TREE]   > THR_TREE_DENSE)    ? 1 : 0;
    ana->tree_approaching = (ana->tree_roi_growth > STDC_TREE_GROWTH_THR) ? 1 : 0;

    ana->collision_risk = 0;
    if (ana->col_person_ratio >= THR_PERSON_COLLISION) ana->collision_risk = 1;
    if (ana->col_car_ratio    >= THR_CAR_COLLISION)    ana->collision_risk = 1;
    if (ana->col_tree_ratio   >= THR_TREE_COLLISION)   ana->collision_risk = 1;
    if (ana->col_bunker_ratio >= THR_BUNKER_COLLISION) ana->collision_risk = 1;
    if (ana->col_pond_ratio   >= THR_POND_COLLISION)   ana->collision_risk = 1;

    /* -------------------------------------------------------
     * 4. Log summary
     * ------------------------------------------------------- */
        printf("[STDC] frame=%u moving=%d diff=%.2f grass_roi=%.1f%% grass_time=%.1fs grass_warn=%d "
            "person=%.1f%% car=%.1f%% tree=%.1f%% tree_growth=%.1f%% collision=%d\n",
            ana->frame_count, ana->is_moving, ana->motion_diff,
            ana->grass_roi_ratio  * 100.0f, ana->grass_duration_sec, ana->grass_warning,
           ana->class_ratio[STDC_CLASS_PERSON] * 100.0f,
           ana->class_ratio[STDC_CLASS_CAR]    * 100.0f,
           ana->class_ratio[STDC_CLASS_TREE]   * 100.0f,
            ana->tree_roi_growth * 100.0f,
           ana->collision_risk);

    return KP_SUCCESS;
}
