#pragma once

#include <string>
#include <vector>
#include <array>
#include <fstream>
#include <algorithm>
#include <numeric>
#include <cmath>
#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>

namespace ANSCENTER {
namespace onnxocr {

// Detection normalization constants (BGR channel order, matching PaddleOCR official)
// PaddleOCR config: mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
// Applied directly to BGR channels WITHOUT BGR→RGB conversion:
//   Channel 0 (B) → mean=0.485, std=0.229
//   Channel 1 (G) → mean=0.456, std=0.224
//   Channel 2 (R) → mean=0.406, std=0.225
constexpr float kDetMean0 = 0.485f;  // B channel
constexpr float kDetMean1 = 0.456f;  // G channel
constexpr float kDetMean2 = 0.406f;  // R channel
constexpr float kDetStd0  = 0.229f;  // B channel
constexpr float kDetStd1  = 0.224f;  // G channel
constexpr float kDetStd2  = 0.225f;  // R channel
constexpr float kScale = 1.0f / 255.0f;

// Detection defaults (PP-OCRv5 server: limit_type=max, limit_side_len=960)
constexpr int   kDetMaxSideLen    = 960;
constexpr int   kDetMaxSideLimit  = 4000;  // Safety cap on max dimension
constexpr float kDetDbThresh      = 0.3f;
constexpr float kDetBoxThresh     = 0.6f;
constexpr float kDetUnclipRatio   = 1.5f;
constexpr int   kDetMaxCandidates = 1000;

// Classifier defaults (PP-LCNet_x1_0_textline_ori model)
// Input: [B, 3, 80, 160], ImageNet normalization, 2-class (0°/180°)
// Direct resize to 80x160 (no aspect ratio preservation)
constexpr int   kClsImageH   = 80;
constexpr int   kClsImageW   = 160;
constexpr float kClsThresh   = 0.9f;

// Recognition defaults
constexpr int kRecImgH    = 48;
constexpr int kRecImgW    = 320;   // Default rec width (PP-OCRv5 rec_image_shape[2]=320, min padded width)
constexpr int kRecImgMaxW = 960;   // Allow wide recognition input for long text lines
constexpr int kRecBatchSize = 6;

// Maximum crops submitted to the recognizer in a single ORT Run call.
// Two things must stay in sync with this value:
//   1. The TRT dynamic profile in PaddleOCRV5Engine::BuildNvidiaOcrOptions —
//      the profile's max-batch dimension is set from kRecMaxBatch so TRT
//      builds a single engine that handles everything up to this size.
//   2. The bucket-chunking loop in ONNXOCRRecognizer::RecognizeBatch —
//      bucket groups larger than this get sliced into multiple Run() calls
//      so we never exceed the profile and fall off the fast batched path.
//
// Raising it increases peak runtime VRAM (the TRT execution context
// allocates worst-case activation buffers), so keep it as low as is
// reasonable for your expected plate count per frame.
constexpr int kRecMaxBatch = 24;

// Runtime batch-padding ladder.  Real batch sizes get rounded UP to one of
// these values before the ORT Run() call, with zero-padded slots in the
// input tensor and the padded outputs discarded after CTC decode.  The goal
// is to bound the number of distinct (batch, bucketW) input shapes the
// recognizer ever sees: with this ladder × kRecBucketWidths there are at
// most 24 unique shapes total, small enough for OpenVINO/DirectML to keep
// in their per-shape kernel cache instead of recompiling mid-stream.
//
// Without this padding, OpenVINO recompiles the recognizer for every new
// (N, W) combination — measured at ~1.7 s per recompile on Intel Iris Xe,
// which is exactly what produced the 1500–1800 ms frame spikes seen in the
// OPENVINO_GPU video tests.
//
// The last entry MUST equal kRecMaxBatch so the whole ladder stays inside
// the TRT dynamic profile [batch=1..kRecMaxBatch].
constexpr int kRecBatchLadder[] = { 1, 2, 4, 8, 16, 24 };
constexpr int kRecNumBatchSizes = sizeof(kRecBatchLadder) / sizeof(kRecBatchLadder[0]);
static_assert(kRecBatchLadder[kRecNumBatchSizes - 1] == kRecMaxBatch,
              "Last kRecBatchLadder entry must equal kRecMaxBatch so that "
              "padded shapes stay inside the TRT profile.");

// Warmup batch ladder — a subset of kRecBatchLadder used at engine init to
// pre-compile the most common runtime shapes.  Kept smaller than the full
// ladder because each warmup Run on OpenVINO costs ~1.7 s of kernel-compile
// time, and {1,2,4,8} covers the realistic plate-count distribution for
// ALPR (>95% of frames have ≤8 detected plates).  Frames with 9+ plates
// will pay a one-time recompile spike on first occurrence of batch=16/24.
constexpr int kRecWarmupBatchLadder[] = { 1, 2, 4, 8 };
constexpr int kRecNumWarmupBatchSizes =
    sizeof(kRecWarmupBatchLadder) / sizeof(kRecWarmupBatchLadder[0]);

// A detected text box: 4 corner points (top-left, top-right, bottom-right, bottom-left)
struct TextBox {
    std::array<cv::Point2f, 4> points;
    float score = 0.0f;
};

// A single recognized text line
struct TextLine {
    std::string text;
    float score = 0.0f;
};

// OCR result matching PaddleOCR::OCRPredictResult format
struct OCRPredictResult {
    std::vector<std::vector<int>> box;  // 4 corner points [[x,y], ...]
    std::string text;
    float score    = -1.0f;
    float cls_score = 0.0f;
    int   cls_label = -1;
};

// Load character dictionary from file
inline std::vector<std::string> LoadDict(const std::string& dictPath) {
    std::vector<std::string> keys;
    std::ifstream file(dictPath);
    if (!file.is_open()) return keys;
    std::string line;
    while (std::getline(file, line)) {
        if (!line.empty() && line.back() == '\r') {
            line.pop_back();
        }
        keys.push_back(line);
    }
    // CTC blank token at index 0
    keys.insert(keys.begin(), "#");
    // Space at end
    keys.push_back(" ");
    return keys;
}

// Compute resize dimensions for detection model.
// limit_type='max': scale down if max side > maxSideLen (PP-OCRv5 server default)
// maxSideLimit: safety cap on final max dimension (default 4000)
//
// Each dimension is rounded UP to a multiple of kDetSizeBucket (96). The
// coarse granularity is deliberate: cuDNN HEURISTIC has to re-select
// convolution algorithms every time it sees a new input shape, and that
// selection costs ~100 ms per shape. With multiples of 32, a typical ALPR
// run produces 30+ unique detector shapes; with multiples of 96 that drops
// to 5–10, which cuDNN can cache and reuse for the rest of the video.
// 96 is divisible by the DBNet down-stride of 32, so feature-map sizes
// stay integer.
inline cv::Size ComputeDetResizeShape(int srcH, int srcW, int maxSideLen,
                                       int maxSideLimit = kDetMaxSideLimit) {
    constexpr int kDetSizeBucket = 96;

    float ratio = 1.0f;
    int maxSide = std::max(srcH, srcW);
    if (maxSide > maxSideLen) {
        ratio = static_cast<float>(maxSideLen) / static_cast<float>(maxSide);
    }
    int newH = static_cast<int>(srcH * ratio);
    int newW = static_cast<int>(srcW * ratio);

    // Safety cap: clamp if either dimension exceeds maxSideLimit
    if (std::max(newH, newW) > maxSideLimit) {
        float clampRatio = static_cast<float>(maxSideLimit) / static_cast<float>(std::max(newH, newW));
        newH = static_cast<int>(newH * clampRatio);
        newW = static_cast<int>(newW * clampRatio);
    }

    auto roundUpToBucket = [](int x) {
        return std::max(kDetSizeBucket,
                        ((x + kDetSizeBucket - 1) / kDetSizeBucket) * kDetSizeBucket);
    };
    newH = roundUpToBucket(newH);
    newW = roundUpToBucket(newW);
    return cv::Size(newW, newH);
}

// Normalize BGR float image to CHW BGR vector for detection
// BGR channel order preserved (matching PaddleOCR official - no BGR→RGB conversion)
inline std::vector<float> NormalizeAndPermute(const cv::Mat& img) {
    int h = img.rows;
    int w = img.cols;
    std::vector<float> result(3 * h * w);
    for (int y = 0; y < h; y++) {
        const float* row = img.ptr<float>(y);
        for (int x = 0; x < w; x++) {
            float b = row[x * 3 + 0];
            float g = row[x * 3 + 1];
            float r = row[x * 3 + 2];
            // BGR order: channel 0=B, 1=G, 2=R (matching PaddleOCR official)
            result[0 * h * w + y * w + x] = (b * kScale - kDetMean0) / kDetStd0;
            result[1 * h * w + y * w + x] = (g * kScale - kDetMean1) / kDetStd1;
            result[2 * h * w + y * w + x] = (r * kScale - kDetMean2) / kDetStd2;
        }
    }
    return result;
}

// Normalize for classifier and recognizer: (x/255 - 0.5) / 0.5
// BGR channel order preserved (matching PaddleOCR official - no BGR→RGB conversion)
inline std::vector<float> NormalizeAndPermuteCls(const cv::Mat& img) {
    int h = img.rows;
    int w = img.cols;
    std::vector<float> result(3 * h * w);
    for (int y = 0; y < h; y++) {
        const float* row = img.ptr<float>(y);
        for (int x = 0; x < w; x++) {
            float b = row[x * 3 + 0];
            float g = row[x * 3 + 1];
            float r = row[x * 3 + 2];
            // BGR order: channel 0=B, 1=G, 2=R (matching PaddleOCR official)
            result[0 * h * w + y * w + x] = (b * kScale - 0.5f) / 0.5f;
            result[1 * h * w + y * w + x] = (g * kScale - 0.5f) / 0.5f;
            result[2 * h * w + y * w + x] = (r * kScale - 0.5f) / 0.5f;
        }
    }
    return result;
}

// Sort text boxes from top to bottom, left to right
inline void SortTextBoxes(std::vector<TextBox>& boxes) {
    std::sort(boxes.begin(), boxes.end(),
        [](const TextBox& a, const TextBox& b) {
            if (std::abs(a.points[0].y - b.points[0].y) < 10.0f) {
                return a.points[0].x < b.points[0].x;
            }
            return a.points[0].y < b.points[0].y;
        });
}

// Get rotated and cropped image from text box polygon
inline cv::Mat GetRotateCropImage(const cv::Mat& srcImage, const TextBox& box) {
    auto pts = box.points;
    float width = static_cast<float>(std::max(
        cv::norm(pts[0] - pts[1]),
        cv::norm(pts[2] - pts[3])));
    float height = static_cast<float>(std::max(
        cv::norm(pts[0] - pts[3]),
        cv::norm(pts[1] - pts[2])));

    std::vector<cv::Point2f> srcPts = { pts[0], pts[1], pts[2], pts[3] };
    std::vector<cv::Point2f> dstPts = {
        {0, 0}, {width, 0}, {width, height}, {0, height}
    };

    cv::Mat M = cv::getPerspectiveTransform(srcPts, dstPts);
    cv::Mat cropped;
    cv::warpPerspective(srcImage, cropped, M,
        cv::Size(static_cast<int>(width), static_cast<int>(height)),
        cv::BORDER_REPLICATE);

    if (cropped.rows > cropped.cols * 1.5f) {
        cv::Mat rotated;
        cv::transpose(cropped, rotated);
        cv::flip(rotated, rotated, 0);
        return rotated;
    }
    return cropped;
}

// Resize recognition image to fixed height, proportional width
inline cv::Mat ResizeRecImage(const cv::Mat& img, int targetH, int maxW) {
    float ratio = static_cast<float>(targetH) / img.rows;
    int targetW = static_cast<int>(img.cols * ratio);
    targetW = std::min(targetW, maxW);
    targetW = std::max(targetW, 1);

    cv::Mat resized;
    cv::resize(img, resized, cv::Size(targetW, targetH));
    return resized;
}

} // namespace onnxocr
} // namespace ANSCENTER
