#include <cassert>
#include <limits>
#include "detection_decode.hpp"

int main()
{
    using namespace contextra;
    const DetectionPolicy policy{};
    const float box[] = {0.5f, 0.5f, 0.4f, 0.4f};
    const float missing = -std::numeric_limits<float>::infinity();
    std::vector<float> logits(91, missing);
    logits[13] = logits[90] = 2;
    auto out = decodeDetections(box, logits.data(), 1, 100, 100, policy);
    assert(out.size() == 2);  // two classes from one query; ties use flat index
    assert(out[0].classId == 11 && out[1].classId == 79);
    assert(std::fabs(out[0].score - 0.88079708f) < 1e-6f);
    assert(std::fabs(out[0].left - 30) < 1e-5f);
    assert(sigmoid(-1000) == 0 && sigmoid(1000) == 1 && sigmoid(0) == 0.5f);
    assert(decodeDetections(nullptr, logits.data(), 1, 100, 100, policy).empty());
    assert(decodeDetections(box, logits.data(), 0, 100, 100, policy).empty());
    assert(decodeDetections(box, logits.data(), 1, 0, 100, policy).empty());
    logits.assign(91, missing);
    logits[1] = 0;
    assert(decodeDetections(box, logits.data(), 1, 100, 100, policy).size() == 1);
    assert(decodeDetections(box, logits.data(), 1, 100, 100, {true, 0.5f}).empty());
    logits[1] = -10;
    assert(decodeDetections(box, logits.data(), 1, 100, 100, policy).size() == 1);
    logits[1] = std::numeric_limits<float>::quiet_NaN();
    logits[2] = std::numeric_limits<float>::infinity();
    assert(decodeDetections(box, logits.data(), 1, 100, 100, policy).empty());
    logits[1] = 2;
    const float edge[] = {1.2f, 0.5f, 0.6f, 0.4f};
    out = decodeDetections(edge, logits.data(), 1, 100, 100, policy);
    assert(out.size() == 1 && std::fabs(out[0].left - 90) < 1e-4f);
    assert(std::fabs(out[0].width - 10) < 1e-4f);
    const float huge[] = {0.5f, 0.5f, 4, 4};
    out = decodeDetections(huge, logits.data(), 1, 100, 100, policy);
    assert(out.size() == 1 && out[0].width == 100 && out[0].height == 100);
    const float invalid[] = {0,0,-1,1};
    assert(decodeDetections(invalid, logits.data(), 1, 100, 100, policy).empty());
    const float nonfinite[] = {0,0,std::numeric_limits<float>::infinity(),1};
    assert(decodeDetections(nonfinite, logits.data(), 1, 100, 100, policy).empty());
    const float tiny[] = {0.5f,0.5f,0.005f,0.005f};
    assert(decodeDetections(tiny, logits.data(), 1, 100, 100, policy).size() == 1);

    std::vector<float> boxes(4*4);
    for (int i=0; i<4; ++i) std::copy(box, box+4, boxes.begin()+i*4);
    logits.assign(4*91, 2);
    out = decodeDetections(boxes.data(), logits.data(), 4, 100, 100, policy);
    // First 300 slots selected before mapping: 3*80 + slots 1..26 (24 valid).
    assert(out.size() == 264);
    assert(out[0].classId == 0 && out[79].classId == 79 && out[80].classId == 0);
    return 0;
}