#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 rtocr {

// ============================================================================
// Engine normalization constants (BGR channel order, matching PaddleOCR official)
// ============================================================================
// PaddleOCR processes images in BGR order (no BGR→RGB conversion).
// Engine applies: (pixel/255.0 - subVals[c]) / divVals[c] per channel.
// When feeding BGR input (no cvtColor), subVals/divVals indices map to:
//   [0]=B channel, [1]=G channel, [2]=R channel
//
// PaddleOCR config: mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
// These are applied to BGR channels: B=0.485/0.229, G=0.456/0.224, R=0.406/0.225

// Detection normalization (BGR order)
constexpr std::array<float, 3> kDetSubVals = { 0.485f, 0.456f, 0.406f };
constexpr std::array<float, 3> kDetDivVals = { 0.229f, 0.224f, 0.225f };

// Classifier normalization: PP-LCNet_x1_0_textline_ori uses ImageNet normalization (BGR order)
// Config: mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], scale=1/255
constexpr std::array<float, 3> kClsSubVals = { 0.485f, 0.456f, 0.406f };
constexpr std::array<float, 3> kClsDivVals = { 0.229f, 0.224f, 0.225f };

// Recognition normalization: (pixel/255 - 0.5) / 0.5 (same as classifier)
constexpr std::array<float, 3> kRecSubVals = { 0.5f, 0.5f, 0.5f };
constexpr std::array<float, 3> kRecDivVals = { 0.5f, 0.5f, 0.5f };

// ============================================================================
// Detection defaults (PP-OCRv5 server)
// kDetMaxSideLen is the default max side length for detection preprocessing.
// TRT engine auto-fallbacks to smaller max if GPU memory is insufficient during build.
// ============================================================================
constexpr int   kDetMaxSideLen    = 2560;
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;

// ============================================================================
// Data structures
// ============================================================================

// 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;
};

// ============================================================================
// Utility functions
// ============================================================================

// 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 (multiples of 32)
// limit_type='max': scale down if max side > maxSideLen (PP-OCRv5 server default)
// maxSideLimit: safety cap on final max dimension (default 4000)
inline cv::Size ComputeDetResizeShape(int srcH, int srcW, int maxSideLen,
                                       int maxSideLimit = kDetMaxSideLimit) {
    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);
    }

    newH = std::max(32, static_cast<int>(std::round(newH / 32.0) * 32));
    newW = std::max(32, static_cast<int>(std::round(newW / 32.0) * 32));
    return cv::Size(newW, newH);
}

// 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 rtocr
} // namespace ANSCENTER
