// Copyright (2019-2023) Paul Scherrer Institute

#include "StrongPixelSet.h"

struct intensity_id {
    int64_t I;
    uint32_t id;
};

struct res_id {
    float res;
    uint32_t id;
};

void FilterSpotsByCount(const DiffractionExperiment& experiment,
                             const std::vector<DiffractionSpot> &input,
                             std::vector<DiffractionSpot> &output) {
    if (input.size() < experiment.GetMaxSpotCount())
        output = input;

    size_t output_size = std::min<size_t>(input.size(), experiment.GetMaxSpotCount());
    std::vector<intensity_id> intensity_id_vector(input.size());


    for (int i = 0; i < input.size(); i++) {
        intensity_id_vector[i].I = input[i].Count();
        intensity_id_vector[i].id = i;
    }

    std::partial_sort(intensity_id_vector.begin(), intensity_id_vector.begin() + output_size, intensity_id_vector.end(),
                      [](const intensity_id& a,const intensity_id& b ) { return a.I > b.I;});

    output.reserve(output_size);
    for (int i = 0; i < output_size; i++)
        output.push_back(input[intensity_id_vector[i].id]);
}

void FilterSpotsByResolution(const DiffractionExperiment& experiment,
                             const std::vector<DiffractionSpot> &input,
                             std::vector<DiffractionSpot> &output) {
    if (input.size() < experiment.GetMaxSpotCount())
        output = input;

    size_t output_size = std::min<size_t>(input.size(), experiment.GetMaxSpotCount());
    std::vector<res_id> res_id_vector(input.size());


    for (int i = 0; i < input.size(); i++) {
        res_id_vector[i].res = input[i].GetResolution(experiment);
        res_id_vector[i].id = i;
    }
    std::partial_sort(res_id_vector.begin(), res_id_vector.begin() + output_size, res_id_vector.end(),
                      [](const res_id& a,const res_id& b ) { return a.res > b.res;});

    output.reserve(output_size);
    for (int i = 0; i < output_size; i++)
        output.push_back(input[res_id_vector[i].id]);
}

StrongPixelSet::StrongPixelSet()
        : strong_pixel_vector(RAW_MODULE_SIZE, false),
          strong_pixel_count(0) {
}

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

void StrongPixelSet::AddSingleStrongPixel(uint16_t col, uint16_t line, int32_t photons) {
    auto key = strong_pixel_coord(col, line);
    auto p = std::pair<uint32_t, int32_t>(key, photons);
    single_pixels.emplace_back(p);
}

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 SpotFindingSettings &settings,
                               std::vector<DiffractionSpot> &spots, uint16_t module_number) {
    if (settings.min_pix_per_spot == 1) {
        for (const auto &i: single_pixels) {
            uint16_t col = col_from_strong_pixel(i.first);
            uint16_t line = line_from_strong_pixel(i.first);

            DiffractionSpot spot(col, line, i.second);
            spots.emplace_back(spot);
        }
    }

    while (!strong_pixel_map.empty()) {
        auto iter = strong_pixel_map.begin();
        DiffractionSpot spot = BuildSpot(iter);
        spot.ConvertToImageCoordinates(experiment, module_number);

        if ((spot.PixelCount() <= settings.max_pix_per_spot)
            && (spot.PixelCount() >= settings.min_pix_per_spot))
            spots.push_back(spot);
    }
}

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 DeviceOutput &output) {
    strong_pixel_count += output.spot_finding_result.strong_pixel_count;

    for (int i = 0; i < std::min<uint32_t>(output.spot_finding_result.max_memory_index * 32,
                                           SPOT_FINDER_MAX_STRONG_PIXEL); i++) {
        uint32_t val = output.spot_finding_result.strong_pixel_number[i];
        if (val != UINT32_MAX) {
            uint32_t pixel = (val & ((1LU << 20) - 1));
            uint8_t conn = ((val >> 24) & UINT8_MAX);
            if (conn == 0)
                AddSingleStrongPixel(pixel % RAW_MODULE_COLS, pixel / RAW_MODULE_COLS, output.pixels[pixel]);
            else
                AddStrongPixel(pixel % RAW_MODULE_COLS, pixel / RAW_MODULE_COLS, output.pixels[pixel]);
        }
    }
}

uint32_t StrongPixelSet::GetStrongPixelCount() const {
    return strong_pixel_count;
}