//
// Created by jungfrau on 2/4/26.
//

#include "RotationSpotAccumulator.h"

RotationSpot3D::RotationSpot3D(SpotToSave s, float phi_deg, int64_t image) {
    x = s.x * s.intensity;
    y = s.y * s.intensity;
    phi = phi_deg * s.intensity;
    intensity = s.intensity;
    maxc = s.maxc;
    first_image = image;
    last_image = image;
}

void RotationSpot3D::Add(const SpotToSave& s, float phi_deg, int64_t image) {
    x += s.x * s.intensity;
    y += s.y * s.intensity;
    phi += phi_deg * s.intensity;
    intensity += s.intensity;
    maxc = std::max(maxc, s.maxc);
    ice_ring = ice_ring || s.ice_ring;
    first_image = std::min(first_image, image);
    last_image = std::max(last_image, image);
}

float RotationSpot3D::calcX() const {
    return x / intensity;
}

float RotationSpot3D::calcY() const {
    return y / intensity;
}

float RotationSpot3D::calcPhi() const {
    return phi / intensity;
}

float RotationSpot3D::calcI() const {
    return intensity;
}

int64_t RotationSpot3D::calcMAXC() const {
    return maxc;
}

int64_t RotationSpot3D::getLastImage() const {
    return last_image;
}

void RotationSpotAccumulator::AddImage(int64_t image, const std::vector<SpotToSave>& in_spots) {
    std::lock_guard<std::mutex> lock(mutex_);
    spots[image] = in_spots;
}


 std::vector<RotationSpot3D> RotationSpotAccumulator::FinalizeAll() {
        std::lock_guard<std::mutex> lock(mutex_);

        int64_t image0 = INT64_MIN;

        std::vector<RotationSpot3D> completed_;
        std::vector<RotationSpot3D> pending_;

        for (auto& [image, v]: spots) {
            float image_angle = axis_.GetIncrement_deg() * image + axis_.GetWedge_deg() / 2.0f;

            std::vector<RotationSpot3D> tmp;
            tmp.reserve(pending_.size() + v.size());

            if (image == image0 + 1) {
                for (const auto &p: pending_) {
                    if (p.getLastImage() != image - 1) {
                        completed_.push_back(p);
                    } else {
                        tmp.push_back(p);
                    }
                }

                // Spatial hash for tmp (pending)
                const float cell_size = config_.grid_cell_size;
                const float inv_cell = (cell_size > 0.0f) ? (1.0f / cell_size) : 1.0f;

                std::unordered_map<int64_t, std::vector<int>> grid;
                grid.reserve(tmp.size() * 2);

                auto cell_key = [](int32_t cx, int32_t cy) -> int64_t {
                    return (static_cast<int64_t>(cx) << 32) ^ static_cast<uint32_t>(cy);
                };

                for (int i = 0; i < static_cast<int>(tmp.size()); ++i) {
                    const int32_t cx = static_cast<int32_t>(std::floor(tmp[i].calcX() * inv_cell));
                    const int32_t cy = static_cast<int32_t>(std::floor(tmp[i].calcY() * inv_cell));
                    grid[cell_key(cx, cy)].push_back(i);
                }

                for (const auto &s: v) {
                    bool matched = false;

                    const int32_t scx = static_cast<int32_t>(std::floor(s.x * inv_cell));
                    const int32_t scy = static_cast<int32_t>(std::floor(s.y * inv_cell));

                    for (int dx = -1; dx <= 1 && !matched; ++dx) {
                        for (int dy = -1; dy <= 1 && !matched; ++dy) {
                            auto it = grid.find(cell_key(scx + dx, scy + dy));
                            if (it == grid.end())
                                continue;

                            for (int idx : it->second) {
                                float dist_x = s.x - tmp[idx].calcX();
                                float dist_y = s.y - tmp[idx].calcY();

                                if (dist_x * dist_x + dist_y * dist_y < config_.xy_tolerance_pxl_sq) {
                                    tmp[idx].Add(s, image_angle, image);
                                    matched = true;
                                    break;
                                }
                            }
                        }
                    }

                    if (!matched)
                        tmp.emplace_back(s, image_angle, image);
                }
            } else {
                for (const auto &p: pending_)
                    completed_.push_back(p);

                for (const auto &s: v)
                    tmp.emplace_back(s, image_angle, image);
            }
            pending_ = std::move(tmp);
            image0 = image;
        }

        for (const auto &p: pending_)
            completed_.push_back(p);

        return completed_;
    }