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

#include <cmath>
#include <utility>
#include "NeuralNetInferenceClient.h"

#include "../common/JFJochException.h"
#include "httplib/httplib.h"

void NeuralNetInferenceClient::AddLogger(Logger *in_logger) {
    logger = in_logger;
}

size_t NeuralNetInferenceClient::GetHostCount() {
    std::unique_lock ul(m);
    return hosts.size();
}

void NeuralNetInferenceClient::AddHost(std::string hostname, uint16_t port) {
    std::unique_lock ul(m);
    hosts.emplace_back(PredictorAddr{.hostname = hostname, .port = port, .busy = false});
    enable = true;
}

void NeuralNetInferenceClient::AddHost(std::string addr) {
    if (addr.empty()) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Host address cannot be empty");
    }

    size_t pos = addr.find(':');
    std::string hostname;
    uint16_t port = 8000; // Default port

    if (pos != std::string::npos) {
        hostname = addr.substr(0, pos);
        if (hostname.empty()) {
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Hostname cannot be empty");
        }

        std::string portStr = addr.substr(pos + 1);
        try {
            int portValue = std::stoi(portStr);
            if (portValue < 0 || portValue > 65535) {
                throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                      "Port must be in range 0-65535, got: " + portStr);
            }
            port = static_cast<uint16_t>(portValue);
        } catch (const std::invalid_argument&) {
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Invalid port number format: " + portStr);
        } catch (const std::out_of_range&) {
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Port number out of range: " + portStr);
        }
    } else {
        hostname = addr;
        if (hostname.empty()) {
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Hostname cannot be empty");
        }
    }

    AddHost(hostname, port);
}


template<class T>
std::vector<float> NeuralNetInferenceClient::PrepareInternal(const DiffractionExperiment& experiment,
                                                             const PixelMask& mask,
                                                             const T* image,
                                                             Quarter q) {
    std::vector<float> ret(512*512);

    int64_t pool_factor = GetMaxPoolFactor(experiment);
    int64_t xpixel = experiment.GetXPixelsNum();
    int64_t ypixel = experiment.GetYPixelsNum();
    int64_t start_x = std::lround(experiment.GetBeamX_pxl());
    int64_t start_y = std::lround(experiment.GetBeamY_pxl());

    for (int64_t y = 0; y < 512; y++) {
        int64_t y0;

        if ((q == Quarter::BottomLeft) || (q == Quarter::BottomRight))
            y0 = start_y + y * pool_factor;
        else
            y0 = start_y - (y + 1) * pool_factor + 1;

        int64_t max_yp = std::min(y0 + pool_factor, ypixel);
        if (y0 < 0)
            y0 = 0;

        for (int64_t x = 0; x < 512; x++) {
            int64_t x0;

            if ((q == Quarter::TopRight) || (q == Quarter::BottomRight))
                x0 = start_x + x * pool_factor;
            else
                x0 = start_x - (x + 1) * pool_factor + 1;

            int64_t max_xp = std::min(x0 + pool_factor, xpixel);
            if (x0 < 0)
                x0 = 0;

            int64_t val = 0.0;
            for (int64_t yp = y0; yp < max_yp; yp++) {
                for (int64_t xp = x0; xp < max_xp; xp++) {
                    int64_t pxl = image[yp * xpixel + xp];
                    if (mask.GetMask().at(yp * xpixel + xp) != 0)
                        pxl = INT64_MAX;
                    else if (pxl > INT16_MAX)
                        pxl = INT16_MAX;

                    if (pxl > val)
                        val = pxl;
                }
            }

            if (val == INT64_MAX)
                ret[512 * y + x] = 0;
            else
                ret[512 * y + x] = floor(sqrt(static_cast<double>(val)));
        }
    }

    return ret;
}

std::vector<float> NeuralNetInferenceClient::Prepare(const DiffractionExperiment& experiment,
                                                     const PixelMask& mask,
                                                     const int16_t *image, Quarter q) {
    return PrepareInternal(experiment, mask, image, q);
}

std::vector<float> NeuralNetInferenceClient::Prepare(const DiffractionExperiment& experiment,
                                                     const PixelMask& mask,
                                                     const int32_t *image, Quarter q) {
    return PrepareInternal(experiment, mask, image, q);
}

std::vector<float> NeuralNetInferenceClient::Prepare(const DiffractionExperiment& experiment,
                                                     const PixelMask& mask,
                                                     const int8_t *image, Quarter q) {
    return PrepareInternal(experiment, mask, image, q);
}

size_t NeuralNetInferenceClient::GetMaxPoolFactor(const DiffractionExperiment& experiment) const {
    float max_direction = std::max(experiment.GetXPixelsNum(), experiment.GetYPixelsNum()) / 2.0;

    size_t pool_factor = std::ceil(max_direction / 512.0f - 0.25);

    if (pool_factor <= 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Detector size is too small");

    if (pool_factor > 8)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Detector size is too large");
    return pool_factor;
}

std::optional<size_t> NeuralNetInferenceClient::GetFreeNode() {
    //Assumes m is locked!

    for (int i = 0; i < hosts.size(); i++) {
        if (!hosts[i].busy) {
            hosts[i].busy = true;
            return i;
        }
    }
    return {};
}

std::optional<size_t> NeuralNetInferenceClient::AcquireNode() {
    std::unique_lock ul(m);

    if (hosts.empty())
        return {};

    auto ret = GetFreeNode();
    if (ret)
        return ret;

    // No free node available, wait for one with a timeout
    c.wait_for(ul, std::chrono::milliseconds(200),
               [this, &ret]() {
                   ret = GetFreeNode();
                   return ret.has_value();
               });

    return ret;
}

void NeuralNetInferenceClient::ReturnNode(size_t node) {
    std::unique_lock ul(m);
    hosts[node].busy = false;
    c.notify_one();
}

std::optional<float> NeuralNetInferenceClient::Run(const std::vector<float> &input) {
    auto h = AcquireNode();

    if (!h.has_value())
        return {};

    if (logger)
        logger->Debug("Using host {:d}", h.value());

    httplib::Result resp;
    try {
        httplib::Client cli(hosts[h.value()].hostname, hosts[h.value()].port);
        cli.set_read_timeout(1);
        cli.set_write_timeout(1);
        resp = cli.Post("/infer", (char *) input.data(), input.size() * sizeof(float), "application/octet-stream");
    } catch (...) {
        ReturnNode(h.value());
    }
    ReturnNode(h.value());

    if (resp && resp->status == httplib::StatusCode::OK_200) {
        auto j = nlohmann::json::parse(resp->body);
        if (j.contains("result") && j["result"].is_number_float())
            return j["result"].get<float>();
    }
    return {};

}

std::optional<float>
NeuralNetInferenceClient::Inference(const DiffractionExperiment &experiment,
                                    const PixelMask& mask,
                                    const void *image,
                                    int nquads) {
    if (!enable)
        return {};

    std::optional<float> quad[4];

    if (nquads >= 1)
        quad[0] = Inference(experiment, mask, image, Quarter::BottomRight);
    if (nquads >= 2)
        quad[1] = Inference(experiment, mask, image, Quarter::BottomRight);
    if (nquads >= 3)
        quad[2] = Inference(experiment, mask, image, Quarter::BottomRight);
    if (nquads >= 4)
        quad[3] = Inference(experiment, mask, image, Quarter::BottomLeft);

    int count = 0;
    float sum = 0.0f;
    for (int i = 0; i < 4; i++) {
        if (quad[i]) {
            count++;
            sum += quad[i].value();
        }
    }
    if (count == 0)
        return {};
    return sum / count;
}

std::optional<float> NeuralNetInferenceClient::Inference(const DiffractionExperiment& experiment,
                                                         const PixelMask& mask,
                                                         const void *image,
                                                         Quarter q) {
    if (!enable)
        return {};

    std::vector<float> v;

    switch (experiment.GetByteDepthImage()) {
        case 1:
            if (experiment.IsPixelSigned())
                v = PrepareInternal(experiment, mask, (int8_t *) image, q);
            else
                v = PrepareInternal(experiment, mask, (uint8_t *) image, q);
            break;
        case 2:
            if (experiment.IsPixelSigned())
                v = PrepareInternal(experiment, mask, (int16_t *) image, q);
            else
                v = PrepareInternal(experiment, mask, (uint16_t *) image, q);
            break;
        case 4:
            if (experiment.IsPixelSigned())
                v = PrepareInternal(experiment, mask, (int32_t *) image, q);
            else
                v = PrepareInternal(experiment, mask, (uint32_t *) image, q);
            break;
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Bit depth not supported");
    }
    auto resp = Run(v);

    if (!resp.has_value())
        return {};

    float two_theta = atanf(
            ((2.0f * GetMaxPoolFactor(experiment) * experiment.GetPixelSize_mm()
              / experiment.GetDetectorDistance_mm()) * resp.value()));
    float stheta = sinf(two_theta * 0.5f);
    float resolution = experiment.GetWavelength_A() / (2.0f * stheta);
    return resolution;
}