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

#include "StillsGPUAnalysis.h"

#include <thread>

#include "../common/CUDAWrapper.h"
#include "../image_analysis/MXAnalysisWithoutFPGA.h"

StillsGPUAnalysis::StillsGPUAnalysis(const DiffractionExperiment &in_experiment, const PixelMask &in_mask,
                                     const StartMessage &in_start, IndexerThreadPool *in_indexing_thread_pool,
                                     int64_t nthreads, int64_t max_held, Logger &in_logger, Ready in_ready,
                                     Sink in_sink, Settings in_settings, Failed in_failed,
                                     RugnuxLiveStatus *in_status)
    : experiment(in_experiment), mask(in_mask), start(in_start), indexing_thread_pool(in_indexing_thread_pool),
      logger(in_logger), ready(std::move(in_ready)), sink(std::move(in_sink)), settings(std::move(in_settings)),
      setup_failed(std::move(in_failed)), status(in_status), queue(std::max<int64_t>(1, max_held)) {
    live.images_needed = std::min(PREPASS_IMAGES, experiment.GetImageNum());
    live.file_beam_x_pxl = experiment.GetBeamX_pxl();
    live.file_beam_y_pxl = experiment.GetBeamY_pxl();
    live.file_distance_mm = experiment.GetDetectorDistance_mm();
    if (status) {
        status->Start(experiment.GetImageNum(),
                      static_cast<double>(experiment.GetImageNum())
                      * std::chrono::duration<double>(experiment.GetImageTime()).count(),
                      std::chrono::steady_clock::now());
        status->Step("Pre-pass");
        Publish();
    }
    prepass_future = std::async(std::launch::async, &StillsGPUAnalysis::PrePassThread, this);
    for (int64_t i = 0; i < nthreads; i++)
        analysis_futures.emplace_back(std::async(std::launch::async, &StillsGPUAnalysis::AnalysisThread, this));
    logger.Info("GPU analysis: {} threads, the first {} images held for the pre-pass", nthreads, PREPASS_IMAGES);
}

StillsGPUAnalysis::~StillsGPUAnalysis() {
    try {
        if (!finished)
            Finish();
    } catch (const std::exception &e) {
        logger.Error("GPU analysis: {}", e.what());
    }
}

void StillsGPUAnalysis::AnalysisError(const std::string &text) {
    logger.Error("GPU analysis: {}", text);
    std::unique_lock ul(error_mutex);
    if (analysis_error.empty())
        analysis_error = text;
}

void StillsGPUAnalysis::Publish() {
    std::unique_lock ul(live_mutex);
    if (live.prepass == RugnuxLivePrePass::Done || live.prepass == RugnuxLivePrePass::Failed) {
        const double t = std::chrono::duration<double>(std::chrono::steady_clock::now() - ready_time).count();
        if (t > 0 && live.images_analysed > 0)
            live.analysis_rate_hz = static_cast<double>(live.images_analysed) / t;
    }
    status->Stills(live);
}

bool StillsGPUAnalysis::Hold(const std::shared_ptr<HeldImage> &image) {
    if (status)
        status->ImageReceived(std::chrono::steady_clock::now());
    if (!queue.Put(image)) {
        // As many frames wait as the buffer is to hold: the pre-pass, if it has not run, runs on what it has.
        StartPrePass();
        return false;
    }
    {
        std::unique_lock ul(prepass_mutex);
        if (!prepass_go) {
            if (image->message.number < PREPASS_IMAGES) {
                prepass_images[image->message.number] = image;
                std::unique_lock ul_live(live_mutex);
                live.images_held = static_cast<int64_t>(prepass_images.size());
            }
            ++images_held;
            // Every one of the first images is there - or, where some never come, twice as many are.
            prepass_go = (prepass_images.size() == std::min(PREPASS_IMAGES, experiment.GetImageNum()))
                         || (images_held >= 2 * PREPASS_IMAGES);
        }
    }
    prepass_cv.notify_all();
    if (status)
        Publish();
    return true;
}

void StillsGPUAnalysis::StartPrePass() {
    {
        std::unique_lock ul(prepass_mutex);
        prepass_go = true;
    }
    prepass_cv.notify_all();
}

void StillsGPUAnalysis::Finish() {
    finished = true;
    StartPrePass();
    prepass_future.get();
    for (size_t i = 0; i < analysis_futures.size(); i++)
        queue.PutBlocking(nullptr);
    for (auto &f: analysis_futures)
        f.get();
    if (status) {
        status->LastImage();
        Publish();
        if (setup_error)
            status->Finish(RugnuxLiveState::Failed, "GPU analysis setup failed");
        else {
            std::unique_lock ul(error_mutex);
            if (analysis_error.empty())
                status->Finish(RugnuxLiveState::Done);
            else
                status->Finish(RugnuxLiveState::Failed, analysis_error);
        }
    }
    if (setup_error)
        std::rethrow_exception(setup_error);
}

void StillsGPUAnalysis::PrePassThread() {
    std::map<int64_t, std::shared_ptr<HeldImage>> sample;
    {
        std::unique_lock ul(prepass_mutex);
        prepass_cv.wait(ul, [this] { return prepass_go; });
        sample.swap(prepass_images);
    }

    const auto masked = [](const PixelMask &m) {
        return std::count_if(m.GetMask().begin(), m.GetMask().end(), [](uint32_t v) { return v != 0; });
    };
    setup.experiment = experiment;
    setup.mask = mask;
    {
        std::unique_lock ul(live_mutex);
        live.prepass = RugnuxLivePrePass::Running;
    }
    if (status)
        Publish();
    bool failed = sample.empty();
    const auto prepass_start = std::chrono::steady_clock::now();
    if (!sample.empty()) {
        try {
            const auto t0 = std::chrono::steady_clock::now();
            const auto r = RunStillsPrePass(start, sample, static_cast<int>(std::thread::hardware_concurrency()));
            setup.experiment.BeamX_pxl(r.beam_x_pxl).BeamY_pxl(r.beam_y_pxl).DetectorDistance_mm(r.distance_mm);
            setup.mask = r.mask;
            const auto &m = r.mask.GetMask();
            std::unique_lock ul(live_mutex);
            live.beam_x_pxl = r.beam_x_pxl;
            live.beam_y_pxl = r.beam_y_pxl;
            live.distance_mm = r.distance_mm;
            live.shadow_pixels = std::count_if(m.begin(), m.end(), [](uint32_t v) {
                return (v & (1u << PixelMask::BeamStopPixelBit)) != 0;
            });
            live.bad_pixels = std::count_if(m.begin(), m.end(), [](uint32_t v) {
                return (v & (1u << PixelMask::HotPixelBit)) != 0;
            });
            logger.Info("Pre-pass on {} images ({:.1f} s): beam centre ({:.2f}, {:.2f}) -> ({:.2f}, {:.2f}) px, "
                        "distance {:.2f} -> {:.2f} mm, masked pixels {} -> {}", r.images,
                        std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count(),
                        experiment.GetBeamX_pxl(), experiment.GetBeamY_pxl(), r.beam_x_pxl, r.beam_y_pxl,
                        experiment.GetDetectorDistance_mm(), r.distance_mm,
                        masked(mask), masked(r.mask));
        } catch (const std::exception &e) {
            failed = true;
            logger.Warning("Pre-pass failed, the images are analysed with the stated beam centre and mask - {}",
                           e.what());
        }
    }
    sample.clear();
    try {
        setup.mapping = std::make_unique<AzimuthalIntegrationMapping>(setup.experiment, setup.mask);
        indexer = std::make_unique<IndexAndRefine>(setup.experiment, indexing_thread_pool, false, true);
        ready(setup);
    } catch (const std::exception &e) {
        // Nothing can be written without the start message: the frames are dropped (AnalysisThread), and
        // the collection is stopped.
        const auto *j = dynamic_cast<const JFJochException *>(&e);
        const JFJochException error = j ? *j : JFJochException(JFJochExceptionCategory::WrongDAQState,
                                                               std::string("GPU analysis setup failed: ") + e.what());
        logger.ErrorException(error);
        setup_error = j ? std::current_exception() : std::make_exception_ptr(error);
        if (setup_failed)
            setup_failed(error);
        prepass_done.count_down();
        return;
    }
    {
        std::unique_lock ul(live_mutex);
        live.prepass = failed ? RugnuxLivePrePass::Failed : RugnuxLivePrePass::Done;
        live.prepass_s = std::chrono::duration<double>(std::chrono::steady_clock::now() - prepass_start).count();
        ready_time = std::chrono::steady_clock::now();
    }
    if (status) {
        status->Step("GPU analysis");
        Publish();
    }
    prepass_done.count_down();
}

void StillsGPUAnalysis::AnalysisThread() {
    try {
        pin_gpu();
    } catch (const std::exception &e) {
        logger.Warning("Error pinning GPU {}", e.what());
    }

    prepass_done.wait();

    if (setup_error) {
        while (queue.GetBlocking()) {}
        return;
    }

    // A thread whose analysis cannot be built still passes its frames on, unanalysed.
    std::unique_ptr<MXAnalysisWithoutFPGA> analysis;
    try {
        analysis = std::make_unique<MXAnalysisWithoutFPGA>(setup.experiment, *setup.mapping, setup.mask, *indexer,
                                                           /*enable_fused_adaptive_gpu=*/true);
    } catch (const std::exception &e) {
        AnalysisError(std::string("analysis thread could not start, its frames are written unanalysed - ")
                      + e.what());
    }

    while (auto held = queue.GetBlocking()) {
        DataMessage &message = held->message;
        AzimuthalIntegrationProfile profile(*setup.mapping);
        const auto start_time = std::chrono::steady_clock::now();
        try {
            if (analysis)
                analysis->Analyze(message, profile, settings());
        } catch (const JFJochException &e) {
            // The frame is still written, without its analysis.
            logger.ErrorException(e);
        } catch (const std::exception &e) {
            AnalysisError(fmt::format("image {} written unanalysed - {}", message.number, e.what()));
        }
        message.processing_time_s = std::chrono::duration<float>(std::chrono::steady_clock::now()
                                                                  - start_time).count();
        try {
            sink(message, profile);
            if (status) {
                int64_t analysed;
                {
                    std::unique_lock ul(live_mutex);
                    analysed = ++live.images_analysed;
                    if (message.indexing_result.value_or(false))
                        ++live.images_indexed;
                }
                status->Progress(analysed, experiment.GetImageNum());
                Publish();
            }
        } catch (const std::exception &e) {
            // The thread goes on with the next frame, so that the queue is drained.
            AnalysisError(fmt::format("image {} lost - {}", message.number, e.what()));
        }
    }
}

std::optional<UnitCell> StillsGPUAnalysis::GetConsensusUnitCell() const {
    return indexer->GetConsensusUnitCell();
}

float StillsGPUAnalysis::GetBeamX_pxl() const {
    return setup.experiment.GetBeamX_pxl();
}

float StillsGPUAnalysis::GetBeamY_pxl() const {
    return setup.experiment.GetBeamY_pxl();
}

float StillsGPUAnalysis::GetDetectorDistance_mm() const {
    return setup.experiment.GetDetectorDistance_mm();
}