// 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

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

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::FindSpotsImage(const SpotFindingSettings &settings, std::vector<DiffractionSpot> &spots) {
    if (!pixels.empty()) {
        for (const auto &spot: sparseccl()) {
            if ((spot.PixelCount() <= settings.max_pix_per_spot)
                && (spot.PixelCount() >= settings.min_pix_per_spot)) {
                spots.push_back(spot);
            }
        }
    }
}

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