// 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 "../common/ThreadAffinity.h"
#include "../image_analysis/MXAnalysisWithoutFPGA.h"
#include "../rugnux/RugnuxDefaults.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, bool in_merge)
    : 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), merge_frames(in_merge),
      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());
    }
    if (status && outcome)
        status->Finish(outcome->first, outcome->second);
}

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) {
        last_image = status->LastImage();
        Publish();
    }
    if (merge_future.valid()) {
        if (status)
            status->Step("Merge");
        merge_queue.Put(-1);
        merge_future.get();
        if (merged && status) {
            PublishMerge(*merged, "the collection");
            status->Headline(merge_experiment.GetSpaceGroupOrP1().xhm(), merged->merge.statistics.overall.d_min,
                             std::nullopt);
        }
    }
    if (status) {
        if (setup_error)
            status->Finish(RugnuxLiveState::Failed, "GPU analysis setup failed");
        else {
            // Reported by Written, once everything is written - the end message and the writer go on meanwhile.
            status->Step("Writing");
            std::unique_lock ul(error_mutex);
            if (analysis_error.empty())
                outcome.emplace(RugnuxLiveState::Done, "");
            else
                outcome.emplace(RugnuxLiveState::Failed, analysis_error);
        }
    }
    if (setup_error)
        std::rethrow_exception(setup_error);
}

void StillsGPUAnalysis::Written() {
    if (!status || !outcome)
        return;
    if (last_image)
        status->Latency(std::chrono::duration<double>(std::chrono::steady_clock::now() - last_image.value()).count());
    status->Finish(outcome->first, outcome->second);
    outcome.reset();
}

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();
    int geometry_frames = 0;
    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;
            geometry_frames = r.geometry_frames;
            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);
        // Retained for the merge, which takes each image's reflections as soon as it is analysed.
        indexer = std::make_unique<IndexAndRefine>(setup.experiment, indexing_thread_pool, merge_frames, true);
        // The pre-pass's fit rests on the frames of its sample that indexed; where those are fewer than the
        // running refinement bundles, the geometry is refined again, from where the pre-pass left it, on the
        // frames as they are analysed.
        if (geometry_frames < StillsRunningGeometry::STRONG_FRAMES && setup.experiment.GetUnitCell()) {
            running_geometry = std::make_unique<StillsRunningGeometry>(setup.experiment, MERGE_THREADS);
            frame_epoch.assign(experiment.GetImageNum(), 0);
            geometry_state = GEOMETRY_PENDING;
            logger.Info("The pre-pass refined the geometry on {} frames: it is refined again on {} strong frames as "
                        "they are analysed", geometry_frames, StillsRunningGeometry::STRONG_FRAMES);
        }
        if (merge_frames) {
            merge_experiment = setup.experiment;
            merge_experiment.ImportScalingSettings(RugnuxDefaultScalingSettings(false));
            merge = std::make_unique<StillsLiveMerge>(merge_experiment, experiment.GetImageNum(),
                                                      setup.experiment.GetUnitCell(), std::nullopt,
                                                      setup.experiment.IsDetectIceRings(),
                                                      SpotFindingSettings().ice_ring_width_Q_recipA, logger);
            merge_future = std::async(std::launch::async, &StillsGPUAnalysis::MergeThread, this);
        }
        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. It is built again at
    // the geometry the running refinement switches to.
    std::unique_ptr<MXAnalysisWithoutFPGA> analysis;
    int epoch = -1;
    const auto build = [&] {
        epoch = geometry_epoch;
        analysis.reset();
        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;
        std::optional<AzimuthalIntegrationProfile> profile;
        bool switch_geometry = false;
        const auto start_time = std::chrono::steady_clock::now();
        {
            std::unique_lock gate(gate_mutex);
            gate_cv.wait(gate, [this] { return !switch_waiting; });
        }
        {
            std::shared_lock sl(geometry_mutex);
            if (epoch != geometry_epoch)
                build();
            profile.emplace(*setup.mapping);
            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()));
            }
            if (geometry_state == GEOMETRY_PENDING) {
                frame_epoch.at(message.number) = static_cast<uint8_t>(epoch + 1);
                switch_geometry = running_geometry->Add(message);
            }
        }
        message.processing_time_s = std::chrono::duration<float>(std::chrono::steady_clock::now()
                                                                  - start_time).count();
        if (switch_geometry)
            SwitchGeometry();
        try {
            sink(message, *profile);
            if (merge)
                merge_queue.Put(message.number);
            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()));
        }
    }
}

void StillsGPUAnalysis::SwitchGeometry() {
    // One fit, in the analysis thread that offered the last frame it needed; the others go on meanwhile.
    const auto t0 = std::chrono::steady_clock::now();
    const auto r = running_geometry->Refine();
    const bool last_round = ++geometry_round == GEOMETRY_ROUNDS;
    if (!r) {
        logger.Warning("Running geometry refinement, round {}: no fit on {} frames ({} strong); the frames are "
                       "analysed at the geometry they were", geometry_round, running_geometry->Offered(),
                       running_geometry->Strong());
        geometry_state = geometry_epoch > 0 ? GEOMETRY_SWITCHED : GEOMETRY_KEPT;
        return;
    }
    const std::string text = fmt::format(
            "Geometry refined on {} analysed frames ({} strong), round {}: beam centre ({:.2f}, {:.2f}) -> "
            "({:.2f}, {:.2f}) px, distance {:.2f} -> {:.2f} mm ({:.1f} s)", running_geometry->Offered(),
            running_geometry->Strong(), geometry_round, setup.experiment.GetBeamX_pxl(),
            setup.experiment.GetBeamY_pxl(), r->beam_x_px, r->beam_y_px, setup.experiment.GetDetectorDistance_mm(),
            r->distance_mm, std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count());
    {
        std::unique_lock gate(gate_mutex);
        switch_waiting = true;
    }
    {
        std::unique_lock ul(geometry_mutex);
        setup.experiment.BeamX_pxl(static_cast<float>(r->beam_x_px)).BeamY_pxl(static_cast<float>(r->beam_y_px))
                .DetectorDistance_mm(static_cast<float>(r->distance_mm));
        retired_mappings.push_back(std::move(setup.mapping));
        setup.mapping = std::make_unique<AzimuthalIntegrationMapping>(setup.experiment, setup.mask);
        ++geometry_epoch;
        if (last_round)
            geometry_state = GEOMETRY_SWITCHED;
        else
            running_geometry = std::make_unique<StillsRunningGeometry>(setup.experiment, MERGE_THREADS);
    }
    {
        std::unique_lock gate(gate_mutex);
        switch_waiting = false;
    }
    gate_cv.notify_all();
    {
        std::unique_lock ul(live_mutex);
        live.beam_x_pxl = r->beam_x_px;
        live.beam_y_pxl = r->beam_y_px;
        live.distance_mm = r->distance_mm;
    }
    logger.Info("{}", text);
    if (status)
        status->Note(text);
}

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

const std::optional<StillsMergeOutput> &StillsGPUAnalysis::GetMerged() const {
    return merged;
}

const DiffractionExperiment &StillsGPUAnalysis::GetMergeExperiment() const {
    return merge_experiment;
}

void StillsGPUAnalysis::MergeThread() {
    constexpr int64_t PREPARE_EVERY = 200;   // frames between two Prepare calls after a checkpoint
    const int64_t images = experiment.GetImageNum();
    // The cell merged in, as rugnux merges it: the consensus of the frames indexed by the first merge (the
    // collection's cell only picks the frames that are merged).
    bool cell_set = false;
    const auto have_cell = [&] {
        if (!cell_set)
            if (const auto c = indexer->GetConsensusUnitCell()) {
                merge_experiment.SetUnitCell(*c);
                merge->Cell(*c);
                cell_set = true;
            }
        return cell_set;
    };
    size_t next = 0;
    try {
        const auto add = [&](int64_t n, IntegrationOutcome &&o) {
            merge->Add(n, std::move(o));
            if (next < CHECKPOINTS.size()
                && merge->Added() >= std::llround(CHECKPOINTS[next] * static_cast<double>(images))) {
                if (have_cell())
                    PublishMerge(merge->Checkpoint(MERGE_THREADS),
                                 fmt::format("{:.0f}% of the images (provisional)", 100.0 * CHECKPOINTS[next]));
                ++next;
            } else if (next > 0 && merge->Added() % PREPARE_EVERY == 0)
                merge->Prepare(MERGE_THREADS);
        };
        // The frames analysed at a geometry the running refinement may still replace wait until it is decided:
        // left out where analysed before its last switch, merged as analysed otherwise.
        std::vector<int64_t> held_back;
        const auto release = [&] {
            for (const int64_t h: held_back)
                if (frame_epoch[h] - 1 < geometry_epoch) {
                    indexer->DropIntegrationOutcome(h);
                    add(h, IntegrationOutcome{});
                } else
                    add(h, indexer->TakeIntegrationOutcome(h));
            held_back.clear();
        };
        for (int64_t n = merge_queue.GetBlocking(); n >= 0; n = merge_queue.GetBlocking()) {
            if (!frame_epoch.empty() && frame_epoch[n] > 0)
                held_back.push_back(n);
            else
                add(n, indexer->TakeIntegrationOutcome(n));
            if (geometry_state != GEOMETRY_PENDING && !held_back.empty())
                release();
        }
        release();
        if (have_cell()) {
            merged = merge->Final(static_cast<int>(AvailableCpus()));
            logger.Info("Stills merge after the last image: {:.2f} s", merged->seconds);
        }
    } catch (const std::exception &e) {
        logger.Warning("Stills merge failed: {}", e.what());
        merged.reset();
    }
}

void StillsGPUAnalysis::PublishMerge(const StillsMergeOutput &out, const std::string &title) {
    const auto &o = out.merge.statistics.overall;
    const std::string text = fmt::format("Merge of {}: {} crystals, {:.2f} A, CC1/2 {:.3f}, R_split {:.3f}, "
                                         "completeness {:.1f}%, <I/sigma> {:.1f} ({:.1f} s)", title, out.images,
                                         o.d_min, o.cc_half, o.r_split,
                                         o.possible_unique_reflections > 0
                                         ? 100.0 * o.unique_reflections / o.possible_unique_reflections : 0.0,
                                         o.mean_i_over_sigma, out.seconds);
    logger.Info("{}", text);
    if (status)
        status->Note(text);
}