// Copyright (2019-2024) Paul Scherrer Institute

#include "MXAnalyzer.h"
#include "CPUSpotFinder.h"
#include "../common/DiffractionGeometry.h"
#include "Regression.h"

double stddev(const std::vector<float> &v) {
    if (v.size() <= 1)
        return 0.0;

    double mean = 0.0f;

    for (const auto &i: v)
        mean += i;
    mean /= v.size();

    double stddev = 0.0f;
    for (const auto &i: v)
        stddev += (i - mean) * (i - mean);

    return sqrt(stddev / (v.size() - 1));
}


MXAnalyzer::MXAnalyzer(const DiffractionExperiment &in_experiment)
: experiment(in_experiment) {
    auto uc = experiment.GetUnitCell();
    if (uc) {
        do_indexing = true;
        indexer.Setup(uc.value());
    }
    if (experiment.IsSpotFindingEnabled())
        find_spots = true;

}

void MXAnalyzer::ReadFromFPGA(const DeviceOutput *output, const SpotFindingSettings &settings, size_t module_number) {
    if (!find_spots || !settings.enable)
        return;
    StrongPixelSet strong_pixel_set;
    strong_pixel_set.ReadFPGAOutput(experiment, *output);
    strong_pixel_set.FindSpots(experiment, settings, spots, module_number);
}

void MXAnalyzer::ReadFromCPU(const int16_t *image, const SpotFindingSettings &settings, size_t module_number) {
    if (!find_spots)
        return;
    if (experiment.GetPixelDepth() == 4)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "CPU spot finder simulation doesn't support 32-bit images");
    std::vector<float> d_map(RAW_MODULE_SIZE);

    DeviceOutput output{};
    memcpy(output.pixels, image, RAW_MODULE_SIZE * sizeof(int16_t));

    CalcSpotFinderResolutionMap(d_map.data(), experiment, module_number);

    FindSpots(output, settings, d_map.data());

    ReadFromFPGA(&output, settings, module_number);
}

void MXAnalyzer::Process(DataMessage &message, const SpotFindingSettings& settings) {
    message.indexing_result = 0;
    if (!find_spots)
        return;

    std::vector<DiffractionSpot> spots_out;
    FilterSpotsByCount(experiment, spots, spots_out);
    spots.clear();

    for (const auto &spot: spots_out)
        message.spots.push_back(spot);

    if (do_indexing && settings.indexing) {
        std::vector<Coord> recip;
        recip.reserve(spots_out.size());
        for (const auto &i: spots_out)
            recip.push_back(i.ReciprocalCoord(experiment));

        auto indexer_result = indexer.Run(recip, settings.indexing_tolerance);

        if (!indexer_result.empty()) {
            message.indexing_result = 1;
            CalculateBeamCenterAndDistance(spots_out,
                                           indexer_result[0],
                                           message);

            // identify indexed spots
            for (int i = 0; i < recip.size(); i++) {
                auto predicted_pos = RecipToDector(experiment, indexer_result[0].predicted_spots[i]);
                float x_diff = predicted_pos.first - spots_out[i].RawCoord().x;
                float y_diff = predicted_pos.second - spots_out[i].RawCoord().y;
                message.spots[i].indexed = (x_diff * x_diff + y_diff * y_diff
                                            < spot_distance_threshold_pxl * spot_distance_threshold_pxl);
            }

            indexer_result[0].l.Save(message.indexing_lattice);
            message.indexing_unit_cell = indexer_result[0].l.GetUnitCell();
        }
    }
}

void MXAnalyzer::CalculateBeamCenterAndDistance(const std::vector<DiffractionSpot> &spots,
                                                const IndexingResult& result,
                                                DataMessage &message) {
    if (spots.empty())
        return;
    assert(result.predicted_spots.size() == spots.size());

    std::vector<float> spot_x_position(spots.size());
    std::vector<float> spot_y_position(spots.size());
    std::vector<float> spot_x_Sd1_over_Sd3(spots.size());
    std::vector<float> spot_y_Sd2_over_Sd3(spots.size());

    for (int i = 0; i < spots.size(); i++) {
        spot_x_position[i] = spots[i].RawCoord().x * experiment.GetPixelSize_mm();
        spot_y_position[i] = spots[i].RawCoord().y * experiment.GetPixelSize_mm();

        Coord S = experiment.GetScatteringVector() + result.predicted_spots[i];
        if (S.z != 0.0f) {
            spot_x_Sd1_over_Sd3[i] = S.x / S.z;
            spot_y_Sd2_over_Sd3[i] = S.y / S.z;
        }
    }

    auto reg_x = regression(spot_x_Sd1_over_Sd3, spot_x_position);
    auto reg_y = regression(spot_y_Sd2_over_Sd3, spot_y_position);

    // Subtract 0.5 pixel, to account for differences in definition of pixel coordinate (center or top-left corner)
    message.corr_beam_x_pxl = experiment.GetBeamX_pxl() - reg_x.intercept / experiment.GetPixelSize_mm() - 0.5;
    message.corr_beam_y_pxl = experiment.GetBeamY_pxl() - reg_y.intercept / experiment.GetPixelSize_mm() - 0.5;
    message.corr_det_dist_mm = experiment.GetDetectorDistance_mm() - (reg_x.slope + reg_y.slope) / 2.0;
}