#include "detector.hpp"
#include <opencv2/imgproc.hpp>
#include <algorithm>
#include <cmath>
namespace cms {
using fangjia::Box;
Detector::Detector(std::filesystem::path const& path,fangjia::TrtLogger& logger,float confidence,int maxPeople):engine_(path,logger),confidence_(confidence),maxPeople_(maxPeople) {
    if(engine_.inputs().size()!=1||engine_.outputs().size()!=1) throw std::runtime_error("Expected YOLO raw ONNX output (one input/output)");
    input_=engine_.inputs()[0];output_=engine_.outputs()[0];
    auto d=engine_.shape(input_);
    if(d.nbDims!=4||d.d[0]!=1||d.d[1]!=3||d.d[2]<=0||d.d[3]<=0) throw std::runtime_error("Expected fixed YOLO NCHW input");
    height_=d.d[2];width_=d.d[3];
    engine_.prepare();
    d=engine_.shape(output_);
    if(d.nbDims!=3||d.d[0]!=1||d.d[1]!=84) throw std::runtime_error("Expected YOLO output [1,84,N]");
    if(engine_.dataType(input_)!=nvinfer1::DataType::kFLOAT || engine_.dataType(output_)!=nvinfer1::DataType::kFLOAT) throw std::runtime_error("Expected float32 engine I/O");
    preprocess_=std::make_unique<YoloPreprocess>(width_,height_);outputHost_.resize(engine_.bytes(output_)/sizeof(float));
}
std::vector<Box> Detector::infer(cv::Mat const& frame,cudaStream_t stream,int filterHeight) {
    double gain=std::min(double(width_)/frame.cols,double(height_)/frame.rows);
    int rw=int(std::nearbyint(frame.cols*gain)),rh=int(std::nearbyint(frame.rows*gain));
    int left=int(std::nearbyint((width_-rw)/2.0-0.1)),top=int(std::nearbyint((height_-rh)/2.0-0.1));
    preprocess_->enqueue(frame,static_cast<float*>(engine_.device(input_)),stream);
    if(!engine_.enqueue(stream)) throw std::runtime_error("YOLO inference failed");
    CUDA_CHECK(cudaMemcpyAsync(outputHost_.data(),engine_.device(output_),engine_.bytes(output_),cudaMemcpyDeviceToHost,stream));
    CUDA_CHECK(cudaStreamSynchronize(stream));
    int n=engine_.shape(output_).d[2];std::vector<Box> candidates;
    for(int i=0;i<n;++i) {
        int cls=0; float score=outputHost_[4*n+i];
        for(int c=1;c<80;++c) if(outputHost_[(4+c)*n+i]>score) {cls=c;score=outputHost_[(4+c)*n+i];}
        if(cls!=0||score<=confidence_||!std::isfinite(score)) continue;
        float x=outputHost_[i],y=outputHost_[n+i],w=outputHost_[2*n+i],h=outputHost_[3*n+i];
        if(!std::isfinite(x)||!std::isfinite(y)||!std::isfinite(w)||!std::isfinite(h)||w<=0||h<=0) continue;
        candidates.push_back({x-w/2,y-h/2,x+w/2,y+h/2,score,0});
    }
    std::stable_sort(candidates.begin(),candidates.end(),[](auto const&a,auto const&b){return a.score>b.score;});
    std::vector<Box> kept;
    for(auto b:candidates) {
        bool suppressed=false;
        for(auto const& a:kept) {
            float inter=std::max(0.F,std::min(a.x2,b.x2)-std::max(a.x1,b.x1))*std::max(0.F,std::min(a.y2,b.y2)-std::max(a.y1,b.y1));
            float uni=(a.x2-a.x1)*(a.y2-a.y1)+(b.x2-b.x1)*(b.y2-b.y1)-inter;
            if(inter/uni>0.30F) {suppressed=true;break;}
        }
        if(!suppressed) kept.push_back(b);
        if(kept.size()==300) break;
    }
    std::vector<Box> filtered;
    for(auto b:kept) {
        b.x1=std::clamp(float((b.x1-left)/gain),0.F,float(frame.cols));b.x2=std::clamp(float((b.x2-left)/gain),0.F,float(frame.cols));
        b.y1=std::clamp(float((b.y1-top)/gain),0.F,float(frame.rows));b.y2=std::clamp(float((b.y2-top)/gain),0.F,float(frame.rows));
        if((b.x1+b.x2)*0.5F>0.80F*frame.cols||b.y2-b.y1<0.20F*(filterHeight>0?filterHeight:frame.rows)) continue;
        filtered.push_back(b);
    }
    if(int(filtered.size())>maxPeople_) filtered.resize(maxPeople_);
    std::stable_sort(filtered.begin(),filtered.end(),[](auto const&a,auto const&b){return a.x1<b.x1;});
    return filtered;
}
}