// Copyright (2019-2023) Paul Scherrer Institute

#include <bitset>
#include "StrongPixelSet.h"

StrongPixelSet::StrongPixelSet(const DiffractionExperiment &experiment)
        : xpixel(experiment.GetXPixelsNum()),
          ypixel(experiment.GetYPixelsNum()),
          strong_pixel_vector(experiment.GetPixelsNum(), false) {
}

void StrongPixelSet::AddStrongPixel(uint16_t col, uint16_t line, int32_t photons) {
    auto key = strong_pixel_coord(col, line);
    strong_pixel_map[key] = photons;
    strong_pixel_vector.at(line * xpixel + col) = true;
}

inline void StrongPixelSet::AddNeighbor(DiffractionSpot &spot, uint16_t col, uint16_t line) {
    if (strong_pixel_vector[line * xpixel + col]) {
        uint64_t coord = strong_pixel_coord(col, line);
        auto iter = strong_pixel_map.find(coord);
        ExtendSpot(spot, iter);
    }
}

// Creates a continuous spot
// strong pixels are loaded into dictionary (one dictionary per frame)
// and routine checks if neighboring pixels are also in dictionary (likely in log(N) time)
DiffractionSpot StrongPixelSet::BuildSpot(std::unordered_map<uint32_t, int32_t>::iterator &it) {

    uint16_t col = col_from_strong_pixel(it->first);
    uint16_t line = line_from_strong_pixel(it->first);

    DiffractionSpot spot(col, line, it->second);

    strong_pixel_vector[line * xpixel + col] = false;
    strong_pixel_map.erase(it); // Remove strong pixel from the dictionary, so it is not processed again

    if (col+1 < xpixel) {
        AddNeighbor(spot, col + 1, line);
        if (line < ypixel - 1)
            AddNeighbor(spot, col + 1, line + 1);
        if (line > 0)
            AddNeighbor(spot, col + 1, line - 1);
    }
    if (col != 0) {
        AddNeighbor(spot, col - 1, line);
        if (line < ypixel - 1)
            AddNeighbor(spot, col - 1, line + 1);
        if (line > 0)
            AddNeighbor(spot, col - 1, line - 1);
    }
    if (line < ypixel - 1)
        AddNeighbor(spot, col, line+1);
    if (line > 0)
        AddNeighbor(spot, col, line-1);

    return spot;
}

void StrongPixelSet::ExtendSpot(DiffractionSpot &spot, std::unordered_map<uint32_t, int32_t>::iterator &it) {

    uint16_t col = col_from_strong_pixel(it->first);
    uint16_t line = line_from_strong_pixel(it->first);

    spot.AddPixel(col, line, it->second);

    strong_pixel_vector[line * xpixel + col] = false;
    strong_pixel_map.erase(it); // Remove strong pixel from the dictionary, so it is not processed again

    if (col+1 < xpixel) {
        AddNeighbor(spot, col + 1, line);
        if (line < ypixel - 1)
            AddNeighbor(spot, col + 1, line + 1);
        if (line > 0)
            AddNeighbor(spot, col + 1, line - 1);
    }
    if (col != 0) {
        AddNeighbor(spot, col - 1, line);
        if (line < ypixel - 1)
            AddNeighbor(spot, col - 1, line + 1);
        if (line > 0)
            AddNeighbor(spot, col - 1, line - 1);
    }
    if (line < ypixel - 1)
        AddNeighbor(spot, col, line+1);
    if (line > 0)
        AddNeighbor(spot, col, line-1);
}

void StrongPixelSet::FindSpots(const DiffractionExperiment &experiment, const JFJochProtoBuf::DataProcessingSettings &settings,
                               std::vector<DiffractionSpot> &spots) {
    std::multimap<double, DiffractionSpot> spots_map;

    while (!strong_pixel_map.empty()) {
        auto iter = strong_pixel_map.begin();
        DiffractionSpot spot = BuildSpot(iter);
        double d = spot.GetResolution(experiment);

        if ((spot.PixelCount() <= settings.max_pix_per_spot())
            && (spot.PixelCount() >= settings.min_pix_per_spot())
            && (!settings.has_low_resolution_limit() || (d <= settings.low_resolution_limit()))
            && (!settings.has_high_resolution_limit() || (d >= settings.high_resolution_limit())))
            spots_map.insert(std::make_pair(-static_cast<float>(d), spot));
    }

    for (auto &[x, spot]: spots_map)
        spots.push_back(spot);

    if (experiment.GetMaxSpotCount() > 0)
        spots.resize(std::min<size_t>(spots.size(), experiment.GetMaxSpotCount()));
}

size_t StrongPixelSet::Count() const {
    return strong_pixel_map.size();
}

size_t StrongPixelSet::Common(const StrongPixelSet &set) const {
    size_t ret = 0;
    for (const auto& pixel: strong_pixel_map) {
        if (set.strong_pixel_map.count(pixel.first) == 1)
            ret++;
    }
    return ret;
}