#include "PaddleOCRV5Engine.h"
#include "EPLoader.h"

#include <opencv2/imgproc.hpp>
#include <iostream>
#include <algorithm>

namespace ANSCENTER {
namespace onnxocr {

bool PaddleOCRV5Engine::Initialize(const std::string& detModelPath,
                                    const std::string& clsModelPath,
                                    const std::string& recModelPath,
                                    const std::string& dictPath,
                                    bool preferTensorRT) {
    std::lock_guard<std::recursive_mutex> lock(_mutex);
    ModelLoadingGuard mlg(_modelLoading);

    // High-perf options.  The OCR sub-models split into two groups:
    //
    //   1. Detector — its input shape varies continuously with every
    //      plate-ROI aspect ratio.  TRT EP is a poor fit because it
    //      builds a fresh engine for each unique shape (minutes each).
    //      We keep it on CUDA EP with the largest cuDNN workspace and
    //      let cuDNN HEURISTIC handle the per-shape algo selection.
    //
    //   2. Classifier + Recognizer — fixed-bucket shapes (cls is
    //      [1,3,80,160], rec is [1,3,48,{320,480,640,960}]).  These
    //      benefit massively from TRT EP because the engine is built
    //      once per shape and reused forever.
    OrtHandlerOptions detectorOpts;
    // Detector uses CUDA EP with *conservative* cuDNN workspace.
    // Empirical: on VRAM-constrained GPUs (LPD TRT engine + rec TRT
    // engine + ORT arena in play) the max-workspace mode causes cuDNN
    // to pick Winograd/implicit-precomp-GEMM variants that silently
    // fall back to slow NO-WORKSPACE algorithms when the big workspace
    // can't be allocated. With "0" cuDNN picks algorithms that are
    // known to fit and runs ~10x faster in practice.
    detectorOpts.useMaxCudnnWorkspace = false;
    detectorOpts.preferTensorRT       = false;   // never TRT for the detector

    // Classifier (fixed [1,3,80,160]): TRT with no profile is fine.
    OrtHandlerOptions classifierOpts;
    classifierOpts.useMaxCudnnWorkspace = true;
    classifierOpts.preferTensorRT       = preferTensorRT;
    classifierOpts.trtFP16              = true;

    // Recognizer: needs a DYNAMIC profile so one TRT engine covers every
    // (batch, bucket_width) pair we generate at runtime. Without this,
    // each new shape triggers a ~80s engine rebuild mid-stream when a
    // new plate appears or the plate count changes.
    //
    // Profile range:
    //   batch  : 1 .. 16       (16 plates worth of crops is generous)
    //   H      : 48 (fixed)
    //   W      : 320 .. 960    (covers all 4 recognizer buckets)
    //
    // Query the actual input name from the .onnx file instead of
    // hardcoding — PaddleOCR usually exports it as "x" but the name can
    // vary across model versions.
    OrtHandlerOptions recognizerOpts;
    recognizerOpts.useMaxCudnnWorkspace = true;
    recognizerOpts.preferTensorRT       = preferTensorRT;
    recognizerOpts.trtFP16              = true;
    if (preferTensorRT) {
        std::string recInputName = BasicOrtHandler::QueryModelInputName(recModelPath);
        if (recInputName.empty()) {
            std::cerr << "[PaddleOCRV5Engine] Could not query recognizer "
                         "input name — defaulting to 'x'" << std::endl;
            recInputName = "x";
        }
        std::cout << "[PaddleOCRV5Engine] Recognizer input name: '"
                  << recInputName << "' — building TRT dynamic profile "
                  << "[batch=1..16, W=320..960]" << std::endl;
        recognizerOpts.trtProfileMinShapes = recInputName + ":1x3x48x320";
        recognizerOpts.trtProfileOptShapes = recInputName + ":4x3x48x480";
        recognizerOpts.trtProfileMaxShapes = recInputName + ":16x3x48x960";
    }

    try {
        // Initialize detector (also triggers EPLoader init in BasicOrtHandler)
        detector_ = std::make_unique<ONNXOCRDetector>(detModelPath, detectorOpts);
        std::cout << "[PaddleOCRV5Engine] Detector initialized: " << detModelPath << std::endl;

        // Ensure this DLL's copy of Ort::Global<void>::api_ is initialized.
        // BasicOrtHandler sets it in ONNXEngine.dll, but each DLL has its own
        // inline-static copy. Without this, inference calls from ANSOCR.dll crash.
        if (Ort::Global<void>::api_ == nullptr) {
            Ort::InitApi(static_cast<const OrtApi*>(EPLoader::GetOrtApiRaw()));
        }

        // Initialize classifier (optional)
        if (!clsModelPath.empty()) {
            classifier_ = std::make_unique<ONNXOCRClassifier>(clsModelPath, classifierOpts);
            std::cout << "[PaddleOCRV5Engine] Classifier initialized: " << clsModelPath << std::endl;
        }
        else {
            classifier_.reset();
            std::cout << "[PaddleOCRV5Engine] Classifier skipped (no model path)" << std::endl;
        }

        // Initialize recognizer
        recognizer_ = std::make_unique<ONNXOCRRecognizer>(recModelPath, recognizerOpts);
        if (!recognizer_->LoadDictionary(dictPath)) {
            std::cerr << "[PaddleOCRV5Engine] Failed to load dictionary" << std::endl;
            return false;
        }
        std::cout << "[PaddleOCRV5Engine] Recognizer initialized: " << recModelPath << std::endl;

        // Pre-warm classifier (fixed [1,3,80,160]) and recognizer (4
        // bucket widths) so the first frame doesn't pay the cuDNN/TRT
        // algorithm-selection tax. The detector is intentionally NOT
        // warmed up: its input shape varies continuously with each
        // plate-ROI aspect ratio, so a warmup at any single canonical
        // shape would cost minutes (TRT) or be useless (CUDA cache miss
        // on the real frame anyway). Real frames will pay the per-shape
        // cuDNN HEURISTIC cost on first use.
        std::cout << "[PaddleOCRV5Engine] Warming up OCR pipeline..." << std::endl;
        if (classifier_) classifier_->Warmup();
        if (recognizer_) recognizer_->Warmup();
        std::cout << "[PaddleOCRV5Engine] Warmup complete" << std::endl;

        _initialized = true;
        std::cout << "[PaddleOCRV5Engine] Pipeline initialized successfully" << std::endl;
        return true;
    }
    catch (const std::exception& e) {
        std::cerr << "[PaddleOCRV5Engine] Initialization failed: " << e.what() << std::endl;
        detector_.reset();
        classifier_.reset();
        recognizer_.reset();
        _initialized = false;
        return false;
    }
}

std::vector<OCRPredictResult> PaddleOCRV5Engine::ocr(const cv::Mat& img) {
    if (_modelLoading.load()) return {};

    std::vector<OCRPredictResult> results;

    {
        auto lk = TryLockWithTimeout("PaddleOCRV5Engine::ocr");
        if (!lk.owns_lock()) return results;
        if (!_initialized || img.empty()) return results;
    }
    // _mutex released — heavy pipeline runs lock-free

    // Step 1: Text Detection
    auto boxes = detector_->Detect(img, _maxSideLen, _detDbThresh, _detBoxThresh, _detUnclipRatio, _useDilation);

    if (boxes.empty()) {
        return results;
    }

    // Step 2: Crop detected text regions
    std::vector<cv::Mat> croppedImages;
    croppedImages.reserve(boxes.size());
    for (auto& box : boxes) {
        cv::Mat cropped = GetRotateCropImage(img, box);
        if (!cropped.empty()) {
            croppedImages.push_back(cropped);
        }
    }

    // Step 3: Classification (optional)
    std::vector<int> cls_labels(croppedImages.size(), 0);
    std::vector<float> cls_scores(croppedImages.size(), 0.0f);

    if (classifier_) {
        classifier_->Classify(croppedImages, cls_labels, cls_scores, _clsThresh);

        // Rotate images classified as upside-down (label=1 and score > threshold)
        for (size_t i = 0; i < croppedImages.size(); i++) {
            if (cls_labels[i] % 2 == 1 && cls_scores[i] > _clsThresh) {
                cv::rotate(croppedImages[i], croppedImages[i], cv::ROTATE_180);
            }
        }
    }

    // Step 4: Text Recognition
    auto textLines = recognizer_->RecognizeBatch(croppedImages);

    // Step 5: Combine results
    for (size_t i = 0; i < boxes.size() && i < textLines.size(); i++) {
        OCRPredictResult result;

        // Convert TextBox points to box format [[x0,y0], [x1,y1], [x2,y2], [x3,y3]]
        result.box.resize(4);
        for (int j = 0; j < 4; j++) {
            result.box[j] = {
                static_cast<int>(boxes[i].points[j].x),
                static_cast<int>(boxes[i].points[j].y)
            };
        }

        result.text      = textLines[i].text;
        result.score      = textLines[i].score;
        result.cls_label  = cls_labels[i];
        result.cls_score  = cls_scores[i];

        results.push_back(result);
    }

    return results;
}

TextLine PaddleOCRV5Engine::recognizeOnly(const cv::Mat& croppedImage) {
    if (_modelLoading.load()) return { "", 0.0f };
    {
        auto lk = TryLockWithTimeout("PaddleOCRV5Engine::recognizeOnly");
        if (!lk.owns_lock()) return { "", 0.0f };
        if (!_initialized || !recognizer_ || croppedImage.empty()) return { "", 0.0f };
    }
    return recognizer_->Recognize(croppedImage);
}

std::vector<TextLine> PaddleOCRV5Engine::recognizeMany(const std::vector<cv::Mat>& croppedImages) {
    if (_modelLoading.load()) return std::vector<TextLine>(croppedImages.size());
    {
        auto lk = TryLockWithTimeout("PaddleOCRV5Engine::recognizeMany");
        if (!lk.owns_lock()) return std::vector<TextLine>(croppedImages.size());
        if (!_initialized || !recognizer_ || croppedImages.empty()) {
            return std::vector<TextLine>(croppedImages.size());
        }
    }
    // Delegates to the bucketed, batched path in ONNXOCRRecognizer.
    return recognizer_->RecognizeBatch(croppedImages);
}

} // namespace onnxocr
} // namespace ANSCENTER