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

#include "JFJochImageReadingWorker.h"
#include "../image_analysis/geom_refinement/AssignSpotsToRings.h"
#include "../image_analysis/spot_finding/StrongPixelSet.h"
#include "../image_analysis/spot_finding/SpotUtils.h"
#include "../image_analysis/spot_finding/ImageSpotFinder.h"
#include <QVector>
#include <QMutexLocker>

#include "../preview/JFJochTIFF.h"

JFJochImageReadingWorker::JFJochImageReadingWorker(const SpotFindingSettings &settings,
                                                   const DiffractionExperiment &experiment, QObject *parent)
    : QObject(parent),
      indexing_settings(experiment.GetIndexingSettings()),
      azint_settings(experiment.GetAzimuthalIntegrationSettings()) {
    spot_finding_settings = settings;;

    indexing = std::make_unique<IndexerThreadPool>(indexing_settings);
    http_reader.Experiment(experiment);
    file_reader.Experiment(experiment);

    autoload_timer = new QTimer(this);
    autoload_timer->setInterval(autoload_interval);
    connect(autoload_timer, &QTimer::timeout, this, &JFJochImageReadingWorker::AutoLoadTimerExpired);
}

void JFJochImageReadingWorker::LoadFile(const QString &filename, qint64 image_number, qint64 summation) {
    QMutexLocker ul(&m);

    try {
        std::shared_ptr<const JFJochReaderDataset> dataset;
        auto start = std::chrono::high_resolution_clock::now();

        if (!http_mode && filename == current_file) {
            logger.Info("File {} already loaded", filename.toStdString());
        } else {
            if (filename.startsWith("http://")) {
                http_mode = true;
                http_reader.ReadURL(filename.toStdString());
                total_images = http_reader.GetNumberOfImages();
                dataset = http_reader.GetDataset();
                if (image_number < 0)
                    setAutoLoadMode_i(AutoloadMode::HTTPSync);
                else
                    setAutoLoadMode_i(AutoloadMode::None);
            } else {
                http_mode = false;
                file_reader.ReadFile(filename.toStdString());
                total_images = file_reader.GetNumberOfImages();
                dataset = file_reader.GetDataset();
                setAutoLoadMode_i(AutoloadMode::None);
            }
            current_image.reset();
            current_summation = 1;
            current_file = filename;
            curr_experiment = dataset->experiment;
            curr_experiment.ImportIndexingSettings(indexing_settings);
            curr_experiment.ImportAzimuthalIntegrationSettings(azint_settings);
            UpdateAzint_i(dataset.get());
            emit datasetLoaded(dataset);
        }

        auto end = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
        logger.Info("Loaded file {} in {} ms", filename.toStdString(), duration);

        LoadImage_i(image_number, summation);
    } catch (std::exception &e) {
        logger.Error("Error loading file {} {}", filename.toStdString(), e.what());
    }
}

void JFJochImageReadingWorker::CloseFile() {
    QMutexLocker ul(&m);

    if (http_mode)
        http_reader.Close();
    else
        file_reader.Close();

    current_image_ptr.reset();
    current_image.reset();
    current_summation = 1;
    total_images = 0;
    current_file = "";
    emit imageLoaded({});
    emit datasetLoaded({});
}

void JFJochImageReadingWorker::LoadImage(int64_t image_number, int64_t summation) {
    QMutexLocker ul(&m);
    setAutoLoadMode_i(AutoloadMode::None);
    if ((image_number == current_image) && (current_summation == summation))
        return;
    LoadImage_i(image_number, summation);
}

void JFJochImageReadingWorker::UpdateAzint_i(const JFJochReaderDataset *dataset) {
    if (dataset) {
        azint_mapping = std::make_unique<AzimuthalIntegration>(curr_experiment, dataset->pixel_mask);
        index_and_refine = std::make_unique<IndexAndRefine>(curr_experiment, indexing.get());
        image_analysis = std::make_unique<MXAnalysisWithoutFPGA>(curr_experiment, *azint_mapping, dataset->pixel_mask,
            *index_and_refine.get());
    }
}

void JFJochImageReadingWorker::LoadImage_i(int64_t image_number, int64_t summation) {
    // Assumes m locked!
    try {
        if (summation <= 0 || image_number + summation > total_images)
            return;

        std::vector<int32_t> image;
        auto start = std::chrono::high_resolution_clock::now();
        if (http_mode) {
            if (image_number < 0 && summation != 1)
                return;
            current_image_ptr = http_reader.LoadImage(image_number, summation);
            total_images = http_reader.GetNumberOfImages();
            emit datasetLoaded(http_reader.GetDataset());
        } else {
            if (image_number < 0)
                return;
            current_image_ptr = file_reader.LoadImage(image_number, summation);
        }

        if (!current_image_ptr) {
            emit imageLoaded({});
            return;
        }

        current_image = current_image_ptr->ImageData().number;
        current_summation = summation;

        auto end = std::chrono::high_resolution_clock::now();

        if (auto_reanalyze)
            ReanalyzeImage_i();

        auto end_analysis = std::chrono::high_resolution_clock::now();

        auto duration_1 = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
        auto duration_2 = std::chrono::duration_cast<std::chrono::milliseconds>(end_analysis - end).count();

        logger.Info("Loaded image {} in {}/{} ms", image_number, duration_1, duration_2);
        emit imageNumberChanged(total_images, current_image.value());
        emit imageLoaded(current_image_ptr);
    } catch (std::exception &e) {
        logger.Error("Error loading image {}: {}", image_number, e.what());
    }
}

void JFJochImageReadingWorker::SetROIBox(QRect box) {
    QMutexLocker ul(&m);

    if (box.width() * box.height() == 0)
        roi.reset();

    roi = std::make_unique<ROIBox>("roi1", box.left(), box.right(), box.bottom(), box.top());
}

void JFJochImageReadingWorker::SetROICircle(double x, double y, double radius) {
    QMutexLocker ul(&m);

    if (radius <= 0)
        roi.reset();
    else
        roi = std::make_unique<ROICircle>("roi1", x, y, radius);
}

void JFJochImageReadingWorker::UpdateDataset_i(const std::optional<DiffractionExperiment> &experiment) {
    if (!current_image_ptr)
        return;

    std::shared_ptr<const JFJochReaderDataset> dataset;

    if (http_mode) {
        if (experiment)
            http_reader.UpdateGeomMetadata(experiment.value());
        dataset = http_reader.GetDataset();
    } else {
        if (experiment)
            file_reader.UpdateGeomMetadata(experiment.value());
        dataset = file_reader.GetDataset();
    }
    curr_experiment = dataset->experiment;
    curr_experiment.ImportIndexingSettings(indexing_settings);
    curr_experiment.ImportAzimuthalIntegrationSettings(azint_settings);
    UpdateAzint_i(dataset.get());

    emit datasetLoaded(dataset);

    current_image_ptr = std::make_shared<JFJochReaderImage>(current_image_ptr->ImageData(), dataset);

    if (auto_reanalyze)
        ReanalyzeImage_i();
    emit imageLoaded(current_image_ptr);
}

void JFJochImageReadingWorker::UpdateDataset(const DiffractionExperiment &experiment) {
    QMutexLocker ul(&m);
    UpdateDataset_i(experiment);
}

void JFJochImageReadingWorker::ReanalyzeImage_i() {
    if (!current_image_ptr || !azint_mapping || !image_analysis)
        return;

    auto start_time = std::chrono::high_resolution_clock::now();
    auto new_image = std::make_shared<JFJochReaderImage>(*current_image_ptr);
    auto new_image_dataset = new_image->CreateMutableDataset();

    new_image_dataset->experiment.ImportIndexingSettings(indexing_settings);
    new_image_dataset->experiment.ImportAzimuthalIntegrationSettings(azint_settings);
    new_image_dataset->az_int_bin_to_phi = azint_mapping->GetBinToPhi();
    new_image_dataset->az_int_bin_to_q = azint_mapping->GetBinToQ();
    new_image_dataset->azimuthal_bins = azint_mapping->GetAzimuthalBinCount();
    new_image_dataset->q_bins = azint_mapping->GetQBinCount();

    std::vector<uint8_t> buffer;
    AzimuthalIntegrationProfile azint_profile(*azint_mapping);
    image_analysis->Analyze(new_image->ImageData(), buffer, azint_profile, spot_finding_settings);

    current_image_ptr = new_image;
    auto end_time = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time);
    logger.Info("Analysis of image in {} ms", duration.count());
}

void JFJochImageReadingWorker::Analyze() {
    QMutexLocker locker(&m);
    ReanalyzeImage_i();
    emit imageLoaded(current_image_ptr);
}

void JFJochImageReadingWorker::FindCenter(const UnitCell& calibrant, bool guess) {
    QMutexLocker locker(&m);
    if (!current_image_ptr)
        return;
    logger.Info("Finding center");

    DiffractionGeometry geom = current_image_ptr->Dataset().experiment.GetDiffractionGeometry();
    try {
        if (guess)
            GuessGeometry(geom, current_image_ptr->ImageData().spots, calibrant);
        else
            OptimizeGeometry(geom, current_image_ptr->ImageData().spots, calibrant);
    } catch (const JFJochException &e) {
        logger.ErrorException(e);
        return;
    }

    logger.Info("Geometry found X: {} pxl Y: {} pxl Dist: {} mm", geom.GetBeamX_pxl(), geom.GetBeamY_pxl(),
                geom.GetDetectorDistance_mm());

    DiffractionExperiment new_experiment = current_image_ptr->Dataset().experiment;
    new_experiment.BeamX_pxl(geom.GetBeamX_pxl()).BeamY_pxl(geom.GetBeamY_pxl())
            .DetectorDistance_mm(geom.GetDetectorDistance_mm())
            .PoniRot1_rad(geom.GetPoniRot1_rad())
            .PoniRot2_rad(geom.GetPoniRot2_rad())
            .PoniRot3_rad(geom.GetPoniRot3_rad());
    UpdateDataset_i(new_experiment);

    std::vector<float> ring_Q = CalculateXtalRings(calibrant);

    QVector<float> rings;
    for (int i = 0; i < 15 && i < ring_Q.size(); i++) {
        rings.push_back(2 * M_PI / ring_Q[i]);
    }
    emit setRings(rings);
}

void JFJochImageReadingWorker::UpdateSpotFindingSettings(const SpotFindingSettings &settings,
                                                         const IndexingSettings &indexing,
                                                         int64_t max_spots) {
    QMutexLocker locker(&m);
    spot_finding_settings = settings;

    indexing_settings.Tolerance(indexing.GetTolerance());
    indexing_settings.ViableCellMinSpots(indexing.GetViableCellMinSpots());
    indexing_settings.IndexIceRings(indexing.GetIndexIceRings());
    indexing_settings.UnitCellDistTolerance(indexing.GetUnitCellDistTolerance());
    curr_experiment.ImportIndexingSettings(indexing_settings);
    curr_experiment.MaxSpotCount(max_spots);

    if (auto_reanalyze) {
        ReanalyzeImage_i();
        emit imageLoaded(current_image_ptr);
    }
}

void JFJochImageReadingWorker::ReanalyzeImages(bool input) {
    QMutexLocker locker(&m);
    auto_reanalyze = input;

    if (auto_reanalyze) {
        ReanalyzeImage_i();
        emit imageLoaded(current_image_ptr);
    }
}


void JFJochImageReadingWorker::UpdateAzintSettings(const AzimuthalIntegrationSettings &settings) {
    QMutexLocker locker(&m);
    azint_settings = settings;
    UpdateDataset_i(std::nullopt);
}

void JFJochImageReadingWorker::UpdateUserMask_i(const std::vector<uint32_t> &mask) {
    std::shared_ptr<const JFJochReaderDataset> dataset;
    if (http_mode) {
        http_reader.UpdateUserMask(mask);
        dataset = http_reader.GetDataset();
    } else {
        file_reader.UpdateUserMask(mask);
        dataset = file_reader.GetDataset();
    }
    UpdateAzint_i(dataset.get());

    emit datasetLoaded(dataset);

    current_image_ptr = std::make_shared<JFJochReaderImage>(current_image_ptr->ImageData(), dataset);

    if (current_image.has_value())
        LoadImage_i(current_image.value(), current_summation);
}


void JFJochImageReadingWorker::AddROIToUserMask() {
    QMutexLocker locker(&m);

    if (!roi || !current_image_ptr)
        return;

    auto user_mask = current_image_ptr->Dataset().pixel_mask.GetUserMask();

    int64_t width = current_image_ptr->Dataset().experiment.GetXPixelsNum();
    int64_t height = current_image_ptr->Dataset().experiment.GetYPixelsNum();
    const auto res = azint_mapping->Resolution();
    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            if (roi->CheckROI(x, y, 0))
                user_mask[x + y * width] = 1;
        }
    }

    UpdateUserMask_i(user_mask);
}

void JFJochImageReadingWorker::SubtractROIFromUserMask() {
    QMutexLocker locker(&m);
    if (!roi || !current_image_ptr)
        return;

    auto user_mask = current_image_ptr->Dataset().pixel_mask.GetUserMask();

    int64_t width = current_image_ptr->Dataset().experiment.GetXPixelsNum();
    int64_t height = current_image_ptr->Dataset().experiment.GetYPixelsNum();

    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            if (roi->CheckROI(x, y, 0))
                user_mask[x + y * width] = 0;
        }
    }

    UpdateUserMask_i(user_mask);
}

void JFJochImageReadingWorker::ClearUserMask() {
    QMutexLocker locker(&m);
    if (!roi || !current_image_ptr)
        return;

    auto user_mask = std::vector<uint32_t>(current_image_ptr->Dataset().experiment.GetXPixelsNum() * current_image_ptr->Dataset().experiment.GetYPixelsNum(), 0);
    UpdateUserMask_i(user_mask);
}

void JFJochImageReadingWorker::SaveUserMaskTIFF(QString filename) {
    QMutexLocker locker(&m);
    if (!current_image_ptr)
        return;
    auto user_mask = current_image_ptr->Dataset().pixel_mask.GetUserMask();
    CompressedImage mask_image(user_mask, current_image_ptr->Dataset().experiment.GetXPixelsNum(), current_image_ptr->Dataset().experiment.GetYPixelsNum());
    WriteTIFFToFile(filename.toStdString(), mask_image);
}

void JFJochImageReadingWorker::UploadUserMask() {
    QMutexLocker locker(&m);
    if (!current_image_ptr)
        return;

    auto user_mask = current_image_ptr->Dataset().pixel_mask.GetUserMask();
    if (http_mode)
        http_reader.UploadUserMask(user_mask);
}

void JFJochImageReadingWorker::LoadCalibration(QString dataset) {
    if (!http_mode) {
        auto tmp = std::make_shared<SimpleImage>();

        try {
            tmp->image = file_reader.ReadCalibration(tmp->buffer, dataset.toStdString());

            std::shared_ptr<const SimpleImage> ctmp = tmp;
            emit simpleImageLoaded(ctmp);
        } catch (const std::exception &e) {
            logger.Info("Error loading calibration: {}", e.what());
        }
    } else
        logger.Info("HTTP mode doesn't allow to read calibration (at the moment");
}

void JFJochImageReadingWorker::AutoLoadTimerExpired() {
    QMutexLocker locker(&m);
    switch (autoload_mode) {
        case AutoloadMode::HTTPSync:
            if (http_mode)
                LoadImage_i(-1 , 1);
            break;
        case AutoloadMode::Movie: {
            if (total_images == 0 || !current_image)
                return;
            int64_t new_image = (current_image.value() + jump_value) % total_images;
            LoadImage_i(new_image, current_summation);
            break;
        }
        case AutoloadMode::None:
            break;
    }
}

void JFJochImageReadingWorker::setAutoLoadMode_i(AutoloadMode in_mode) {
    autoload_mode = in_mode;
    if (autoload_mode == AutoloadMode::None)
        autoload_timer->stop();
    else
        autoload_timer->start();
    emit autoloadChanged(autoload_mode);
}

void JFJochImageReadingWorker::setAutoLoadMode(AutoloadMode mode) {
    QMutexLocker ul(&m);

    switch (mode) {
        case AutoloadMode::HTTPSync:
            if (http_mode)
                setAutoLoadMode_i(mode);
            else
                setAutoLoadMode_i(AutoloadMode::None);
            break;
        case AutoloadMode::Movie:
            setAutoLoadMode_i(mode);
            break;
        case AutoloadMode::None:
            setAutoLoadMode_i(mode);
            break;
    }
}

void JFJochImageReadingWorker::setAutoLoadJump(int64_t val) {
    QMutexLocker ul(&m);
    if (val > 0)
        jump_value = val;
}