/**
 * sliding_window_detector.h - 通用滑动窗口检测工具
 * 
 * 功能：
 * 1. 支持配置多窗口或自动计算窗口
 * 2. 从源图裁剪窗口、resize 到模型输入尺寸
 * 3. 检测结果坐标映射回原图
 * 4. NMS 合并多窗口结果
 * 
 * 用法：
 *   SlidingWindowDetector swd;
 *   swd.Init(config);  // 从配置初始化窗口
 *   auto windows = swd.GetWindows(src_w, src_h);
 *   for (auto& win : windows) {
 *       auto input = swd.PrepareInput(frame, win, model_w, model_h);
 *       // ... 推理 ...
 *       auto dets = swd.MapDetectionsToOriginal(raw_dets, win, model_w, model_h);
 *   }
 */

#pragma once

#include <algorithm>
#include <cstdint>
#include <cstring>
#include <vector>

#include "frame/frame.h"
#include "utils/simple_json.h"

namespace rk3588 {

// 窗口定义
struct DetectionWindow {
    int x = 0;
    int y = 0;
    int w = 640;
    int h = 640;
    
    bool IsValid() const { return w > 0 && h > 0; }
};

// 检测框（通用格式）
struct DetectionBox {
    float x, y, w, h;  // 左上角 + 宽高
    float confidence;
    int class_id;
};

/**
 * 滑动窗口检测器
 */
class SlidingWindowDetector {
public:
    SlidingWindowDetector() = default;
    
    /**
     * 从配置初始化
     * @param config SimpleJson 配置，支持 "windows" 数组
     * @return 是否成功
     */
    bool InitFromConfig(const SimpleJson& config) {
        windows_.clear();
        
        // 解析预配置窗口
        if (const SimpleJson* win_arr = config.Find("windows"); win_arr && win_arr->IsArray()) {
            for (const auto& w : win_arr->AsArray()) {
                if (w.IsObject()) {
                    DetectionWindow win;
                    win.x = w.ValueOr<int>("x", 0);
                    win.y = w.ValueOr<int>("y", 0);
                    win.w = w.ValueOr<int>("w", 640);
                    win.h = w.ValueOr<int>("h", 640);
                    if (win.IsValid()) {
                        windows_.push_back(win);
                    }
                }
            }
        }
        
        // 目标 resize 高度（用于自动计算窗口时）
        target_height_ = config.ValueOr<int>("target_height", 640);
        
        return true;
    }
    
    /**
     * 获取窗口列表（预配置或自动计算）
     * @param src_w 源图宽度
     * @param src_h 源图高度
     * @return 窗口列表
     */
    std::vector<DetectionWindow> GetWindows(int src_w, int src_h) const {
        if (!windows_.empty()) {
            return windows_;
        }
        return CalculateWindowsAuto(src_w, src_h);
    }
    
    /**
     * 准备模型输入
     * 从源图裁剪窗口区域，resize 到模型输入尺寸
     * 
     * @param frame 源帧
     * @param win 窗口定义
     * @param model_w 模型输入宽
     * @param model_h 模型输入高
     * @param output 输出缓冲区（model_w * model_h * 3）
     * @return 是否成功
     */
    bool PrepareInput(const FramePtr& frame, 
                      const DetectionWindow& win,
                      int model_w, int model_h,
                      uint8_t* output) const {
        if (!frame || !frame->data || !output) return false;
        
        const int src_w = frame->width;
        const int src_h = frame->height;
        
        // 获取源数据指针
        const uint8_t* src = frame->planes[0].data ? frame->planes[0].data : frame->data;
        const int src_stride = frame->planes[0].stride > 0 ? frame->planes[0].stride
                              : (frame->stride > 0 ? frame->stride : frame->width * 3);
        
        // 限制窗口在源图范围内
        int win_x = std::max(0, std::min(win.x, src_w - 1));
        int win_y = std::max(0, std::min(win.y, src_h - 1));
        int win_w = std::min(win.w, src_w - win_x);
        int win_h = std::min(win.h, src_h - win_y);
        
        if (win_w <= 0 || win_h <= 0) return false;
        
        // 裁剪窗口
        std::vector<uint8_t> crop_buf(static_cast<size_t>(win_w) * win_h * 3);
        for (int row = 0; row < win_h; ++row) {
            const uint8_t* src_row = src + (win_y + row) * src_stride + win_x * 3;
            uint8_t* dst_row = crop_buf.data() + row * win_w * 3;
            memcpy(dst_row, src_row, static_cast<size_t>(win_w) * 3);
        }
        
        // Resize 到模型输入尺寸
        ResizeRgbBilinear(crop_buf.data(), win_w, win_h, win_w * 3,
                          output, model_w, model_h, false);
        
        return true;
    }
    
    /**
     * 将检测结果从模型坐标映射回原图坐标
     * 
     * @param detections 模型输出的检测框（在 model_w x model_h 坐标系中）
     * @param win 窗口定义
     * @param model_w 模型输入宽
     * @param model_h 模型输入高
     * @return 映射后的检测框（在原图坐标系中）
     */
    std::vector<DetectionBox> MapDetectionsToOriginal(
            const std::vector<DetectionBox>& detections,
            const DetectionWindow& win,
            int model_w, int model_h) const {
        
        std::vector<DetectionBox> mapped = detections;
        
        float scale_x = static_cast<float>(win.w) / static_cast<float>(model_w);
        float scale_y = static_cast<float>(win.h) / static_cast<float>(model_h);
        
        for (auto& det : mapped) {
            det.x = win.x + det.x * scale_x;
            det.y = win.y + det.y * scale_y;
            det.w *= scale_x;
            det.h *= scale_y;
        }
        
        return mapped;
    }
    
    /**
     * 获取预配置窗口数量
     */
    size_t GetConfiguredWindowCount() const {
        return windows_.size();
    }

private:
    /**
     * 自动计算窗口（覆盖整个图像）
     * 策略：生成重叠的 640x640 窗口网格
     */
    std::vector<DetectionWindow> CalculateWindowsAuto(int src_w, int src_h) const {
        std::vector<DetectionWindow> windows;
        
        const int win_size = 640;
        
        // 计算步长（带重叠）
        int step_x = (src_w <= win_size) ? src_w : (src_w - win_size) / ((src_w + win_size - 1) / win_size - 1);
        int step_y = (src_h <= win_size) ? src_h : (src_h - win_size) / ((src_h + win_size - 1) / win_size - 1);
        
        if (step_x < win_size) step_x = win_size;
        if (step_y < win_size) step_y = win_size;
        
        for (int y = 0; y < src_h; y += step_y) {
            for (int x = 0; x < src_w; x += step_x) {
                DetectionWindow win;
                win.x = x;
                win.y = y;
                win.w = win_size;
                win.h = win_size;
                windows.push_back(win);
                
                if (x + win_size >= src_w) break;
            }
            if (y + win_size >= src_h) break;
        }
        
        return windows;
    }
    
    /**
     * RGB 图像双线性 resize
     * @param swap_rb 是否交换 R/B 通道
     */
    static void ResizeRgbBilinear(const uint8_t* src, int src_w, int src_h, int src_stride,
                                   uint8_t* dst, int dst_w, int dst_h, bool swap_rb) {
        const float scale_x = static_cast<float>(src_w) / dst_w;
        const float scale_y = static_cast<float>(src_h) / dst_h;
        
        for (int y = 0; y < dst_h; ++y) {
            float fy = y * scale_y;
            int y0 = static_cast<int>(fy);
            int y1 = std::min(y0 + 1, src_h - 1);
            float dy = fy - y0;
            
            for (int x = 0; x < dst_w; ++x) {
                float fx = x * scale_x;
                int x0 = static_cast<int>(fx);
                int x1 = std::min(x0 + 1, src_w - 1);
                float dx = fx - x0;
                
                // 双线性插值
                for (int c = 0; c < 3; ++c) {
                    int src_c = swap_rb ? (2 - c) : c;
                    
                    float v00 = src[(y0 * src_stride) + (x0 * 3) + src_c];
                    float v01 = src[(y0 * src_stride) + (x1 * 3) + src_c];
                    float v10 = src[(y1 * src_stride) + (x0 * 3) + src_c];
                    float v11 = src[(y1 * src_stride) + (x1 * 3) + src_c];
                    
                    float v0 = v00 * (1 - dx) + v01 * dx;
                    float v1 = v10 * (1 - dx) + v11 * dx;
                    float v = v0 * (1 - dy) + v1 * dy;
                    
                    dst[(y * dst_w + x) * 3 + c] = static_cast<uint8_t>(v);
                }
            }
        }
    }
    
    std::vector<DetectionWindow> windows_;
    int target_height_ = 640;
};

}  // namespace rk3588