// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

// 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
//
// The union-find and the two-scan structure are theirs. How a pixel's earlier neighbours are FOUND
// is not: see sparseccl below.

#include <bitset>

#include "StrongPixelSet.h"

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});
    ++strong_pixel_count;
}

uint32_t StrongPixelSet::find_root(uint32_t e) {
    uint32_t r = e;
    while (L[r] != r)
        r = L[r];
    return r;
}

uint32_t StrongPixelSet::make_union(uint32_t e1, uint32_t e2) {
    uint32_t e;
    if (e1 < e2) {
        e = e1;
        L[e2] = e;
    } else {
        e = e2;
        L[e1] = e;
    }
    return e;
}

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

    unsigned int labels = 0;

    // First scan: pixel association. The pixels arrive in raster order - line ascending, column
    // ascending within a line - which upstream uses to walk a sliding window of the last two lines,
    // testing every pixel in it for adjacency. That is quadratic in how many strong pixels a line
    // pair holds: fine for the silicon-tracker hits it was written for, but a flooded detector line
    // holds four thousand of them, and labelling a fully lit frame took 76 seconds.
    //
    // Since the columns ascend, the window need not be walked. A pixel's earlier 8-neighbours are
    // exactly the one to its left and the at most three above it, so keep the previous line's range
    // and a cursor into it that only ever moves forward - the same four neighbours the GPU extractor
    // finds by binary search. Same edge set, same unions in the same order, therefore the same
    // labels; the flooded frame now takes 0.16 s.
    uint32_t line_begin = 0;                  // first pixel of the line being scanned
    uint32_t prev_begin = 0, prev_end = 0;    // the pixels of the line above it
    uint32_t up = 0;                          // cursor into [prev_begin, prev_end)
    for (uint32_t i = 0; i < pixels.size(); ++i) {
        L[i] = i;
        if (i > 0 && pixels[i].line != pixels[i - 1].line) {
            // The line above is the previous one only if it really is the line above: a line with no
            // strong pixel at all leaves nothing to join to.
            prev_begin = (pixels[i].line == pixels[i - 1].line + 1) ? line_begin : i;
            prev_end = i;
            line_begin = i;
            up = prev_begin;
        }
        uint32_t ai = i;
        while (up < prev_end && pixels[up].col + 1 < pixels[i].col)
            ++up;
        for (uint32_t j = up; j < prev_end && pixels[j].col <= pixels[i].col + 1; ++j)
            ai = make_union(ai, find_root(j));
        // The pixel to the left comes last, as it did when the window was walked in order.
        if (i > line_begin && pixels[i - 1].col + 1 == pixels[i].col)
            ai = make_union(ai, find_root(i - 1));
    }

    // second scan: transitive closure
    for (uint32_t 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 (uint32_t i = 0; i < L.size(); i++)
        spots[L[i]].AddPixel(pixels[i].col, pixels[i].line, pixels[i].counts);

    return spots;
}


void StrongPixelSet::FindComponentsImage(const SpotFindingSettings &settings, std::vector<DiffractionSpot> &spots) {
    // No StrongPixelLimit test here: the caller knows how big the image is and has already applied it.
    for (const auto &spot: sparseccl()) {
        if (spot.PixelCount() <= settings.max_pix_per_spot)
            spots.push_back(spot);
    }
}

void StrongPixelSet::FindSpots(const DiffractionExperiment &experiment, const SpotFindingSettings &settings,
                               std::vector<DiffractionSpot> &spots, uint16_t module_number) {
    // Per module, so the bar is the module's own - and ReadFPGAOutput has already refused anything
    // past max_strong_pixel_per_module, far below it.
    if (!pixels.empty() && (strong_pixel_count < StrongPixelLimit(RAW_MODULE_SIZE))) {
        for (const auto &spot: sparseccl()) {
            if ((spot.PixelCount() <= settings.max_pix_per_spot)
                && (spot.PixelCount() >= settings.min_pix_per_spot.value_or(2))) {
                auto s = spot;
                s.ConvertToImageCoordinates(experiment, module_number);
                spots.push_back(s);
            }
        }
    }
}

void StrongPixelSet::ReadFPGAOutput(const DiffractionExperiment & experiment,
                                    const DeviceOutput &output) {
    // 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)) {
        // If max strong pixel per module condition kicks-in, still report correct strong pixel count
        strong_pixel_count = output.spot_finding_result.strong_pixel_count;
        return;
    }

    auto pixel_depth = experiment.GetByteDepthImage();
    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++) {
        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;
                    if (pixel_depth == 2)
                        AddStrongPixel(col, line, output.pixels[npixel | j]);
                    else if (pixel_depth == 1)
                        AddStrongPixel(col, line, ((int8_t *)output.pixels)[npixel | j]);
                    else if (pixel_depth == 4)
                        AddStrongPixel(col, line, ((int32_t *)output.pixels)[npixel | j]);
                }
            }
        }
    }

}

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