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

#include "JFJochReceiverPlots.h"

void JFJochReceiverPlots::Setup(const DiffractionExperiment &experiment, const AzimuthalIntegration &mapping) {
    {
        std::unique_lock ul(az_int_profile_mutex);
        az_int_profile = std::make_unique<AzimuthalIntegrationProfile>(mapping);
        az_int_profile->SetTitle("dataset");
    }

    default_binning = experiment.GetDefaultPlotBinning();

    // Reset all status vectors
    xfel_pulse_id.Clear();
    xfel_event_code.Clear();
    bkg_estimate.Clear();
    spot_count.Clear();
    indexing_solution.Clear();
    indexing_unit_cell_angle.Clear();
    indexing_unit_cell_len.Clear();
    error_pixels.Clear();
    strong_pixels.Clear();
    receiver_delay.Clear();
    receiver_free_send_buf.Clear();
    image_collection_efficiency.Clear();
    roi_sum.Clear();
    roi_max_count.Clear();
    roi_pixels.Clear();
    roi_x.Clear();
    roi_y.Clear();
    roi_mean.Clear();
    packets_received.Clear();
    max_value.Clear();
    resolution_estimate.Clear();
}

void JFJochReceiverPlots::Add(const DataMessage &msg, const AzimuthalIntegrationProfile &profile) {
    if ( msg.bkg_estimate.has_value())
        bkg_estimate.AddElement(msg.number, msg.bkg_estimate.value());

    if (msg.resolution_estimate.has_value())
        resolution_estimate.AddElement(msg.number, msg.resolution_estimate.value());

    spot_count.AddElement("Spots (crystal)", msg.number, msg.spots.size());
    spot_count.AddElement("Spots (rings)", msg.number, msg.spot_count_in_rings);

    error_pixels.AddElement("Error", msg.number, msg.error_pixel_count);
    error_pixels.AddElement("Saturated", msg.number, msg.saturated_pixel_count);

    strong_pixels.AddElement(msg.number, msg.strong_pixel_count);

    packets_received.AddElement(msg.number, msg.packets_received);
    image_collection_efficiency.AddElement(msg.number, msg.image_collection_efficiency);
    receiver_delay.AddElement(msg.number, msg.receiver_aq_dev_delay);
    receiver_free_send_buf.AddElement(msg.number, msg.receiver_free_send_buf);

    max_value.AddElement(msg.number, msg.max_viable_pixel_value);

    if (msg.indexing_unit_cell) {
        indexing_unit_cell_len.AddElement("a", msg.number, msg.indexing_unit_cell->a);
        indexing_unit_cell_len.AddElement("b", msg.number, msg.indexing_unit_cell->b);
        indexing_unit_cell_len.AddElement("c", msg.number, msg.indexing_unit_cell->c);
        indexing_unit_cell_angle.AddElement("alpha", msg.number, msg.indexing_unit_cell->alpha);
        indexing_unit_cell_angle.AddElement("beta", msg.number, msg.indexing_unit_cell->beta);
        indexing_unit_cell_angle.AddElement("gamma", msg.number, msg.indexing_unit_cell->gamma);
    }

    indexing_solution.AddElement(msg.number, msg.indexing_result);
    {
        std::unique_lock ul(az_int_profile_mutex);
        *az_int_profile += profile;
    }

    xfel_pulse_id.AddElement(msg.number, msg.xfel_pulse_id);
    xfel_event_code.AddElement(msg.number, msg.xfel_event_code);

    for (const auto &[key, value] : msg.roi) {
        roi_sum.AddElement(key, msg.number, value.sum);
        roi_mean.AddElement(key, msg.number, static_cast<double>(value.sum) / static_cast<double>(value.pixels));
        roi_max_count.AddElement(key, msg.number, value.max_count);
        roi_pixels.AddElement(key, msg.number, value.pixels);
        roi_x.AddElement(key, msg.number, static_cast<double>(value.x_weighted) / static_cast<double>(value.sum));
        roi_y.AddElement(key, msg.number, static_cast<double>(value.y_weighted) / static_cast<double>(value.sum));
    }
}

void JFJochReceiverPlots::AddEmptyImage(const DataMessage &msg) {
    image_collection_efficiency.AddElement(msg.number, msg.image_collection_efficiency);
}

MultiLinePlot JFJochReceiverPlots::GetPlots(const PlotRequest &request) {
    MultiLinePlot ret;
    auto nbins = default_binning;
    if (request.binning > 0)
        nbins = request.binning;
    MultiLinePlot tmp;
    switch (request.type) {
        case PlotType::SpotCount:
            return spot_count.GetMeanPlot(nbins);
        case PlotType::IndexingRate:
            return indexing_solution.GetMeanPlot(nbins);
        case PlotType::BkgEstimate:
            return bkg_estimate.GetMeanPlot(nbins);
        case PlotType::ResolutionEstimate:
            return resolution_estimate.GetMeanPlot(nbins);
        case PlotType::ErrorPixels:
            return error_pixels.GetMeanPlot(nbins);
        case PlotType::ImageCollectionEfficiency:
            return image_collection_efficiency.GetMeanPlot(nbins);
        case PlotType::ReceiverDelay:
            return receiver_delay.GetMeanPlot(nbins);
        case PlotType::ReceiverFreeSendBuf:
            return receiver_free_send_buf.GetMeanPlot(nbins);
        case PlotType::StrongPixels:
            return strong_pixels.GetMeanPlot(nbins);
        case PlotType::ROISum:
            return roi_sum.GetMeanPlot(nbins);
        case PlotType::ROIMaxCount:
            return roi_max_count.GetMeanPlot(nbins);
        case PlotType::ROIPixels:
            return roi_pixels.GetMeanPlot(nbins);
        case PlotType::ROIMean:
            return roi_mean.GetMeanPlot(nbins);
        case PlotType::ROIWeightedX:
            return roi_x.GetMeanPlot(nbins);
        case PlotType::ROIWeightedY:
            return roi_y.GetMeanPlot(nbins);
        case PlotType::AzInt:
            return GetAzIntProfilePlot();
        case PlotType::IndexingUnitCellLength:
            return indexing_unit_cell_len.GetMeanPlot(nbins);
        case PlotType::IndexingUnitCellAngle:
            return indexing_unit_cell_angle.GetMeanPlot(nbins);
        case PlotType::PacketsReceived:
            return packets_received.GetMeanPlot(nbins);
        case PlotType::MaxValue:
            return max_value.GetMaxPlot(nbins); // doesn't make sense to give mean here
        default:
            // Do nothing
            return ret;
    }
}

std::optional<float> JFJochReceiverPlots::GetIndexingRate() const {
    auto tmp = indexing_solution.Mean();
    if (std::isfinite(tmp))
        return tmp;
    else
        return {};
}

std::optional<float> JFJochReceiverPlots::GetBkgEstimate() const {
    auto tmp = bkg_estimate.Mean();
    if (std::isfinite(tmp))
        return tmp;
    else
        return {};
}

void JFJochReceiverPlots::GetXFELPulseID(std::vector<uint64_t> &v) const {
    v = xfel_pulse_id.ExportArray(0);
}

void JFJochReceiverPlots::GetXFELEventCode(std::vector<uint64_t> &v) const {
    v = xfel_event_code.ExportArray(0);
}

std::vector<float> JFJochReceiverPlots::GetAzIntProfile() const {
    std::unique_lock ul(az_int_profile_mutex);
    if (!az_int_profile)
        return {};
    return az_int_profile->GetResult();
}

MultiLinePlot JFJochReceiverPlots::GetAzIntProfilePlot() const {
    std::unique_lock ul(az_int_profile_mutex);
    if (!az_int_profile)
        return {};
    return az_int_profile->GetPlot();
}