#include "target_tracker.h"

#include <float.h>
#include <string.h>

#define TRACKER_DEFAULT_POSITION_ALPHA      0.35F
#define TRACKER_DEFAULT_SPEED_ALPHA         0.45F
#define TRACKER_DEFAULT_MOTION_ALPHA        0.60F
#define TRACKER_DEFAULT_ASSOCIATION_MM      1800.0F
#define TRACKER_DEFAULT_MAX_MISSED_FRAMES   6U

typedef struct {
    const target_tracker_t *tracker;
    const rd03d_target_t *detections;
    size_t active_indices[TARGET_TRACKER_CAPACITY];
    size_t active_count;
    size_t detection_count;
    float gate_squared;
    int current_assignment[TARGET_TRACKER_CAPACITY];
    int best_assignment[TARGET_TRACKER_CAPACITY];
    int best_match_count;
    float best_cost;
} association_search_t;

static bool alpha_is_valid(float alpha)
{
    return alpha >= 0.0F && alpha <= 1.0F;
}

void target_tracker_default_config(target_tracker_config_t *config)
{
    if (config == NULL) {
        return;
    }

    config->position_ema_alpha = TRACKER_DEFAULT_POSITION_ALPHA;
    config->speed_ema_alpha = TRACKER_DEFAULT_SPEED_ALPHA;
    config->motion_ema_alpha = TRACKER_DEFAULT_MOTION_ALPHA;
    config->max_association_distance_mm = TRACKER_DEFAULT_ASSOCIATION_MM;
    config->max_missed_frames = TRACKER_DEFAULT_MAX_MISSED_FRAMES;
}

static target_tracker_config_t sanitise_config(const target_tracker_config_t *input)
{
    target_tracker_config_t result;

    target_tracker_default_config(&result);
    if (input == NULL) {
        return result;
    }

    if (alpha_is_valid(input->position_ema_alpha)) {
        result.position_ema_alpha = input->position_ema_alpha;
    }
    if (alpha_is_valid(input->speed_ema_alpha)) {
        result.speed_ema_alpha = input->speed_ema_alpha;
    }
    if (alpha_is_valid(input->motion_ema_alpha)) {
        result.motion_ema_alpha = input->motion_ema_alpha;
    }
    if (input->max_association_distance_mm > 0.0F) {
        result.max_association_distance_mm = input->max_association_distance_mm;
    }
    result.max_missed_frames = input->max_missed_frames;
    return result;
}

void target_tracker_init(target_tracker_t *tracker,
                         const target_tracker_config_t *config)
{
    if (tracker == NULL) {
        return;
    }

    memset(tracker, 0, sizeof(*tracker));
    tracker->config = sanitise_config(config);
    tracker->next_id = 1U;
}

void target_tracker_reset(target_tracker_t *tracker)
{
    target_tracker_config_t config;

    if (tracker == NULL) {
        return;
    }

    config = tracker->config;
    memset(tracker, 0, sizeof(*tracker));
    tracker->config = config;
    tracker->next_id = 1U;
}

static float association_cost(const association_search_t *search,
                              size_t active_offset,
                              size_t detection_index)
{
    const target_track_state_t *track =
        &search->tracker->tracks[search->active_indices[active_offset]];
    const rd03d_target_t *detection = &search->detections[detection_index];
    const float frames_ahead =
        (float)track->snapshot.last_seen_age_frames + 1.0F;
    const float predicted_x = track->last_observed_x_mm +
                              track->velocity_x_mm_per_frame * frames_ahead;
    const float predicted_y = track->last_observed_y_mm +
                              track->velocity_y_mm_per_frame * frames_ahead;
    const float dx = predicted_x - (float)detection->x_mm;
    const float dy = predicted_y - (float)detection->y_mm;

    return dx * dx + dy * dy;
}

static void search_associations(association_search_t *search,
                                size_t active_offset,
                                unsigned int used_detection_mask,
                                int match_count,
                                float total_cost)
{
    size_t detection_index;

    if (active_offset == search->active_count) {
        if (match_count > search->best_match_count ||
            (match_count == search->best_match_count &&
             total_cost < search->best_cost)) {
            search->best_match_count = match_count;
            search->best_cost = total_cost;
            memcpy(search->best_assignment,
                   search->current_assignment,
                   sizeof(search->best_assignment));
        }
        return;
    }

    for (detection_index = 0U;
         detection_index < search->detection_count;
         ++detection_index) {
        const unsigned int detection_bit = 1U << detection_index;
        float cost;

        if ((used_detection_mask & detection_bit) != 0U) {
            continue;
        }
        cost = association_cost(search, active_offset, detection_index);
        if (cost > search->gate_squared) {
            continue;
        }

        search->current_assignment[active_offset] = (int)detection_index;
        search_associations(search,
                            active_offset + 1U,
                            used_detection_mask | detection_bit,
                            match_count + 1,
                            total_cost + cost);
    }

    search->current_assignment[active_offset] = -1;
    search_associations(search,
                        active_offset + 1U,
                        used_detection_mask,
                        match_count,
                        total_cost);
}

static uint32_t allocate_id(target_tracker_t *tracker)
{
    uint32_t id = tracker->next_id++;

    if (id == 0U) {
        id = tracker->next_id++;
    }
    if (tracker->next_id == 0U) {
        tracker->next_id = 1U;
    }
    return id;
}

static void initialise_track(target_tracker_t *tracker,
                             target_track_state_t *track,
                             const rd03d_target_t *detection)
{
    memset(track, 0, sizeof(*track));
    track->active = true;
    track->last_observed_x_mm = (float)detection->x_mm;
    track->last_observed_y_mm = (float)detection->y_mm;
    track->snapshot.id = allocate_id(tracker);
    track->snapshot.x_mm = (float)detection->x_mm;
    track->snapshot.y_mm = (float)detection->y_mm;
    track->snapshot.speed_cm_s = (float)detection->speed_cm_s;
    track->snapshot.resolution_mm = detection->resolution_mm;
    track->snapshot.age_frames = 1U;
    track->snapshot.last_seen_age_frames = 0U;
    track->snapshot.visible = true;
    track->snapshot.observed_this_frame = true;
}

static void update_matched_track(target_track_state_t *track,
                                 const rd03d_target_t *detection,
                                 const target_tracker_config_t *config)
{
    const float gap = (float)track->snapshot.last_seen_age_frames + 1.0F;
    const float measured_velocity_x =
        ((float)detection->x_mm - track->last_observed_x_mm) / gap;
    const float measured_velocity_y =
        ((float)detection->y_mm - track->last_observed_y_mm) / gap;
    const float position_keep = 1.0F - config->position_ema_alpha;
    const float speed_keep = 1.0F - config->speed_ema_alpha;
    const float motion_keep = 1.0F - config->motion_ema_alpha;

    track->velocity_x_mm_per_frame =
        config->motion_ema_alpha * measured_velocity_x +
        motion_keep * track->velocity_x_mm_per_frame;
    track->velocity_y_mm_per_frame =
        config->motion_ema_alpha * measured_velocity_y +
        motion_keep * track->velocity_y_mm_per_frame;
    track->last_observed_x_mm = (float)detection->x_mm;
    track->last_observed_y_mm = (float)detection->y_mm;

    track->snapshot.x_mm = config->position_ema_alpha * (float)detection->x_mm +
                           position_keep * track->snapshot.x_mm;
    track->snapshot.y_mm = config->position_ema_alpha * (float)detection->y_mm +
                           position_keep * track->snapshot.y_mm;
    track->snapshot.speed_cm_s =
        config->speed_ema_alpha * (float)detection->speed_cm_s +
        speed_keep * track->snapshot.speed_cm_s;
    track->snapshot.resolution_mm = detection->resolution_mm;
    track->snapshot.last_seen_age_frames = 0U;
    track->snapshot.observed_this_frame = true;
}

size_t target_tracker_update(target_tracker_t *tracker,
                             const rd03d_target_t *detections,
                             size_t detection_count)
{
    rd03d_target_t valid_detections[TARGET_TRACKER_CAPACITY];
    bool detection_used[TARGET_TRACKER_CAPACITY] = {false, false, false};
    association_search_t search;
    size_t valid_count = 0U;
    size_t track_index;
    size_t active_offset;
    size_t detection_index;

    if (tracker == NULL || (detections == NULL && detection_count != 0U)) {
        return 0U;
    }

    for (detection_index = 0U;
         detection_index < detection_count && valid_count < TARGET_TRACKER_CAPACITY;
         ++detection_index) {
        if (detections[detection_index].valid) {
            valid_detections[valid_count++] = detections[detection_index];
        }
    }

    memset(&search, 0, sizeof(search));
    search.tracker = tracker;
    search.detections = valid_detections;
    search.detection_count = valid_count;
    search.gate_squared = tracker->config.max_association_distance_mm *
                          tracker->config.max_association_distance_mm;
    search.best_match_count = -1;
    search.best_cost = FLT_MAX;

    for (track_index = 0U; track_index < TARGET_TRACKER_CAPACITY; ++track_index) {
        target_track_state_t *track = &tracker->tracks[track_index];

        if (!track->active) {
            continue;
        }
        track->snapshot.observed_this_frame = false;
        if (track->snapshot.age_frames != UINT32_MAX) {
            ++track->snapshot.age_frames;
        }
        search.active_indices[search.active_count++] = track_index;
    }

    search_associations(&search, 0U, 0U, 0, 0.0F);

    for (active_offset = 0U;
         active_offset < search.active_count;
         ++active_offset) {
        target_track_state_t *track =
            &tracker->tracks[search.active_indices[active_offset]];
        const int assigned_detection = search.best_assignment[active_offset];

        if (assigned_detection >= 0) {
            detection_used[(size_t)assigned_detection] = true;
            update_matched_track(track,
                                 &valid_detections[(size_t)assigned_detection],
                                 &tracker->config);
            continue;
        }

        if (track->snapshot.last_seen_age_frames != UINT16_MAX) {
            ++track->snapshot.last_seen_age_frames;
        }
        if (track->snapshot.last_seen_age_frames >
            tracker->config.max_missed_frames) {
            memset(track, 0, sizeof(*track));
        }
    }

    for (detection_index = 0U; detection_index < valid_count; ++detection_index) {
        if (detection_used[detection_index]) {
            continue;
        }
        target_track_state_t *destination = NULL;
        for (track_index = 0U;
             track_index < TARGET_TRACKER_CAPACITY;
             ++track_index) {
            if (!tracker->tracks[track_index].active) {
                destination = &tracker->tracks[track_index];
                break;
            }
        }
        if (destination == NULL) {
            /* Current detections take precedence over retained missing tracks. */
            for (track_index = 0U;
                 track_index < TARGET_TRACKER_CAPACITY;
                 ++track_index) {
                target_track_state_t *candidate = &tracker->tracks[track_index];
                if (candidate->snapshot.observed_this_frame) {
                    continue;
                }
                if (destination == NULL ||
                    candidate->snapshot.last_seen_age_frames >
                        destination->snapshot.last_seen_age_frames) {
                    destination = candidate;
                }
            }
        }
        if (destination != NULL) {
            initialise_track(tracker, destination,
                             &valid_detections[detection_index]);
        }
    }

    return target_tracker_snapshot(tracker, NULL, 0U);
}

size_t target_tracker_update_frame(target_tracker_t *tracker,
                                   const rd03d_frame_t *frame)
{
    if (frame == NULL) {
        return target_tracker_update(tracker, NULL, 0U);
    }
    return target_tracker_update(tracker, frame->targets, RD03D_MAX_TARGETS);
}

size_t target_tracker_snapshot(const target_tracker_t *tracker,
                               target_snapshot_t *out,
                               size_t capacity)
{
    const target_snapshot_t *ordered[TARGET_TRACKER_CAPACITY];
    size_t count = 0U;
    size_t track_index;
    size_t index;

    if (tracker == NULL) {
        return 0U;
    }

    for (track_index = 0U; track_index < TARGET_TRACKER_CAPACITY; ++track_index) {
        if (tracker->tracks[track_index].active) {
            size_t insertion = count;

            while (insertion > 0U &&
                   ordered[insertion - 1U]->id >
                       tracker->tracks[track_index].snapshot.id) {
                ordered[insertion] = ordered[insertion - 1U];
                --insertion;
            }
            ordered[insertion] = &tracker->tracks[track_index].snapshot;
            ++count;
        }
    }

    if (out == NULL) {
        return count;
    }
    if (capacity > count) {
        capacity = count;
    }
    for (index = 0U; index < capacity; ++index) {
        out[index] = *ordered[index];
    }
    return count;
}