// Copyright (2019-2024) Paul Scherrer Institute

// SparseCCL code taken from https://github.com/acts-project/traccc/blob/main/core/include/traccc/clusterization/detail/sparse_ccl.hpp
// (c) 2021-2022 CERN for the benefit of the ACTS project
// Mozilla Public License Version 2.0

#include <bitset>

#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) {

    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_count(0) {
    pixels.reserve(max_strong_pixel_per_module);
}

void StrongPixelSet::AddStrongPixel(uint16_t col, uint16_t line, int32_t photons) {
    pixels.push_back(strong_pixel{.col = col, .line = line, .counts = photons});
}

bool is_far_enough(strong_pixel pixel0, strong_pixel pixel1) {
    return (pixel1.line - pixel0.line) > 1;
}

bool is_adjacent(strong_pixel pixel0, strong_pixel pixel1) {
    return (fabs(pixel0.line - pixel1.line) <= 1) and
            (fabs(pixel0.col - pixel1.col) <= 1);
}

uint16_t StrongPixelSet::find_root(uint16_t e) {
    uint16_t r = e;
    //assert(r < L.size());
    while (L[r] != r) {
        r = L[r];
        //assert(r < L.size());
    }
    return r;
}

uint16_t StrongPixelSet::make_union(uint16_t e1, uint16_t e2) {
    uint16_t e;
    if (e1 < e2) {
        e = e1;
        //assert(e2 < L.size());
        L[e2] = e;
    } else {
        e = e2;
        //assert(e1 < L.size());
        L[e1] = e;
    }
    return e;
}

std::vector<DiffractionSpot> StrongPixelSet::sparseccl() {
    L.resize(pixels.size());

    unsigned int labels = 0;

    // first scan: pixel association
    uint16_t start_j = 0;
    for (uint16_t i = 0; i < pixels.size(); ++i) {
        L[i] = i;
        uint16_t ai = i;
        for (uint16_t j = start_j; j < i; ++j) {
            if (is_adjacent(pixels[i], pixels[j])) {
                ai = make_union(ai, find_root(j));
            } else if (is_far_enough(pixels[j], pixels[i])) {
                ++start_j;
            }
        }
    }

    // second scan: transitive closure
    for (unsigned int i = 0; i < L.size(); ++i) {
        if (L[i] == i) {
            L[i] = labels++;
        } else {
            L[i] = L[L[i]];
        }
    }

    std::vector<DiffractionSpot> spots(labels);

    for (unsigned int i = 0; i < L.size(); i++)
        spots[L[i]].AddPixel(pixels[i].col, pixels[i].line, pixels[i].counts);

    return spots;
}


void StrongPixelSet::FindSpots(const DiffractionExperiment &experiment, const SpotFindingSettings &settings,
                               std::vector<DiffractionSpot> &spots, uint16_t module_number) {
    if (!pixels.empty()) {
        for (const auto &spot: sparseccl()) {
            if ((spot.PixelCount() <= settings.max_pix_per_spot)
                && (spot.PixelCount() >= settings.min_pix_per_spot)) {
                auto s = spot;
                s.ConvertToImageCoordinates(experiment, module_number);
                spots.push_back(s);
            }
        }
    }
}

void StrongPixelSet::ReadFPGAOutput(const DiffractionExperiment & experiment,
                                    const DeviceOutput &output) {
    strong_pixel_count += output.spot_finding_result.strong_pixel_count;

    // Too many strong pixels will kill performance in data processing, so protection is needed
    // Also if there are no strong pixels, there is no point in looking for them
    if ((output.spot_finding_result.strong_pixel_count == 0) ||
        (output.spot_finding_result.strong_pixel_count > max_strong_pixel_per_module))
        return;

    auto pixel_depth = experiment.GetPixelDepth();
    auto out_ptr = (uint32_t *) output.spot_finding_result.strong_pixel;

    if (pixel_depth == 2) {
        for (int i = 0; i < RAW_MODULE_SIZE / (8 * sizeof(out_ptr[0])); i++) {
            size_t npixel = i * 8 * sizeof(out_ptr[0]);
            size_t line = npixel / RAW_MODULE_COLS;

            if (out_ptr[i] != 0) {
                std::bitset<32> bitset(out_ptr[i]);

                for (int j = 0; j < 32; j++) {
                    if (bitset.test(j)) {
                        size_t col = (npixel | j) % RAW_MODULE_COLS;
                        AddStrongPixel(col, line, output.pixels[npixel | j]);
                    }
                }
            }
        }
    } else {
        for (int i = 0; i < RAW_MODULE_SIZE / (8 * sizeof(out_ptr[0])); i++) {
            size_t npixel = i * 8 * sizeof(out_ptr[0]);
            size_t line = npixel / RAW_MODULE_COLS;

            if (out_ptr[i] != 0) {
                std::bitset<32> bitset(out_ptr[i]);

                for (int j = 0; j < 32; j++) {
                    if (bitset.test(j)) {
                        size_t col = (npixel | j) % RAW_MODULE_COLS;
                        AddStrongPixel(col, line, ((int32_t *) output.pixels)[npixel | j]);
                    }
                }
            }
        }
    }
}

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