// Copyright (2019-2023) Paul Scherrer Institute

#include <bitset>
#include "StrongPixelSet.h"
#include "../common/RawToConvertedGeometry.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 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.low_resolution_limit < 0) || (d <= settings.low_resolution_limit))
            && ((settings.high_resolution_limit < 0) || (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;
}

void StrongPixelSet::ReadFPGAOutput(const DiffractionExperiment& experiment,
                                    const DeviceOutput &output,
                                    uint16_t module) {
    auto out_ptr = (uint32_t *) output.spot_finding_result.strong_pixel;
    for (int i = 0; i < RAW_MODULE_SIZE / (8 * sizeof(out_ptr[0])); i++) {
        if (out_ptr[i]) {
            for (int j = 0; j < 8 * sizeof(out_ptr[0]); j++) {
                if (out_ptr[i] & (1 << j)) {
                    size_t npixel = i * 8 * sizeof(out_ptr[0])| j;
                    auto [col, line] = RawToConvertedCoordinate(experiment, module, npixel);
                    AddStrongPixel(col, line);
                }
            }
        }
    }
}