// 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(m);

    az_int_profile = std::make_unique<AzimuthalIntegrationProfile>(mapping);
    az_int_profile->SetTitle("dataset");

    goniometer = experiment.GetGoniometer();
    grid_scan = experiment.GetGridScan();
    default_binning = experiment.GetDefaultPlotBinning();

    // Reset all status vectors
    xfel_pulse_id.Clear();
    xfel_event_code.Clear();
    bkg_estimate.Clear();
    spot_count.Clear();
    spot_count_low_res.Clear();
    spot_count_indexed.Clear();
    spot_count_ice.Clear();

    indexing_solution.Clear();
    indexing_unit_cell_angle.Clear();
    indexing_unit_cell_len.Clear();
    error_pixels.Clear();
    saturated_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();
    indexing_time.Clear();
    profile_radius.Clear();
    b_factor.Clear();
    processing_time.Clear();
    beam_center_x.Clear();
    beam_center_y.Clear();
    pixel_sum.Clear();
}

void JFJochReceiverPlots::Add(const DataMessage &msg, const AzimuthalIntegrationProfile &profile) {
    bkg_estimate.AddElement(msg.number, msg.bkg_estimate);
    resolution_estimate.AddElement(msg.number, msg.resolution_estimate);
    spot_count.AddElement(msg.number, msg.spot_count);
    spot_count_low_res.AddElement(msg.number, msg.spot_count_low_res);
    spot_count_indexed.AddElement(msg.number, msg.spot_count_indexed);
    spot_count_ice.AddElement(msg.number, msg.spot_count_ice_rings);
    error_pixels.AddElement(msg.number, msg.error_pixel_count);
    saturated_pixels.AddElement(msg.number, msg.saturated_pixel_count);
    pixel_sum.AddElement(msg.number, msg.pixel_sum);
    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);
    indexing_time.AddElement(msg.number, msg.indexing_time_s);
    processing_time.AddElement(msg.number, msg.processing_time_s);

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

    beam_center_x.AddElement(msg.number, msg.beam_corr_x);
    beam_center_y.AddElement(msg.number, msg.beam_corr_y);

    profile_radius.AddElement(msg.number, msg.profile_radius);
    b_factor.AddElement(msg.number, msg.b_factor);
    indexing_solution.AddElement(msg.number, msg.indexing_result);

    {
        std::unique_lock ul(m);
        *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;
    MultiLinePlotUnits units = MultiLinePlotUnits::ImageNumber;
    int64_t nbins;

    std::optional<GridScanSettings> local_grid_scan;
    float start = 0.0;
    float incr = 1.0;

    if (request.type != PlotType::AzInt && request.type != PlotType::AzInt1D) {
        std::unique_lock ul(m);

        if (request.experimental_coord && grid_scan) {
            local_grid_scan = grid_scan;
            units = MultiLinePlotUnits::Grid_um;
            nbins = 1;
        } else {
            nbins = default_binning;
            if (request.binning > 0)
                nbins = request.binning;

            if (request.experimental_coord && goniometer) {
                start = goniometer->GetStart_deg();
                incr = goniometer->GetIncrement_deg();
                units = MultiLinePlotUnits::Angle_deg;
            }
        }
    } else {
        switch (request.azint_unit) {
            case PlotAzintUnit::Q_recipA:
                units = MultiLinePlotUnits::Q_recipA;
                break;
            case PlotAzintUnit::TwoTheta_deg:
                units = MultiLinePlotUnits::Angle_deg;
                break;
            case PlotAzintUnit::D_A:
                units = MultiLinePlotUnits::d_A;
                break;
        }
    }

    switch (request.type) {
        case PlotType::SpotCount:
            ret = spot_count.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::SpotCountLowRes:
            ret = spot_count_low_res.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::SpotCountIndexed:
            ret = spot_count_indexed.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::SpotCountIceRing:
            ret = spot_count_ice.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::IndexingRate:
            ret = indexing_solution.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::BkgEstimate:
            ret = bkg_estimate.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::ResolutionEstimate:
            ret = resolution_estimate.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::ErrorPixels:
            ret = error_pixels.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::SaturatedPixels:
            ret = saturated_pixels.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::IndexingTime:
            ret = indexing_time.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::ProfileRadius:
            ret = profile_radius.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::BFactor:
            ret = b_factor.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::ImageCollectionEfficiency:
            ret = image_collection_efficiency.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::ReceiverDelay:
            ret = receiver_delay.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::ReceiverFreeSendBuf:
            ret = receiver_free_send_buf.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::StrongPixels:
            ret = strong_pixels.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::ROISum:
            ret = roi_sum.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::ROIMaxCount:
            ret = roi_max_count.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::ROIPixels:
            ret = roi_pixels.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::ROIMean:
            ret = roi_mean.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::ROIWeightedX:
            ret = roi_x.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::ROIWeightedY:
            ret = roi_y.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::AzInt:
            ret = GetAzIntProfilePlot(false, request.azint_unit);
            break;
        case PlotType::AzInt1D:
            ret = GetAzIntProfilePlot(true, request.azint_unit);
            break;
        case PlotType::IndexingUnitCellLength:
            ret = indexing_unit_cell_len.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::IndexingUnitCellAngle:
            ret = indexing_unit_cell_angle.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::PacketsReceived:
            ret = packets_received.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::MaxValue:
            ret = max_value.GetMaxPlot(nbins, start, incr, request.fill_value);
            break; // doesn't make sense to give mean here
        case PlotType::PixelSum:
            ret = pixel_sum.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::ImageProcessingTime:
            ret = processing_time.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::RefinementBeamX:
            ret = beam_center_x.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        case PlotType::RefinementBeamY:
            ret = beam_center_y.GetMeanPlot(nbins, start, incr, request.fill_value);
            break;
        default:
            // Do nothing
            break;
    }

    ret.SetUnits(units);
    if (local_grid_scan
        && request.type != PlotType::AzInt
        && request.type != PlotType::AzInt1D)
        ret.Convert2D(local_grid_scan.value());

    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(m);
    if (!az_int_profile)
        return {};
    return az_int_profile->GetResult();
}

MultiLinePlot JFJochReceiverPlots::GetAzIntProfilePlot(bool force_1d, PlotAzintUnit azint_unit) const {
    std::unique_lock ul(m);
    if (!az_int_profile)
        return {};
    return az_int_profile->GetPlot(force_1d, azint_unit);
}