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

static void event_reset(CVState& s) {
    s.event_active       = false;
    s.event_start_frame  = 0;
    s.event_frame_count  = 0;
    s.event_peak_n       = 0;
    s.event_first_c      = -1.0f;
    s.event_last_c       = -1.0f;
    s.event_min_c        = (float)CV_H;
    s.event_max_c        = -1.0f;
    s.event_min_y_seen   = CV_H;
    s.event_max_y_seen   = -1;
    s.event_quiet_count  = 0;
}

void cv_init(CVState& state) {
    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_fire_frame = 0;
    event_reset(state);
}

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

struct Point { int x, y; };

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];
    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;
    }
}

// Decide whether the just-ended event should fire and in which direction.
// Up-through-frame (centroid excursion from high y toward low y) maps to
// ENTRY per mount convention.
static void finalize_event(CVState& s, CVResult& result) {
    if (s.event_frame_count < CV_EVENT_MIN_FRAMES) return;
    // Note: no MAX_FRAMES rejection here. An event that runs the full duration
    // may still be a valid walker whose fg_count stayed above EXIT_THRESH due
    // to a stale bg or an AEC-driven lighting shift. Extent + MIN_TRAJ gates
    // below already reject stationary-person / wobble events.
    if (s.event_min_y_seen > CV_EVENT_EXTENT_TOP)  return;
    if (s.event_max_y_seen < CV_EVENT_EXTENT_BOT)  return;

    // Direction from centroid excursion relative to event start.
    // up_score: how far centroid excursed upward (smaller y) from first_c.
    // down_score: how far it excursed downward (larger y) from first_c.
    float up_score   = s.event_first_c - s.event_min_c;
    float down_score = s.event_max_c   - s.event_first_c;

    float winning = (up_score >= down_score) ? up_score : down_score;
    if (winning < CV_EVENT_MIN_TRAJ) return;

    // Timeout-aware direction. Quiet-exit events (fg fell below EXIT_THRESH)
    // have walker fully out of frame → min/max excursion bracket the true
    // traversal and up/down scores are reliable. Timeout events (event hit
    // MAX_FRAMES while still elevated) captured both an approach and a
    // departure within the window, so excursion measures the walker's
    // *range in frame* rather than direction — an entry walker who paused
    // near the top, then drifted back toward the middle before timeout
    // gets (wrongly) called an entry by up-score even though net motion is
    // mixed. For those, the net first→last centroid displacement is a
    // better direction signal (it's where the walker ended up, not just
    // where they peaked).
    bool timed_out = (s.event_frame_count > CV_EVENT_MAX_FRAMES);
    bool is_entry;
    if (timed_out) {
        is_entry = (s.event_last_c < s.event_first_c);
    } else {
        is_entry = (up_score >= down_score);
    }
    if (is_entry) {
        s.entries++;
        result.entries_delta++;
    } else {
        s.exits++;
        result.exits_delta++;
    }
    s.last_fire_frame = s.frame_index;
    result.fire_first_c  = s.event_first_c;
    result.fire_min_c    = s.event_min_c;
    result.fire_max_c    = s.event_max_c;
    result.fire_last_c   = s.event_last_c;
    result.fire_duration = s.event_frame_count;
}

CVResult cv_process(CVState& state, const uint8_t* frame, uint8_t /*line_pct*/) {
    CVResult result = {0, 0, 0, -1, -1, -1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 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];
    frame_diff(frame, state.background, fg, CV_PIXELS);

    // Running-average background blend: bg = (31*bg + frame)/32. Adapts to
    // slow scene drift during idle periods. Frozen during an active event so
    // the walker's signature is never absorbed — otherwise bg retains a
    // "ghost" of the walker for ~30 frames after they leave, keeping fg_count
    // elevated and preventing subsequent walkers from producing a clean
    // trajectory.
    if (!state.event_active) {
        for (int i = 0; i < CV_PIXELS; i++) {
            state.background[i] = (uint8_t)(((uint16_t)state.background[i] * 31 + frame[i]) >> 5);
        }
    }

    int fg_count = 0;
    int min_y = CV_H, max_y = -1;
    long sum_y = 0;
    for (int y = 0; y < CV_H; y++) {
        const uint8_t* row = &fg[y * CV_W];
        int row_count = 0;
        for (int x = 0; x < CV_W; x++) row_count += row[x];
        if (row_count > 0) {
            if (y < min_y) min_y = y;
            if (y > max_y) max_y = y;
            sum_y += (long)row_count * y;
            fg_count += row_count;
        }
    }
    result.fg_count      = fg_count;
    result.fg_min_y      = (fg_count > 0) ? min_y : -1;
    result.fg_max_y      = (fg_count > 0) ? max_y : -1;
    result.fg_centroid_y = (fg_count > 0) ? ((float)sum_y / fg_count) : -1.0f;

    // Hard self-heal: if more than half the frame is fg, bg is catastrophically
    // wrong. Snap and skip the event machine this frame.
    if (fg_count > CV_PIXELS / 2) {
        memcpy(state.background, frame, CV_PIXELS);
        state.last_motion_frame = state.frame_index;
        if (state.event_active) event_reset(state);
        return result;
    }

    // Diagnostic track management (no effect on counting).
    bool motion = fg_count > CV_MIN_BLOB_PX;
    if (motion) {
        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());
        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.missed  = 0;
            state.tracks.push_back(t);
        }
    } else {
        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());
    }

    // Event state machine. Refractory period after a fire blocks new events
    // for CV_EVENT_REFRACTORY_FRAMES frames — absorbs lingering-walker motion
    // that would otherwise re-trigger a second count.
    bool in_refractory = state.last_fire_frame != 0 &&
        (state.frame_index - state.last_fire_frame) < CV_EVENT_REFRACTORY_FRAMES;

    if (!state.event_active) {
        if (!in_refractory && fg_count >= CV_EVENT_ENTER_THRESH) {
            state.event_active      = true;
            state.event_start_frame = state.frame_index;
            state.event_frame_count = 1;
            state.event_peak_n      = fg_count;
            state.event_first_c     = result.fg_centroid_y;
            state.event_last_c      = result.fg_centroid_y;
            state.event_min_c       = result.fg_centroid_y;
            state.event_max_c       = result.fg_centroid_y;
            state.event_min_y_seen  = min_y;
            state.event_max_y_seen  = max_y;
            state.event_quiet_count = 0;
        }
    } else {
        state.event_frame_count++;
        if (fg_count > state.event_peak_n) state.event_peak_n = fg_count;
        if (fg_count > 0) {
            state.event_last_c = result.fg_centroid_y;
            if (result.fg_centroid_y < state.event_min_c) state.event_min_c = result.fg_centroid_y;
            if (result.fg_centroid_y > state.event_max_c) state.event_max_c = result.fg_centroid_y;
            if (min_y < state.event_min_y_seen) state.event_min_y_seen = min_y;
            if (max_y > state.event_max_y_seen) state.event_max_y_seen = max_y;
        }
        if (fg_count < CV_EVENT_EXIT_THRESH) {
            state.event_quiet_count++;
            if (state.event_quiet_count >= CV_EVENT_QUIET_FRAMES) {
                finalize_event(state, result);
                event_reset(state);
                memcpy(state.background, frame, CV_PIXELS);
            }
        } else {
            state.event_quiet_count = 0;
            if (state.event_frame_count > CV_EVENT_MAX_FRAMES) {
                // Timeout end: fg still elevated. Snap bg anyway — in practice
                // a stuck-high event means bg is stale (walker has merged
                // with stale bg, or AEC shifted). Leaving bg stale permanently
                // poisons subsequent events. If a walker truly is mid-frame
                // they'll get absorbed into bg, but that's a rare corner
                // beaten by the common case of stale bg chaining events.
                finalize_event(state, result);
                event_reset(state);
                memcpy(state.background, frame, CV_PIXELS);
            }
        }
    }

    return result;
}