// Copyright (2019-2022) Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-or-later

#include <catch2/catch.hpp>
#include "../common/ZMQPreviewPublisher.h"
#include "../common/jsonToGrpc.h"

TEST_CASE("ZMQPreviewPublisher","[ZMQ]") {
    ZMQContext context;
    ZMQPreviewPublisher publisher(context, "inproc://#5");

    ZMQSocket socket(context, ZMQSocketType::Sub);
    socket.Connect("inproc://#5");
    socket.SubscribeAll();

    DiffractionExperiment experiment(DetectorGeometry(1, 1, 0, 0, false));

    JFCalibration calibration(experiment);
    publisher.Start(experiment, calibration);

    std::vector<int16_t> image(experiment.GetPixelsNum());

    // Predictable random number generator
    std::mt19937 g1(19876);
    std::uniform_int_distribution<int16_t> distribution(-200,25000);
    for (auto &i: image)
        i = distribution(g1);

    std::vector<SpotToSave> spots;
    spots.push_back(SpotToSave{.x = 7, .y = 8, .intensity = 34, .indexed = false});
    spots.push_back(SpotToSave{.x = 37, .y = 48, .intensity = 123, .indexed = true});

    DataMessage message{.number = 564, .spots = spots};
    publisher.Publish(experiment, image.data(), message);

    std::string s;

    // Pixel mask
    REQUIRE(socket.Receive(s, false) > 0);
    JFJochProtoBuf::PreviewFrame frame;
    REQUIRE_NOTHROW(frame = jsonToGrpc<JFJochProtoBuf::PreviewFrame>(s));
    REQUIRE(frame.pixel_depth() == 4);
    REQUIRE(frame.image_number() == -1);

    // Frame
    REQUIRE(socket.Receive(s, false) > 0);

    REQUIRE_NOTHROW(frame = jsonToGrpc<JFJochProtoBuf::PreviewFrame>(s));

    REQUIRE(frame.pixel_depth() == 2);
    REQUIRE(frame.image_number() == 564);
    std::vector<char> image_out = {frame.data().begin(), frame.data().end()};

    REQUIRE(memcmp(image.data(), image_out.data(), experiment.GetPixelsNum() * experiment.GetPixelDepth()) == 0);
    REQUIRE(frame.spots_size() == 2);
    REQUIRE(frame.spots(0).x() == 7);
    REQUIRE(!frame.spots(0).indexed());
    REQUIRE(frame.spots(1).y() == 48);
    REQUIRE(frame.spots(1).indexed());

    REQUIRE(socket.Receive(s, false) < 0);
}

TEST_CASE("ZMQPreviewPublisher_FrameNumbers","[ZMQ]") {
    ZMQContext context;
    ZMQPreviewPublisher publisher(context, "inproc://#5");

    ZMQSocket socket(context, ZMQSocketType::Sub);
    socket.Connect("inproc://#5");
    socket.SubscribeAll();

    DiffractionExperiment experiment(DetectorGeometry(1, 1, 0, 0, false));
    experiment.PreviewPeriod(std::chrono::milliseconds(10)).ImageTimeUs(std::chrono::milliseconds(1));
    REQUIRE(experiment.GetPreviewStride() == 10);

    JFCalibration calibration(experiment);
    publisher.Start(experiment, calibration);

    std::vector<int16_t> image(experiment.GetPixelsNum());

    publisher.Publish(experiment, image.data(), DataMessage{.number = 0});
    publisher.Publish(experiment, image.data(), DataMessage{.number = 9});
    publisher.Publish(experiment, image.data(), DataMessage{.number = 12});
    publisher.Publish(experiment, image.data(), DataMessage{.number = 10});
    publisher.Publish(experiment, image.data(), DataMessage{.number = 25});
    publisher.Publish(experiment, image.data(), DataMessage{.number = 20});
    publisher.Publish(experiment, image.data(), DataMessage{.number = 18});

    std::string s;
    JFJochProtoBuf::PreviewFrame frame;

    // Pixel mask
    REQUIRE(socket.Receive(s, false) > 0);

    // Frame
    REQUIRE(socket.Receive(s, false) > 0);
    REQUIRE_NOTHROW(frame = jsonToGrpc<JFJochProtoBuf::PreviewFrame>(s));
    REQUIRE(frame.image_number() == 0);
    std::vector<char> image_out = {frame.data().begin(), frame.data().end()};

    REQUIRE(socket.Receive(s, false) > 0);

    REQUIRE_NOTHROW(frame = jsonToGrpc<JFJochProtoBuf::PreviewFrame>(s));
    REQUIRE(frame.image_number() == 12);

    REQUIRE(socket.Receive(s, false) > 0);

    REQUIRE_NOTHROW(frame = jsonToGrpc<JFJochProtoBuf::PreviewFrame>(s));
    REQUIRE(frame.image_number() == 25);

    REQUIRE(socket.Receive(s, false) < 0);
}

TEST_CASE("ZMQPreviewPublisher_GetPreviewImage","") {
    ZMQContext context;
    ZMQPreviewPublisher publisher(context, "inproc://#5");

    DiffractionExperiment experiment(DetectorGeometry(1, 1, 0, 0, false));

    JFCalibration calibration(experiment);
    publisher.Start(experiment, calibration);

    std::vector<int16_t> image(experiment.GetPixelsNum());

    // Predictable random number generator
    std::mt19937 g1(19876);
    std::uniform_int_distribution<int16_t> distribution(-200,25000);
    for (auto &i: image)
        i = distribution(g1);

    DataMessage message{.number = 564};
    publisher.Publish(experiment, image.data(), message);

    JFJochProtoBuf::PreviewFrame frame = publisher.GetPreviewImage();

    REQUIRE(frame.pixel_depth() == 2);
    REQUIRE(frame.image_number() == 564);
    std::vector<char> image_out = {frame.data().begin(), frame.data().end()};

    REQUIRE(memcmp(image.data(), image_out.data(), experiment.GetPixelsNum() * experiment.GetPixelDepth()) == 0);
}