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

#include <algorithm>
#include <cmath>

#include "ScanResultGenerator.h"

namespace {
    template<class T>
    T value_or_zero(const std::optional<T>& v) {
        return v.value_or(static_cast<T>(0));
    }
}

ScanResultGenerator::ScanResultGenerator(const DiffractionExperiment &experiment) {
    grid_scan = experiment.GetGridScan();
    goniometer_axis = experiment.GetGoniometer();
    if (grid_scan)
        v.resize(grid_scan->GetNElem());
    else
        v.resize(experiment.GetImageNum());
    file_prefix = experiment.GetFilePrefix();
}

void ScanResultGenerator::Add(const DataMessage &message) {
    std::unique_lock ul(m);

    int64_t image_number = message.number;

    if (grid_scan)
        image_number = grid_scan->Rearrange(image_number);

    if (image_number >= 0 && static_cast<size_t>(image_number) < v.size()) {
        if (grid_scan) {
            v[image_number].x = grid_scan->GetElementPosX_step(message.number);
            v[image_number].y = grid_scan->GetElementPosY_step(message.number);
        } else if (goniometer_axis) {
            v[image_number].angle_deg = goniometer_axis->GetAngle_deg(message.number);
        }

        v[image_number].number = message.number;
        v[image_number].pixel_sum = message.pixel_sum;
        v[image_number].collection_efficiency = message.image_collection_efficiency.value_or(1.0);
        v[image_number].bkg = message.bkg_estimate;
        v[image_number].spot_count = message.spot_count;
        v[image_number].indexing_solution = message.indexing_result;
        v[image_number].profile_radius = message.profile_radius;
        v[image_number].mosaicity = message.mosaicity_deg;
        v[image_number].b_factor = message.b_factor;
        v[image_number].uc = message.indexing_unit_cell;
        v[image_number].xfel_pulse_id = message.xfel_pulse_id;
        v[image_number].err_pixels = message.error_pixel_count;
        v[image_number].min_viable_pixel = message.min_viable_pixel_value;
        v[image_number].max_viable_pixel = message.max_viable_pixel_value;
        v[image_number].sat_pixels = message.saturated_pixel_count;
        v[image_number].spot_count_ice = message.spot_count_ice_rings;
        v[image_number].spot_count_low_res = message.spot_count_low_res;
        v[image_number].spot_count_indexed = message.spot_count_indexed;
        v[image_number].res = message.resolution_estimate;
        v[image_number].integrated_reflections = message.integrated_reflections;

        if (message.lattice_type)
            v[image_number].niggli_class = message.lattice_type->niggli_class;
    }
}

ScanResult ScanResultGenerator::GetResult() const {
    std::unique_lock ul(m);
    ScanResult ret;
    ret.file_prefix = file_prefix;
    ret.images = v;
    return ret;
}

void ScanResultGenerator::FillEndMessage(EndMessage &message) const {
 std::unique_lock ul(m);

    size_t n = 0;
    for (const auto &e: v) {
        if (e.number >= 0)
            n = std::max(n, static_cast<size_t>(e.number) + 1);
    }
    if (n == 0)
        return;

    message.data_collection_efficiency.resize(n);
    message.spot_count.resize(n);
    message.spot_count_ice_ring.resize(n);
    message.spot_count_low_res.resize(n);
    message.spot_count_indexed.resize(n);
    message.image_indexed.resize(n);
    message.v_bkg_estimate.resize(n);
    message.profile_radius.resize(n);
    message.mosaicity.resize(n);
    message.bFactor.resize(n);
    message.resolution_estimate.resize(n);
    message.min_viable_pixel_value.resize(n);
    message.max_viable_pixel_value.resize(n);
    message.saturated_pixel_count.resize(n);
    message.error_pixel_count.resize(n);
    message.image_scale_factor.resize(n);
    message.integrated_reflections.resize(n);
    message.niggli_class.resize(n);
    message.pixel_sum.resize(n);

    for (const auto &e: v) {
        if (e.number < 0)
            continue;

        const auto number = static_cast<size_t>(e.number);
        if (number >= n)
            continue;

        message.data_collection_efficiency[number] = e.collection_efficiency;
        message.spot_count[number] = static_cast<int32_t>(value_or_zero(e.spot_count));
        message.spot_count_ice_ring[number] = static_cast<int32_t>(value_or_zero(e.spot_count_ice));
        message.spot_count_low_res[number] = static_cast<int32_t>(value_or_zero(e.spot_count_low_res));
        message.spot_count_indexed[number] = static_cast<int32_t>(value_or_zero(e.spot_count_indexed));
        message.image_indexed[number] = static_cast<uint8_t>(e.indexing_solution.value_or(0));
        message.v_bkg_estimate[number] = e.bkg.value_or(NAN);
        message.profile_radius[number] = e.profile_radius.value_or(NAN);
        message.mosaicity[number] = e.mosaicity.value_or(NAN);
        message.bFactor[number] = e.b_factor.value_or(NAN);
        message.resolution_estimate[number] = e.res.value_or(NAN);
        message.min_viable_pixel_value[number] = value_or_zero(e.min_viable_pixel);
        message.max_viable_pixel_value[number] = value_or_zero(e.max_viable_pixel);
        message.saturated_pixel_count[number] = static_cast<int32_t>(value_or_zero(e.sat_pixels));
        message.error_pixel_count[number] = static_cast<int32_t>(value_or_zero(e.err_pixels));
        message.image_scale_factor[number] = e.image_scale_factor.value_or(NAN);
        message.integrated_reflections[number] = static_cast<int32_t>(value_or_zero(e.integrated_reflections));
        message.niggli_class[number] = static_cast<uint8_t>(value_or_zero(e.niggli_class));
        message.pixel_sum[number] = static_cast<int64_t>(value_or_zero(e.pixel_sum));
    }
}