#pragma once
#include <algorithm>
#include <array>
#include <cmath>
#include <stdexcept>
#include <vector>

namespace contextra {
// POSEPATH-01: exact legacy checkpoint [0,17], keypoints [Q,34,8]; not COCO Det slots.
struct PoseDetection {
    float score;
    std::array<float,4> box;
    std::array<std::array<float,3>,17> joints;
};
inline double poseSigmoid(double x) {
    return x >= 0 ? 1/(1+std::exp(-x)) : std::exp(x)/(1+std::exp(x));
}
inline double logSum(const std::vector<double>& values) {
    const double m=*std::max_element(values.begin(),values.end());
    double sum=0;
    for (double v:values) sum+=std::exp(v-m);
    return m+std::log(sum);
}
inline std::vector<PoseDetection> decodePose(const float* boxes,const float* logits,const float* kp) {
    if (!boxes || !logits || !kp) throw std::invalid_argument("missing Pose output");
    for (int i=0;i<400;++i) if (!std::isfinite(boxes[i])) throw std::invalid_argument("nonfinite box");
    for (int i=0;i<200;++i) if (!std::isfinite(logits[i])) throw std::invalid_argument("nonfinite label");
    for (int i=0;i<100*34*8;++i) if (!std::isfinite(kp[i])) throw std::invalid_argument("nonfinite joint");
    std::vector<int> ranked(200);
    for (int i=0;i<200;++i) ranked[i]=i;
    std::stable_sort(ranked.begin(),ranked.end(),[&](int a,int b){
        // Official selection sorts float32 sigmoid, including saturation ties.
        return static_cast<float>(poseSigmoid(logits[a])) > static_cast<float>(poseSigmoid(logits[b]));
    });
    std::vector<PoseDetection> out;
    for (int n=0;n<100;++n) {
        int index=ranked[n];
        if (index%2!=1) continue;
        int q=index/2;
        const float* b=boxes+q*4;
        PoseDetection p{};
        auto clamp=[](float x){return std::max(0.f,std::min(1.f,x));};
        p.box={{clamp(b[0]-b[2]/2),clamp(b[1]-b[3]/2),clamp(b[0]+b[2]/2),clamp(b[1]+b[3]/2)}};
        if (p.box[2]<=p.box[0] || p.box[3]<=p.box[1]) continue;
        std::vector<double> weighted,weights;
        for (int j=0;j<17;++j) {
            const float* k=kp+(q*34+17+j)*8;
            // Preserve out-of-frame coordinates; consumers decide geometric usability.
            p.joints[j]={{k[0],k[1],static_cast<float>(poseSigmoid(k[2]))}};
            double a=-2.*k[4],b2=-2.*k[6];
            double logw=std::log(std::max(poseSigmoid(k[2]),1e-12));
            weighted.push_back(logSum({a,b2,2*std::log(std::max(std::abs(double(k[5])),1e-12))+a+b2})+logw);
            weights.push_back(logw);
        }
        double logsig=logits[index]>=0 ? -std::log1p(std::exp(-double(logits[index])))
            : double(logits[index])-std::log1p(std::exp(double(logits[index])));
        p.score=std::min(std::nextafter(1.f,0.f),float(poseSigmoid(logsig-.2*(logSum(weighted)-logSum(weights)))));
        out.push_back(p);
    }
    return out;
}
}