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

#ifndef JUNGFRAUJOCH_NEURALNETRESPREDICTOR_H
#define JUNGFRAUJOCH_NEURALNETRESPREDICTOR_H

#include <vector>
#include <mutex>
#include <condition_variable>

#include "../common/Logger.h"
#include "../common/DiffractionExperiment.h"

// Based on model described in:
// Mendez, D., Holton, J. M., Lyubimov, A. Y., Hollatz, S., Mathews, I. I., Cichosz, A., Martirosyan, V.,
// Zeng, T., Stofer, R., Liu, R., Song, J., McPhillips, S., Soltis, M. & Cohen, A. E. (2024).
// Acta Cryst. D80, 26-43.

enum class Quarter {TopLeft, TopRight, BottomLeft, BottomRight};

class NeuralNetInferenceClient {
    struct PredictorAddr {
        std::string hostname;
        uint16_t port;
        bool busy = false;
    };

    bool enable = false;

    std::vector<PredictorAddr> hosts;

    std::mutex m;
    std::condition_variable c;

    Logger *logger = nullptr;

    std::optional<float> Run(const std::vector<float> &input);

    std::optional<size_t> GetFreeNode();
    std::optional<size_t> AcquireNode();
    void ReturnNode(size_t node);

    template<class T>
    std::vector<float> PrepareInternal(const DiffractionExperiment& experiment, const T* image, Quarter q);
public:
    void AddHost(std::string hostname, uint16_t port);
    void AddHost(std::string addr);
    size_t GetHostCount();
    void AddLogger(Logger *logger);

    size_t GetMaxPoolFactor(const DiffractionExperiment& experiment) const;
    std::vector<float> Prepare(const DiffractionExperiment& experiment, const int16_t* image, Quarter q);
    std::vector<float> Prepare(const DiffractionExperiment& experiment, const int32_t* image, Quarter q);
    std::vector<float> Prepare(const DiffractionExperiment& experiment, const int8_t* image, Quarter q);

    std::optional<float> Inference(const DiffractionExperiment& experiment, const void* image, int nquads);
    std::optional<float> Inference(const DiffractionExperiment& experiment, const void* image, Quarter q);
};


#endif //JUNGFRAUJOCH_NEURALNETRESPREDICTOR_H