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

#include "JFJochMPCCDReader.h"

#include <algorithm>
#include <cmath>
#include <filesystem>
#include <fstream>
#include <map>
#include <regex>

#include <zlib.h>

#include "../common/JFJochException.h"
#include "../common/Logger.h"
#include "../writer/HDF5Objects.h"

namespace {

constexpr size_t PANELS = 8;
constexpr size_t PANEL_FS = 512;    // fast axis, the stored image's columns
constexpr size_t PANEL_SS = 1024;   // slow axis: the eight sensors are stacked along it
constexpr size_t STACK_PIXELS = PANELS * PANEL_FS * PANEL_SS;

// Pixels on each sensor that measure nothing usable, as the geometry files the SACLA Cheetah
// pipeline generates mask them: a 5-pixel frame on every side, and the last 23 rows of the slow
// axis, which sit under the amplifier shield.
constexpr size_t EDGE = 5;
constexpr size_t SHIELD_FIRST_SS = 1001;

// The Cheetah-SACLA output is gain-corrected so that 1 eV deposited reads 0.00125 ADU - the
// adu_per_eV its geometry generator writes for every panel - about 10 ADU per 8 keV photon.
constexpr double ADU_PER_EV = 0.00125;

bool IsMPCCDFile(const std::string &path) {
    if (!std::filesystem::is_regular_file(path))
        return false;
    bool ret = false;
    std::unique_lock ul(hdf5_mutex);
    H5E_BEGIN_TRY {
        if (H5Fis_hdf5(path.c_str()) > 0) {
            try {
                HDF5ReadOnlyFile file(path);
                ret = file.Exists("/metadata/sensor_id") && file.Exists("/metadata/posx_in_um")
                      && file.Exists("/metadata/angle_in_rad");
            } catch (...) {
                ret = false;
            }
        }
    } H5E_END_TRY;
    return ret;
}

std::vector<std::string> CollectFiles(const std::string &path) {
    std::vector<std::string> ret;
    if (std::filesystem::is_directory(path)) {
        for (const auto &e: std::filesystem::directory_iterator(path)) {
            const auto p = e.path().string();
            if (e.path().extension() == ".h5" && IsMPCCDFile(p))
                ret.push_back(p);
        }
        std::sort(ret.begin(), ret.end());
    } else if (IsMPCCDFile(path))
        ret.push_back(path);
    return ret;
}

// "+0.000390x -1.000000y" as CrystFEL writes a panel's fs and ss directions.
void ParseDirection(const std::string &txt, double &x, double &y) {
    x = 0;
    y = 0;
    static const std::regex term(R"(([+-]?\s*[0-9.eE+-]*)\s*([xyz]))");
    for (auto it = std::sregex_iterator(txt.begin(), txt.end(), term); it != std::sregex_iterator(); ++it) {
        std::string num = (*it)[1].str();
        num.erase(std::remove(num.begin(), num.end(), ' '), num.end());
        double v = 1.0;
        if (num == "-")
            v = -1.0;
        else if (!num.empty() && num != "+")
            v = std::stod(num);
        if ((*it)[2] == "x")
            x = v;
        else if ((*it)[2] == "y")
            y = v;
    }
}

}  // namespace

JFJochMPCCDReader::JFJochMPCCDReader() = default;
JFJochMPCCDReader::~JFJochMPCCDReader() {
    Close();
}

bool JFJochMPCCDReader::CanRead(const std::string &path) {
    if (std::filesystem::is_directory(path)) {
        for (const auto &e: std::filesystem::directory_iterator(path))
            if (e.path().extension() == ".h5")
                return IsMPCCDFile(e.path().string());
        return false;
    }
    return IsMPCCDFile(path);
}

void JFJochMPCCDReader::ReadFiles(const std::string &path) {
    Close();
    files_ = CollectFiles(path);
    if (files_.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No SACLA MPCCD (Cheetah) files found for " + path);

    std::array<Panel, PANELS> panels;
    double distance_mm = 0, photon_energy_eV = 0;
    std::vector<uint8_t> file_mask;
    {
        std::unique_lock ul(hdf5_mutex);
        for (size_t f = 0; f < files_.size(); f++) {
            open_files_.emplace_back(std::make_unique<HDF5ReadOnlyFile>(files_[f]));
            for (const auto &name: open_files_.back()->FindLeafs("/"))
                if (name.starts_with("tag-"))
                    shots_.push_back({f, "/" + name});
        }
        if (shots_.empty())
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "No images (tag-* groups) in " + path);

        auto &file = *open_files_[0];
        const auto posx = file.ReadVector<double>("/metadata/posx_in_um");
        const auto posy = file.ReadVector<double>("/metadata/posy_in_um");
        // Degrees, whatever the name says: the sensors stand at +-90.
        const auto angle = file.ReadVector<double>("/metadata/angle_in_rad");
        const auto pixel = file.ReadVector<double>("/metadata/pixelsizex_in_um");
        if (posx.size() != PANELS || posy.size() != PANELS || angle.size() != PANELS || pixel.size() != PANELS)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  files_[0] + " does not describe eight MPCCD sensors");
        // The same placement the SACLA pipeline's geometry generator writes for CrystFEL.
        for (size_t i = 0; i < PANELS; i++) {
            const double a = angle[i] * M_PI / 180.0;
            panels[i].corner_x = -posx[i] / pixel[i];
            panels[i].corner_y = posy[i] / pixel[i];
            panels[i].fs_x = std::cos(a);
            panels[i].fs_y = std::sin(a);
            panels[i].ss_x = -std::sin(a);
            panels[i].ss_y = std::cos(a);
        }
        distance_mm = HDF5DataSet(file, "/metadata/distance_in_mm").ReadScalar<double>();
        photon_energy_eV = HDF5DataSet(file, shots_[0].group + "/photon_energy_ev").ReadScalar<double>();
        if (file.Exists("/metadata/pixelmask"))
            file_mask = file.ReadVector<uint8_t>("/metadata/pixelmask");
    }

    if (!(photon_energy_eV > 0) || !(distance_mm > 0))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              files_[0] + " states no photon energy or no camera length");
    Build(panels, distance_mm, photon_energy_eV, file_mask);
}

void JFJochMPCCDReader::ApplyCrystFELGeometry(const std::string &geom_file) {
    if (!dataset_)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "No MPCCD files read");

    std::ifstream in(geom_file);
    if (!in)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Cannot open " + geom_file);

    std::map<std::string, std::string> keys;
    std::string line;
    while (std::getline(in, line)) {
        line = line.substr(0, line.find(';'));
        const auto eq = line.find('=');
        if (eq == std::string::npos)
            continue;
        auto trim = [](std::string s) {
            s.erase(0, s.find_first_not_of(" \t"));
            s.erase(s.find_last_not_of(" \t\r") + 1);
            return s;
        };
        keys[trim(line.substr(0, eq))] = trim(line.substr(eq + 1));
    }

    std::array<Panel, PANELS> panels;
    double coffset = 0;
    for (size_t i = 0; i < PANELS; i++) {
        const std::string q = "q" + std::to_string(i + 1) + "/";
        if (!keys.contains(q + "corner_x") || !keys.contains(q + "fs") || !keys.contains(q + "ss"))
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  geom_file + " does not describe panels q1..q8");
        panels[i].corner_x = std::stod(keys[q + "corner_x"]);
        panels[i].corner_y = std::stod(keys[q + "corner_y"]);
        ParseDirection(keys[q + "fs"], panels[i].fs_x, panels[i].fs_y);
        ParseDirection(keys[q + "ss"], panels[i].ss_x, panels[i].ss_y);
        if (keys.contains(q + "coffset"))
            coffset = std::stod(keys[q + "coffset"]);
    }
    if (!keys.contains("clen"))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, geom_file + " states no clen");
    const double distance_mm = (std::stod(keys["clen"]) + coffset) * 1000.0;

    std::vector<uint8_t> file_mask;
    {
        std::unique_lock ul(hdf5_mutex);
        if (open_files_[0]->Exists("/metadata/pixelmask"))
            file_mask = open_files_[0]->ReadVector<uint8_t>("/metadata/pixelmask");
    }
    Build(panels, distance_mm, dataset_->experiment.GetIncidentEnergy_keV() * 1000.0, file_mask);
}

// CrystFEL's frame has y up and x to the left looking downstream; the flat image has columns to the
// right and rows down. The two differ by a half turn about the beam, which keeps the handedness, so
// a position (X, Y) lands on column beam_x - X and row beam_y - Y.
void JFJochMPCCDReader::Build(const std::array<Panel, 8> &panels, double distance_mm,
                              double photon_energy_eV, const std::vector<uint8_t> &file_mask) {
    auto position = [&](size_t p, double fs, double ss) {
        return std::pair{panels[p].corner_x + fs * panels[p].fs_x + ss * panels[p].ss_x,
                         panels[p].corner_y + fs * panels[p].fs_y + ss * panels[p].ss_y};
    };
    double max_x = -1e9, min_x = 1e9, max_y = -1e9, min_y = 1e9;
    for (size_t p = 0; p < PANELS; p++)
        for (double fs: {0.0, double(PANEL_FS)})
            for (double ss: {0.0, double(PANEL_SS)}) {
                const auto [x, y] = position(p, fs, ss);
                max_x = std::max(max_x, x);
                min_x = std::min(min_x, x);
                max_y = std::max(max_y, y);
                min_y = std::min(min_y, y);
            }
    const double beam_x = std::ceil(max_x) + 1;
    const double beam_y = std::ceil(max_y) + 1;
    width_ = static_cast<size_t>(beam_x - std::floor(min_x)) + 2;
    height_ = static_cast<size_t>(beam_y - std::floor(min_y)) + 2;

    // Everything starts as a gap; each sensor's usable pixels are then placed on top.
    std::vector<uint32_t> mask(width_ * height_, 1u << PixelMask::ModuleGapPixelBit);
    destination_.assign(STACK_PIXELS, -1);
    for (size_t p = 0; p < PANELS; p++) {
        for (size_t ss = 0; ss < PANEL_SS; ss++) {
            for (size_t fs = 0; fs < PANEL_FS; fs++) {
                const size_t src = (p * PANEL_SS + ss) * PANEL_FS + fs;
                const auto [x, y] = position(p, fs + 0.5, ss + 0.5);
                const auto col = std::lround(beam_x - x);
                const auto row = std::lround(beam_y - y);
                const size_t dst = static_cast<size_t>(row) * width_ + static_cast<size_t>(col);
                if (fs < EDGE || fs >= PANEL_FS - EDGE || ss < EDGE || ss >= SHIELD_FIRST_SS)
                    continue;
                destination_[src] = static_cast<int64_t>(dst);
                mask[dst] = (src < file_mask.size() && file_mask[src] != 0) ? (1u << PixelMask::ErrorPixelBit) : 0;
            }
        }
    }

    adu_per_photon_ = ADU_PER_EV * photon_energy_eV;

    dataset_ = std::make_shared<JFJochReaderDataset>();
    dataset_->experiment = default_experiment;
    DetectorSetup detector = DetDECTRIS(static_cast<int64_t>(width_), static_cast<int64_t>(height_), "SACLA MPCCD", {});
    detector.PixelSize_um(50.0f);
    // A CCD: nothing for a parallax or sensor-efficiency correction to work with is stated in the file.
    detector.SensorThickness_um(0);
    // The pipeline stores signed 16-bit ADU, so that is where a pixel stops counting.
    detector.SaturationLimit(static_cast<int64_t>(std::floor(INT16_MAX / adu_per_photon_)));
    detector.BitDepthImage(32);
    detector.MinFrameTime(std::chrono::microseconds(0));
    detector.MinCountTime(std::chrono::microseconds(0));
    detector.ReadOutTime(std::chrono::nanoseconds(0));
    dataset_->experiment.Detector(detector);
    dataset_->experiment.BeamX_pxl(static_cast<float>(beam_x));
    dataset_->experiment.BeamY_pxl(static_cast<float>(beam_y));
    dataset_->experiment.DetectorDistance_mm(static_cast<float>(distance_mm));
    dataset_->experiment.IncidentEnergy_keV(static_cast<float>(photon_energy_eV / 1000.0));
    dataset_->experiment.ImagesPerTrigger(static_cast<int64_t>(shots_.size()));
    dataset_->pixel_mask = std::make_shared<const PixelMask>(mask);

    Logger("MPCCDReader").Info("{} shots in {} files; {} x {} image, beam at ({}, {}) px, {:.2f} mm, "
                               "{:.1f} eV, {:.2f} ADU per photon", shots_.size(), files_.size(),
                               width_, height_, beam_x, beam_y, distance_mm, photon_energy_eV, adu_per_photon_);
    SetStartMessage(dataset_);
}

uint64_t JFJochMPCCDReader::GetNumberOfImages() const {
    return shots_.size();
}

void JFJochMPCCDReader::Close() {
    std::unique_lock ul(hdf5_mutex);
    open_files_.clear();
    files_.clear();
    shots_.clear();
    dataset_.reset();
}

// The stored chunk is the whole image, byte-shuffled and deflated. Only reading it needs the HDF5 lock;
// it is inflated, unshuffled and placed outside it, so the decoding runs on as many threads as read.
template <class Buffer>
void JFJochMPCCDReader::DecodeInto(int64_t image_number, Buffer &buffer) const {
    if (image_number < 0 || static_cast<size_t>(image_number) >= shots_.size())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Image number out of range");
    const auto &shot = shots_[image_number];

    std::vector<uint8_t> compressed;
    {
        std::unique_lock ul(hdf5_mutex);
        HDF5DataSet dataset(*open_files_[shot.file], shot.group + "/data");
        dataset.ReadDirectChunk(compressed, {0, 0});
    }

    std::vector<uint8_t> shuffled(STACK_PIXELS * sizeof(int16_t));
    uLongf out_size = shuffled.size();
    if (uncompress(shuffled.data(), &out_size, compressed.data(), compressed.size()) != Z_OK
        || out_size != shuffled.size())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Cannot inflate image " + shot.group + " of " + files_[shot.file]);

    buffer.assign(width_ * height_ * sizeof(int32_t), 0);
    auto out = reinterpret_cast<int32_t *>(buffer.data());
    for (size_t i = 0; i < STACK_PIXELS; i++) {
        if (destination_[i] < 0)
            continue;
        const auto adu = static_cast<int16_t>(shuffled[i] | (shuffled[STACK_PIXELS + i] << 8));
        out[destination_[i]] = static_cast<int32_t>(std::lround(adu / adu_per_photon_));
    }
}

bool JFJochMPCCDReader::LoadImage_i(std::shared_ptr<JFJochReaderDataset> &dataset,
                                    DataMessage &message,
                                    std::vector<uint8_t> &buffer,
                                    int64_t image_number,
                                    bool update_dataset) {
    (void) update_dataset;
    if (!dataset || image_number < 0 || static_cast<size_t>(image_number) >= shots_.size())
        return false;
    DecodeInto(image_number, buffer);
    message.image = CompressedImage(buffer.data(), buffer.size(), width_, height_,
                                    CompressedImageMode::Int32, CompressionAlgorithm::NO_COMPRESSION);
    message.number = image_number;
    return true;
}

bool JFJochMPCCDReader::ReadRawImage(int64_t image_number, JFJochReaderRawImage &image) {
    if (image_number < 0 || static_cast<size_t>(image_number) >= shots_.size())
        return false;
    DecodeInto(image_number, image.image_buffer);
    image.image = CompressedImage(image.image_buffer.data(), image.image_buffer.size(), width_, height_,
                                  CompressedImageMode::Int32, CompressionAlgorithm::NO_COMPRESSION);
    return true;
}

std::vector<SpotToSave> JFJochMPCCDReader::ReadSpots(int64_t) const {
    return {};
}