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

#include "DiagnosticOutput.h"

#include <algorithm>
#include <cmath>
#include <fstream>

#include <spdlog/fmt/fmt.h>

#include "../common/ColorScale.h"
#include "../common/CompressedImage.h"
#include "../common/JFJochException.h"
#include "../image_analysis/scale_merge/Merge.h"
#include "../preview/JFJochJPEG.h"

namespace {
    // The top of the ramp is eight times the MEDIAN of the measured pixels, not the brightest 0.1 %
    // the preview and the viewer saturate at (preview/PreviewImage.cpp). That percentile belongs to
    // the Bragg peaks, and with it setting the range the background lands in the first few per cent
    // of the scale and reads as a flat white - which is exactly what the picture has to be read for,
    // since a shadow is a place the background is MISSING. The median is pure background whatever
    // the crystal, and eight times it lands between the 94th and the 99th percentile of the
    // projection on every test set measured, which is the top of the background: the background
    // spans the ramp radially, and the spots clip, which costs nothing here as they are not what is
    // being judged. Twice the median - the first value tried - lands at the 73rd to 90th percentile
    // instead and floods the whole inner half of the detector to a flat indigo, hiding the deficit
    // just as the old contrast did, at the other end of the scale.
    float SaturationValue(const std::vector<float> &projection, const std::vector<uint32_t> &mask) {
        std::vector<float> valid;
        valid.reserve(projection.size());
        for (size_t i = 0; i < projection.size(); i++)
            if (mask[i] == 0 && std::isfinite(projection[i]))
                valid.push_back(projection[i]);
        if (valid.empty())
            return 1.0f;
        const size_t mid = valid.size() / 2;
        std::nth_element(valid.begin(), valid.begin() + mid, valid.end());
        return std::max(8.0f * valid[mid], 1.0f);
    }

    float Element(const std::vector<float> &v, size_t i) {
        return i < v.size() ? v[i] : NAN;
    }

    int Element(const std::vector<int> &v, size_t i) {
        return i < v.size() ? v[i] : 0;
    }
}

std::vector<rgb> DetectorDiagnosticRGB(const PixelMask &pixel_mask,
                                       const std::vector<float> &mean_projection) {
    constexpr uint32_t gap_bits = (1u << PixelMask::ModuleGapPixelBit)
                                  | (1u << PixelMask::ChipGapPixelBit)
                                  | (1u << PixelMask::ModuleEdgePixelBit);
    constexpr uint32_t beam_stop_bit = 1u << PixelMask::BeamStopPixelBit;

    const auto &mask = pixel_mask.GetMask();
    if (mask.size() != mean_projection.size())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mask and projection are of different size");

    // Untouched: the viewer's defaults are the point, so that the JPEG and the picture the user
    // knows from jfjoch_viewer say the same thing with the same colours.
    ColorScale scale;
    const rgb gap_color = scale.Apply(ColorScaleSpecial::Gap);
    const rgb bad_color = scale.Apply(ColorScaleSpecial::BadPixel);
    const rgb beam_stop_color = scale.Apply(ColorScaleSpecial::BeamStop);

    const auto &lut = scale.LUTData();
    const auto lut_size = static_cast<int64_t>(lut.size());
    const float inv_range = static_cast<float>(lut_size - 1) / SaturationValue(mean_projection, mask);

    // The shadow is tested BEFORE the rest of the mask, so a pixel the detection claims is never
    // hidden under a noisy- or user-masked bit it also carries - the whole point is to see what
    // the detection did. Gaps come first because nothing can be shadowed there: they carry no
    // background for the comparison to miss.
    std::vector<rgb> image(mean_projection.size());
    for (size_t i = 0; i < image.size(); i++) {
        if ((mask[i] & gap_bits) != 0)
            image[i] = gap_color;
        else if ((mask[i] & beam_stop_bit) != 0)
            image[i] = beam_stop_color;
        else if (mask[i] != 0 || !std::isfinite(mean_projection[i]))
            image[i] = bad_color;
        else {
            const auto idx = static_cast<int64_t>(mean_projection[i] * inv_range + 0.5f);
            image[i] = lut[std::clamp<int64_t>(idx, 0, lut_size - 1)];
        }
    }

    return image;
}

std::string RenderDetectorDiagnostic(const DiffractionExperiment &experiment,
                                     const PixelMask &pixel_mask,
                                     const std::vector<float> &mean_projection) {
    // Named, because CompressedImage does not own the pixels it is handed.
    const std::vector<rgb> image = DetectorDiagnosticRGB(pixel_mask, mean_projection);
    const CompressedImage rgb_image(image, experiment.GetXPixelsNumConv(),
                                    experiment.GetYPixelsNumConv());
    return WriteJPEGToMem(rgb_image, 85);
}

void WritePerImagePlot(const std::string &output_prefix, const PerImagePlot &plot,
                       const std::optional<GoniometerAxis> &goniometer, size_t image_count) {
    const std::string path = output_prefix + "_plot.txt";
    std::ofstream file(path, std::ofstream::out | std::ofstream::trunc);
    if (!file)
        throw JFJochException(JFJochExceptionCategory::FileWriteError,
                              "Cannot open " + path + " for writing");

    // One comment line, its names over their own columns, and nothing else a plotting program has
    // to be told to skip. The '#' takes the first character of the image-number field.
    file << "#" << fmt::format("{:>5} {:>10} {:>12} {:>12} {:>8} {:>12} {:>13} {:>12} {:>8}"
                               " {:>7}\n",
                               "image", "angle_deg", "bkg", "resolution_A", "spots", "scale_G",
                               "sigma_M_deg", "cc_to_merge", "cc_n", "merged");

    // Whether "merged" has anything to answer at all. With no merge behind the run - --no-merge, or
    // a sweep nothing indexed on - a column of zeros would read as a sweep every frame of which was
    // thrown out, which is a different statement from "no merge was made".
    const bool merge_ran = !plot.frame_disposition.empty()
                           || std::any_of(plot.scale_g.begin(), plot.scale_g.end(),
                                          [](float g) { return std::isfinite(g); });

    for (size_t i = 0; i < image_count; i++) {
        const float angle = goniometer
                                ? goniometer->GetAngle_deg(static_cast<float>(i))
                                      + goniometer->GetWedge_deg() / 2.0f
                                : NAN;
        const float scale_g = Element(plot.scale_g, i);
        // Did this image's observations reach the merged data? The sweep-quality diagnostic says so
        // outright; where it did not run, an image with no fitted scale contributed nothing, which
        // is the same statement by a weaker route.
        const int merged = i < plot.frame_disposition.size()
                               ? plot.frame_disposition[i]
                                     != static_cast<uint8_t>(FrameDisposition::Rejected)
                               : std::isfinite(scale_g);

        file << fmt::format("{:>6} {:>10.4f} {:>12.4g} {:>12.4g} {:>8.0f} {:>12.5g} {:>13.4f}"
                            " {:>12.4f} {:>8} {:>7}\n",
                            i, angle, Element(plot.bkg_estimate, i), Element(plot.resolution_A, i),
                            Element(plot.spot_count, i), scale_g, Element(plot.mosaicity_deg, i),
                            Element(plot.cc_to_merge, i), Element(plot.cc_n, i),
                            merge_ran ? std::to_string(merged) : "nan");
    }
}