// firmware/lib/cv/cv.cpp
#include "cv.h"
#include <string.h>
#include <math.h>
#include <algorithm>
#include <vector>

void cv_init(CVState& state) {
    // Initialize members directly — avoid CVState{} temporary which puts 9KB on stack
    memset(state.background, 0, sizeof(state.background));
    state.bg_valid          = false;
    state.last_motion_frame = 0;
    state.frame_index       = 0;
    state.next_id           = 1;
    state.tracks.clear();
    state.entries = 0;
    state.exits   = 0;
    state.last_entry_frame = 0;
    state.last_exit_frame  = 0;
}

void cv_reset_counts(CVState& state) {
    state.entries = 0;
    state.exits   = 0;
}

struct Point { int x, y; };

// Note: queue may grow to CV_PIXELS entries (~72KB) on large blobs.
// Requires PSRAM (enabled via -DBOARD_HAS_PSRAM in platformio.ini).
// BFS flood fill. Marks visited pixels (sets fg to 0). Returns {-1,-1} if blob < CV_MIN_BLOB_PX.
static std::pair<float,float> extract_blob(uint8_t* fg, int start_x, int start_y) {
    std::vector<Point> queue;
    queue.reserve(512);
    queue.push_back({start_x, start_y});
    fg[start_y * CV_W + start_x] = 0;

    float sum_x = 0, sum_y = 0;
    int count = 0;

    while (!queue.empty()) {
        Point p = queue.back(); queue.pop_back();
        sum_x += p.x; sum_y += p.y; count++;

        const int dx[] = {-1, 1, 0, 0};
        const int dy[] = {0, 0, -1, 1};
        for (int d = 0; d < 4; d++) {
            int nx = p.x + dx[d], ny = p.y + dy[d];
            if (nx < 0 || nx >= CV_W || ny < 0 || ny >= CV_H) continue;
            int ni = ny * CV_W + nx;
            if (!fg[ni]) continue;
            fg[ni] = 0;
            queue.push_back({nx, ny});
        }
    }

    if (count < CV_MIN_BLOB_PX) return {-1.0f, -1.0f};
    return {sum_x / count, sum_y / count};
}

static std::vector<std::pair<float,float>> find_centroids(const uint8_t* fg) {
    std::vector<std::pair<float,float>> result;
    static uint8_t fg_copy[CV_PIXELS];  // static to avoid 9KB stack allocation
    memcpy(fg_copy, fg, CV_PIXELS);

    for (int y = 0; y < CV_H; y++) {
        for (int x = 0; x < CV_W; x++) {
            if (!fg_copy[y * CV_W + x]) continue;
            auto c = extract_blob(fg_copy, x, y);
            if (c.first >= 0) result.push_back(c);
        }
    }
    return result;
}

static void frame_diff(const uint8_t* frame, const uint8_t* bg,
                       uint8_t* fg, int pixels) {
    for (int i = 0; i < pixels; i++) {
        int diff = (int)frame[i] - (int)bg[i];
        if (diff < 0) diff = -diff;
        fg[i] = (diff > CV_DIFF_THRESH) ? 1 : 0;
    }
}

CVResult cv_process(CVState& state, const uint8_t* frame, uint8_t line_pct) {
    CVResult result = {0, 0};
    state.frame_index++;

    if (!state.bg_valid) {
        memcpy(state.background, frame, CV_PIXELS);
        state.bg_valid = true;
        return result;
    }

    static uint8_t fg[CV_PIXELS];  // static: avoids 9KB on task stack
    frame_diff(frame, state.background, fg, CV_PIXELS);

    int fg_count = 0;
    for (int i = 0; i < CV_PIXELS; i++) fg_count += fg[i];

    bool motion = fg_count > CV_MIN_BLOB_PX;
    if (!motion) {
        if (state.frame_index - state.last_motion_frame > 10) {
            memcpy(state.background, frame, CV_PIXELS);
        }
        for (auto& t : state.tracks) t.missed++;
        state.tracks.erase(
            std::remove_if(state.tracks.begin(), state.tracks.end(),
                [](const CVTrack& t){ return t.missed > CV_MAX_MISSED; }),
            state.tracks.end());
        return result;
    }

    state.last_motion_frame = state.frame_index;

    auto centroids = find_centroids(fg);

    std::vector<bool> centroid_matched(centroids.size(), false);

    for (auto& track : state.tracks) {
        float best_dist = CV_MAX_MOVE * CV_MAX_MOVE;
        int   best_idx  = -1;

        for (int i = 0; i < (int)centroids.size(); i++) {
            if (centroid_matched[i]) continue;
            float dx = centroids[i].first  - track.x;
            float dy = centroids[i].second - track.y;
            float d2 = dx*dx + dy*dy;
            if (d2 < best_dist) { best_dist = d2; best_idx = i; }
        }

        if (best_idx >= 0) {
            centroid_matched[best_idx] = true;
            track.x = centroids[best_idx].first;
            track.y = centroids[best_idx].second;
            track.missed = 0;
        } else {
            track.missed++;
        }
    }

    state.tracks.erase(
        std::remove_if(state.tracks.begin(), state.tracks.end(),
            [](const CVTrack& t){ return t.missed > CV_MAX_MISSED; }),
        state.tracks.end());

    float line_y = (line_pct / 100.0f) * CV_H;
    for (int i = 0; i < (int)centroids.size(); i++) {
        if (centroid_matched[i]) continue;
        CVTrack t;
        t.id         = state.next_id++;
        t.x          = centroids[i].first;
        t.y          = centroids[i].second;
        t.spawn_y    = t.y;
        t.above_line = (t.y < line_y);
        t.counted    = false;
        t.missed     = 0;
        state.tracks.push_back(t);
    }
    // Directional crossing check. A track counts at most once, and only if it
    // spawned clearly on one side of the line AND is now clearly on the other.
    // This rejects blobs that wobble around the line (shadows, body straddling
    // the line, track churn at spawn) — only a true traversal fires an event.
    for (auto& track : state.tracks) {
        if (track.missed > 0)       continue; // only check tracks matched this frame
        if (track.counted)          continue; // one track = one trip

        bool spawned_above = track.spawn_y < (line_y - CV_TRAVERSAL_MARGIN_PX);
        bool spawned_below = track.spawn_y > (line_y + CV_TRAVERSAL_MARGIN_PX);
        bool now_above_firm = track.y   < (line_y - CV_TRAVERSAL_MARGIN_PX);
        bool now_below_firm = track.y   > (line_y + CV_TRAVERSAL_MARGIN_PX);

        if (spawned_above && now_below_firm) {
            bool in_cooldown = state.last_entry_frame != 0 &&
                (state.frame_index - state.last_entry_frame) < CV_CROSSING_COOLDOWN_FRAMES;
            if (!in_cooldown) {
                state.entries++;
                result.entries_delta++;
                state.last_entry_frame = state.frame_index;
                track.counted = true;
            }
        } else if (spawned_below && now_above_firm) {
            bool in_cooldown = state.last_exit_frame != 0 &&
                (state.frame_index - state.last_exit_frame) < CV_CROSSING_COOLDOWN_FRAMES;
            if (!in_cooldown) {
                state.exits++;
                result.exits_delta++;
                state.last_exit_frame = state.frame_index;
                track.counted = true;
            }
        }

        track.above_line = (track.y < line_y);
    }

    return result;
}