#include "door_risk.hpp"
#include <opencv2/imgproc.hpp>
#include <opencv2/imgcodecs.hpp>
#include <algorithm>
#include <cmath>
#include <fstream>
#include <regex>
#include <stdexcept>
namespace cms {
double rounded(double v, int d) { double s=std::pow(10.0,d); return std::nearbyint(v*s)/s; }
static std::string lower(std::string s) { for(auto& c:s) if(static_cast<unsigned char>(c)<128) c=static_cast<char>(std::tolower(c)); return s; }
static bool raster(std::filesystem::path const& p) { auto e=lower(p.extension().string()); return e==".png"||e==".bmp"||e==".tif"||e==".tiff"; }
DoorRisk::DoorRisk(DoorOptions o, cv::Size size):options_(std::move(o)),size_(size) {
    if(options_.constant.empty()==options_.directory.empty()) throw std::runtime_error("Provide exactly one constant mask or door-mask directory");
    if(options_.warning<0 || options_.release<=options_.warning) throw std::runtime_error("Risk thresholds require 0 <= warning < release");
    format_=options_.format;
    if(!options_.constant.empty()) {
        if(!std::filesystem::is_regular_file(options_.constant)) throw std::runtime_error("Missing mask: "+options_.constant.string());
        if(format_=="auto") format_=lower(options_.constant.extension().string())==".json"?"xanylabeling":"raster";
    } else {
        std::vector<std::filesystem::path> files;
        for(auto const& e:std::filesystem::recursive_directory_iterator(options_.directory)) if(e.is_regular_file()) files.push_back(e.path());
        std::sort(files.begin(),files.end());
        if(format_=="auto") format_=std::any_of(files.begin(),files.end(),[](auto const& p){return p.extension()==".json";})?"xanylabeling":"raster";
        std::regex digits("[0-9]+");
        for(auto const& p:files) {
            if(format_=="xanylabeling"?p.extension()!=".json":!raster(p)) continue;
            int index=-1; std::string stem=p.stem().string();
            for(std::sregex_iterator i(stem.begin(),stem.end(),digits),end;i!=end;++i) index=std::stoi(i->str());
            if(index<0 && format_=="xanylabeling") { std::ifstream in(p); auto j=Json::parse(in); if(j.contains("frame_index")) index=j["frame_index"].get<int>(); else if(j.contains("frameIndex")) index=j["frameIndex"].get<int>(); }
            if(index<0) continue;
            if(!paths_.emplace(index,p).second) throw std::runtime_error("Duplicate mask frame: "+std::to_string(index));
        }
        if(paths_.empty()) throw std::runtime_error("No indexed door masks in "+options_.directory.string());
    }
    if(format_!="raster" && format_!="xanylabeling") throw std::runtime_error("Unknown door mask format");
}
cv::Mat DoorRisk::rasterize(std::filesystem::path const& path) const {
    cv::Mat mask(size_,CV_8U,cv::Scalar(0));
    if(format_=="raster") {
        auto img=cv::imread(path.string(),cv::IMREAD_UNCHANGED);
        if(img.empty()) throw std::runtime_error("Cannot read mask: "+path.string());
        if(img.channels()>1) { std::vector<cv::Mat> ch; cv::split(img,ch); img=ch[0]; for(size_t i=1;i<std::min(ch.size(),size_t(3));++i) cv::max(img,ch[i],img); }
        if(img.size()!=size_) cv::resize(img,img,size_,0,0,cv::INTER_NEAREST);
        return img>0;
    }
    std::ifstream in(path); auto j=Json::parse(in);
    double w=j.contains("imageWidth")&&!j["imageWidth"].is_null()?j["imageWidth"].get<double>():size_.width;
    double h=j.contains("imageHeight")&&!j["imageHeight"].is_null()?j["imageHeight"].get<double>():size_.height;
    if(w==0) w=size_.width;
    if(h==0) h=size_.height;
    float sx=static_cast<float>(size_.width/std::max(w,1.0)),sy=static_cast<float>(size_.height/std::max(h,1.0));
    for(auto const& s:j.value("shapes",Json::array())) {
        if(!options_.labels.count(lower(s.value("label",std::string())))) continue;
        std::vector<cv::Point> points;
        for(auto const& p:s.value("points",Json::array())) points.emplace_back(static_cast<int>(std::nearbyint(p[0].get<float>()*sx)),static_cast<int>(std::nearbyint(p[1].get<float>()*sy)));
        auto type=lower(s.value("shape_type",std::string("polygon")));
        if(points.size()<2) continue;
        if(type=="rectangle") cv::rectangle(mask,points[0],points[1],255,-1);
        else if(type=="circle") cv::circle(mask,points[0],static_cast<int>(std::nearbyint(cv::norm(points[1]-points[0]))),255,-1);
        else if(points.size()>=3) cv::fillPoly(mask,std::vector<std::vector<cv::Point>>{points},255);
    }
    return mask;
}
void DoorRisk::load(int frame) {
    if(cached_) return;
    loadedPath_.clear(); mask_.release(); distance_.release();
    if(!options_.constant.empty()) loadedPath_=options_.constant;
    else { auto it=paths_.find(frame+options_.offset); if(it!=paths_.end()) loadedPath_=it->second; }
    if(loadedPath_.empty()) {
        status_="missing";
        if(options_.require) throw std::runtime_error("No door mask for frame "+std::to_string(frame));
        return;
    }
    mask_=rasterize(loadedPath_);
    status_=cv::countNonZero(mask_)?"loaded":"empty";
    if(status_=="loaded") cv::distanceTransform(mask_==0,distance_,cv::DIST_L2,cv::DIST_MASK_PRECISE);
    cached_=!options_.constant.empty();
}
Json DoorRisk::evaluate(int frame, Json& record) {
    load(frame);
    Json regions=Json::array(),fallbacks=Json::array(),nearest=nullptr;
    std::set<int> withHands;
    auto measure=[&](cv::Mat const& region, std::string const& source,int track,std::string const& side,std::string const& state) {
        Json m={{"status","door_mask_"+status_},{"risk_score",nullptr},{"risk_level","unavailable"},{"decision","EVIDENCE_UNAVAILABLE"},{"source",source},{"track_id",track},{"side",side}};
        if(status_=="loaded") {
            int area=cv::countNonZero(region); cv::Mat intersection; cv::bitwise_and(region,mask_,intersection); int overlap=cv::countNonZero(intersection);
            double distance=0; cv::minMaxLoc(distance_,&distance,nullptr,nullptr,nullptr,region);
            double scale=std::max(std::sqrt(double(area)),1.0),relative=distance/scale;
            double score=overlap?1:std::clamp(1-relative/options_.release,0.0,1.0);
            std::string level=overlap?"danger":relative<=options_.warning?"warning":relative<=options_.release?"watch":"clear";
            std::string decision=overlap?"DANGER_INTERSECTION":relative<=options_.warning?"WARNING_NEAR":relative<=options_.release?"WATCH":"CLEAR";
            m={{"status","calculated"},{"risk_score",rounded(score,5)},{"risk_level",level},{"decision",decision},{"source",source},{"track_id",track},{"side",side},{"state",state},{"region_area_px",area},{"door_overlap_px",overlap},{"region_overlap_ratio",rounded(double(overlap)/std::max(area,1),6)},{"min_distance_px",rounded(distance,3)},{"region_scale_px",rounded(scale,3)},{"distance_relative",rounded(relative,6)}};
            if(nearest.is_null() || m["risk_score"].get<double>()>nearest["risk_score"].get<double>()) nearest=m;
        }
        regions.push_back(m); return m;
    };
    auto const ids=record.value("track_ids",Json::array());
    auto trackFor=[&](int pos){return pos>=0&&pos<int(ids.size())?ids[pos].get<int>():pos+1;};
    if(!record.contains("global_hands")) record["global_hands"]=Json::array();
    for(auto& hand:record["global_hands"]) {
        int pos=hand.value("person_id",0)-1; std::vector<cv::Point> points;
        bool finite=true;
        for(auto const& p:hand.value("keypoints",Json::array())) {
            float x=p[0].get<float>(),y=p[1].get<float>();
            if(!std::isfinite(x)||!std::isfinite(y)) {finite=false;break;}
            points.emplace_back(static_cast<int>(std::nearbyint(x)),static_cast<int>(std::nearbyint(y)));
        }
        if(!finite||points.size()<3) continue;
        std::vector<cv::Point> hull; cv::convexHull(points,hull);
        cv::Mat region(size_,CV_8U,cv::Scalar(0)); cv::fillConvexPoly(region,hull,255);
        if(!cv::countNonZero(region)) continue;
        withHands.insert(pos);
        hand["risk"]=measure(region,"RTMW_HAND_KEYPOINTS",trackFor(pos),hand.value("side",std::string("unknown")),hand.value("state",std::string("unknown")));
    }
    auto const boxes=record.value("boxes",Json::array());
    for(int pos=0;pos<int(boxes.size());++pos) {
        if(withHands.count(pos)) continue;
        auto const& b=boxes[pos]; if(b.size()<4) continue;
        int x1=std::clamp(int(std::nearbyint(b[0].get<double>())),0,size_.width-1),y1=std::clamp(int(std::nearbyint(b[1].get<double>())),0,size_.height-1);
        int x2=std::clamp(int(std::nearbyint(b[2].get<double>())),0,size_.width-1),y2=std::clamp(int(std::nearbyint(b[3].get<double>())),0,size_.height-1);
        if(x1>x2) std::swap(x1,x2);
        if(y1>y2) std::swap(y1,y2);
        if(x2<=x1||y2<=y1) continue;
        cv::Mat region(size_,CV_8U,cv::Scalar(0)); cv::rectangle(region,{x1,y1},{x2,y2},255,-1);
        Json fallback={{"track_id",trackFor(pos)},{"side","body"},{"state","rtmw_hand_unavailable"},{"body_zone_box",b}};
        fallback["risk"]=measure(region,"YOLO_BODY_BOX_FALLBACK",trackFor(pos),"body","rtmw_hand_unavailable"); fallbacks.push_back(fallback);
    }
    record["body_fallbacks"]=fallbacks;
    Json result={{"status",status_!="loaded"?"door_mask_"+status_:nearest.is_null()?"person_region_missing":"calculated"}};
    for(auto const* key:{"risk_score","risk_level","decision","source","track_id","side"}) result[key]=nearest.is_null()?Json(nullptr):nearest[key];
    if(nearest.is_null()) {result["risk_level"]="unavailable";result["decision"]="EVIDENCE_UNAVAILABLE";}
    result["door_mask"]={{"status",status_},{"annotation_index",options_.constant.empty()?frame+options_.offset:frame},{"path",loadedPath_.empty()?Json(nullptr):Json(std::filesystem::absolute(loadedPath_).string())},{"format",format_},{"area_px",mask_.empty()?0:cv::countNonZero(mask_)}};
    result["warning_distance_relative"]=options_.warning; result["release_distance_relative"]=options_.release; result["regions"]=regions;
    return result;
}
}