// Copyright (2019-2023) Paul Scherrer Institute

#include "JFJochReceiverPlots.h"

JFJochReceiverPlots::JFJochReceiverPlots(const DiffractionExperiment &experiment,
                                         const AzimuthalIntegrationMapping &mapping)
        : indexing_solution_per_file(experiment.GetTimePointNumber()),
          default_binning(experiment.GetDefaultPlotBinning()),
          az_int_profile(mapping), resolution_estimation(50, 1.0, 5.0) {
    az_int_profile.SetTitle("dataset");
    for (int i = 0; i < experiment.GetTimePointNumber(); i++) {
        az_int_profile_per_file.emplace_back(mapping);
        az_int_profile_per_file[i].SetTitle(std::to_string(i));
    }

    for (const auto &[x, y]: experiment.ROI().GetROINameMap()) {
        roi_sum[x] = std::make_unique<StatusVector<int64_t>>();
        roi_max_count[x] = std::make_unique<StatusVector<int64_t>>();
        roi_pixels[x] = std::make_unique<StatusVector<uint64_t>>();
    }
}

void JFJochReceiverPlots::Add(const DataMessage &msg, uint64_t file_number, const AzimuthalIntegrationProfile &profile) {
    bkg_estimate.AddElement(msg.number, msg.bkg_estimate);
    spot_count.AddElement(msg.number, msg.spots.size());
    saturated_pixels.AddElement(msg.number, msg.saturated_pixel_count);
    error_pixels.AddElement(msg.number, msg.error_pixel_count);
    strong_pixels.AddElement(msg.number, msg.strong_pixel_count);

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

    indexing_solution.AddElement(msg.number, msg.indexing_result);
    indexing_solution_per_file.Add(file_number, msg.indexing_result);

    if (msg.corr_beam_x_pxl)
        corr_beam_x_pxl.AddElement(msg.number, msg.corr_beam_x_pxl.value());
    if (msg.corr_beam_y_pxl)
        corr_beam_y_pxl.AddElement(msg.number, msg.corr_beam_y_pxl.value());
    if (msg.corr_det_dist_mm)
        corr_det_dist_mm.AddElement(msg.number, msg.corr_det_dist_mm.value());

    if (msg.resolution_estimation)
        resolution_estimation.Add(msg.resolution_estimation.value());

    az_int_profile += profile;
    az_int_profile_per_file.at(file_number) += profile;

    if (msg.xfel_pulse_id)
        xfel_pulse_id.AddElement(msg.number, msg.xfel_pulse_id.value());

    if (msg.xfel_event_code)
        xfel_event_code.AddElement(msg.number, msg.xfel_event_code.value());

    for (const auto &[key, value] : msg.roi) {
        if (roi_sum.contains(key) && roi_sum[key])
            roi_sum[key]->AddElement(msg.number, value.sum);
        if (roi_max_count.contains(key) && roi_max_count[key])
            roi_max_count[key]->AddElement(msg.number, value.max_count);
        if (roi_pixels.contains(key) && roi_pixels[key])
            roi_pixels[key]->AddElement(msg.number, value.pixels);
    }
}

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::RadInt:
            return az_int_profile.GetPlot();
        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::IndexingRatePerFile:
            return indexing_solution_per_file.GetPlot();
        case PlotType::IndexingDrift:
            tmp = corr_beam_x_pxl.GetMeanPlot(nbins);
            tmp.at(0).title = "X beam drift (pxl)";
            ret.emplace_back(tmp.at(0));
            tmp = corr_beam_y_pxl.GetMeanPlot(nbins);
            tmp.at(0).title = "Y beam drift (pxl)";
            ret.emplace_back(tmp.at(0));
            tmp = corr_det_dist_mm.GetMeanPlot(nbins);
            tmp.at(0).title = "Det. distance drift (mm)";
            ret.emplace_back(tmp.at(0));
            return ret;
        case PlotType::ErrorPixels:
            tmp = saturated_pixels.GetMeanPlot(nbins);
            tmp.at(0).title = "Saturated pixels (mean)";
            ret.emplace_back(tmp.at(0));

            tmp = saturated_pixels.GetMaxPlot(nbins);
            tmp.at(0).title = "Saturated pixels (max)";
            ret.emplace_back(tmp.at(0));

            tmp = error_pixels.GetMeanPlot(nbins);
            tmp.at(0).title = "Error pixels (mean)";
            ret.emplace_back(tmp.at(0));
            return ret;
        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:
            for (const auto &[key, value]: roi_sum) {
                tmp = value->GetMeanPlot(nbins);
                tmp.at(0).title = key;
                ret.emplace_back(tmp.at(0));
            }
            return ret;
        case PlotType::ROIMaxCount:
            for (const auto &[key, value]: roi_max_count) {
                tmp = value->GetMeanPlot(nbins);
                tmp.at(0).title = key;
                ret.emplace_back(tmp.at(0));
            }
            return ret;
        case PlotType::ROIPixels:
            for (const auto &[key, value]: roi_pixels) {
                tmp = value->GetMeanPlot(nbins);
                tmp.at(0).title = key;
                ret.emplace_back(tmp.at(0));
            }
            return ret;
        case PlotType::ResEstimation:
            return resolution_estimation.GetPlot();
        case PlotType::RadIntPerFile:
            tmp = az_int_profile.GetPlot();
            ret.emplace_back(tmp.at(0));

            for (const auto & i : az_int_profile_per_file) {
                tmp = i.GetPlot();
                ret.emplace_back(tmp.at(0));
            }
            return ret;
        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 {};
}

std::optional<std::pair<float, float>> JFJochReceiverPlots::GetBeamCenterDriftPxl() const {
    auto diffx = corr_beam_x_pxl.Mean();
    auto diffy = corr_beam_y_pxl.Mean();
    if (std::isfinite(diffx) && std::isfinite(diffy))
        return std::make_pair(diffx, diffy);
    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 {
    return az_int_profile.GetResult();
}

std::vector<float> JFJochReceiverPlots::GetAzIntProfilePerFile(uint64_t file_number) {
    if (file_number < az_int_profile_per_file.size())
        return az_int_profile_per_file.at(file_number).GetResult();
    else
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds, "JFJochReceiverPlots::GetAzIntProfilePerFile out of bounds");
}