#include "film_detector.hpp"

#include <openvino/openvino.hpp>

#include <algorithm>
#include <cmath>
#include <numeric>
#include <stdexcept>

namespace film {
namespace {

// reflect 填充的下标映射：不重复边缘像素，填充量超过边长时反复折返。
// 例：长度 5、位置 -1 -> 1，位置 5 -> 3，位置 -6 -> 2。
int reflect_index(int index, int length) {
    if (length == 1) return 0;
    const int period = 2 * length - 2;
    int wrapped = ((index % period) + period) % period;
    return wrapped < length ? wrapped : period - wrapped;
}

}  // namespace

std::vector<int> axis_starts(int length, int size, int overlap) {
    if (length <= size) return {0};
    const int stride = std::max(1, size - overlap);
    std::vector<int> starts;
    for (int start = 0; start <= length - size; start += stride) starts.push_back(start);
    if (starts.back() != length - size) starts.push_back(length - size);
    return starts;
}

std::vector<Tile> plan_tiles(int width, int height, int size, int overlap) {
    if (size <= 0) throw std::invalid_argument("size 必须为正");
    if (overlap < 0 || overlap >= size) throw std::invalid_argument("overlap 需在 [0, size) 内");

    Tile base;
    base.crop_width = std::min(width, size);
    base.crop_height = std::min(height, size);
    // 不足一片时居中放置：缺陷贴边时两侧都有镜像出来的上下文
    base.pad_left = (size - base.crop_width) / 2;
    base.pad_top = (size - base.crop_height) / 2;

    std::vector<Tile> tiles;
    for (int y : axis_starts(height, size, overlap)) {
        for (int x : axis_starts(width, size, overlap)) {
            Tile tile = base;
            tile.x = x;
            tile.y = y;
            tiles.push_back(tile);
        }
    }
    return tiles;
}

void cut_tile(const ImageView& image, const Tile& tile, int size, std::vector<uint8_t>& out) {
    const int channels = image.channels;
    out.resize(static_cast<size_t>(size) * size * channels);
    for (int row = 0; row < size; ++row) {
        // 片内坐标 -> 裁剪区内坐标；落在填充区时按 reflect 折返
        const int local_y = reflect_index(row - tile.pad_top, tile.crop_height);
        const int src_y = tile.y + local_y;
        for (int col = 0; col < size; ++col) {
            const int local_x = reflect_index(col - tile.pad_left, tile.crop_width);
            const int src_x = tile.x + local_x;
            const uint8_t* src = image.data +
                (static_cast<size_t>(src_y) * image.width + src_x) * channels;
            uint8_t* dst = out.data() + (static_cast<size_t>(row) * size + col) * channels;
            for (int c = 0; c < channels; ++c) dst[c] = src[c];
        }
    }
}

float iou(const Box& a, const Box& b) {
    const float left = std::max(a.x1, b.x1), top = std::max(a.y1, b.y1);
    const float right = std::min(a.x2, b.x2), bottom = std::min(a.y2, b.y2);
    if (right <= left || bottom <= top) return 0.f;
    const float overlap = (right - left) * (bottom - top);
    const float united = a.area() + b.area() - overlap;
    return united > 0.f ? overlap / united : 0.f;
}

float containment(const Box& inner, const Box& outer) {
    const float left = std::max(inner.x1, outer.x1), top = std::max(inner.y1, outer.y1);
    const float right = std::min(inner.x2, outer.x2), bottom = std::min(inner.y2, outer.y2);
    if (right <= left || bottom <= top) return 0.f;
    const float area = inner.area();
    return area > 0.f ? (right - left) * (bottom - top) / area : 0.f;
}

std::vector<Box> nms(std::vector<Box> boxes, float iou_threshold, bool per_class,
                     float contained_threshold) {
    std::vector<int> order(boxes.size());
    std::iota(order.begin(), order.end(), 0);
    std::stable_sort(order.begin(), order.end(),
                     [&](int a, int b) { return boxes[a].score > boxes[b].score; });

    std::vector<Box> kept;
    for (int index : order) {
        const Box& box = boxes[index];
        bool duplicate = false;
        for (const Box& other : kept) {
            if (per_class && other.class_id != box.class_id) continue;
            if (iou(box, other) > iou_threshold ||
                (contained_threshold > 0.f && containment(box, other) >= contained_threshold)) {
                duplicate = true;
                break;
            }
        }
        if (!duplicate) kept.push_back(box);
    }
    return kept;
}

struct Detector::Impl {
    ov::Core core;
    ov::CompiledModel compiled;
    mutable ov::InferRequest request;
    int size = 0;
    int channels = 3;
    int classes = 0;
    int anchors = 0;
};

Detector::Detector(const std::string& model_path, const std::string& device)
    : impl_(std::make_unique<Impl>()) {
    auto model = impl_->core.read_model(model_path);
    impl_->compiled = impl_->core.compile_model(model, device);
    impl_->request = impl_->compiled.create_infer_request();

    const ov::Shape input = impl_->compiled.input(0).get_shape();    // [1, 3, N, N]
    const ov::Shape output = impl_->compiled.output(0).get_shape();  // [1, 4+nc, anchors]
    impl_->channels = static_cast<int>(input[1]);
    impl_->size = static_cast<int>(input[2]);
    impl_->classes = static_cast<int>(output[1]) - 4;
    impl_->anchors = static_cast<int>(output[2]);
}

Detector::~Detector() = default;

int Detector::input_size() const { return impl_->size; }
int Detector::num_classes() const { return impl_->classes; }

std::vector<Box> Detector::infer_tile(const ImageView& tile, float conf,
                                      float iou_threshold) const {
    const int size = impl_->size;
    if (tile.width != size || tile.height != size) {
        throw std::invalid_argument("infer_tile 只接受与网络输入等大的图，请先用 plan_tiles 切片");
    }

    // 预处理：RGB、除以 255、BCHW。没有 mean/std 标准化。
    // 灰度输入补成三通道——训练数据里就有灰度图，Python 侧也是这样处理的。
    ov::Tensor input(ov::element::f32, ov::Shape{1, static_cast<size_t>(impl_->channels),
                                                 static_cast<size_t>(size),
                                                 static_cast<size_t>(size)});
    float* dst = input.data<float>();
    const size_t plane = static_cast<size_t>(size) * size;
    for (int row = 0; row < size; ++row) {
        for (int col = 0; col < size; ++col) {
            const uint8_t* src = tile.data +
                (static_cast<size_t>(row) * size + col) * tile.channels;
            const size_t offset = static_cast<size_t>(row) * size + col;
            for (int c = 0; c < impl_->channels; ++c) {
                const uint8_t value = tile.channels == 1 ? src[0] : src[c % tile.channels];
                dst[c * plane + offset] = static_cast<float>(value) / 255.f;
            }
        }
    }

    impl_->request.set_input_tensor(input);
    impl_->request.infer();
    const float* raw = impl_->request.get_output_tensor(0).data<const float>();

    // 输出 [1, 4+nc, anchors]：前 4 行是像素 xywh（可能越界，要裁剪），
    // 其余每行一个类别的分数，已过 sigmoid。end2end=false，必须自己做 NMS。
    const int anchors = impl_->anchors;
    std::vector<Box> candidates;
    for (int i = 0; i < anchors; ++i) {
        int best_class = 0;
        float best_score = 0.f;
        for (int c = 0; c < impl_->classes; ++c) {
            const float score = raw[(4 + c) * anchors + i];
            if (score > best_score) {
                best_score = score;
                best_class = c;
            }
        }
        if (best_score < conf) continue;

        const float cx = raw[0 * anchors + i], cy = raw[1 * anchors + i];
        const float w = raw[2 * anchors + i], h = raw[3 * anchors + i];
        Box box;
        box.class_id = best_class;
        box.score = best_score;
        box.x1 = std::clamp((cx - w / 2.f) / size, 0.f, 1.f);
        box.y1 = std::clamp((cy - h / 2.f) / size, 0.f, 1.f);
        box.x2 = std::clamp((cx + w / 2.f) / size, 0.f, 1.f);
        box.y2 = std::clamp((cy + h / 2.f) / size, 0.f, 1.f);
        candidates.push_back(box);
    }
    return nms(std::move(candidates), iou_threshold);
}

std::vector<Box> Detector::infer(const ImageView& image, float conf, float iou_threshold,
                                 int overlap) const {
    const int size = impl_->size;
    const auto tiles = plan_tiles(image.width, image.height, size, overlap);

    std::vector<uint8_t> patch;
    std::vector<Box> collected;
    for (const Tile& tile : tiles) {
        cut_tile(image, tile, size, patch);
        ImageView view{patch.data(), size, size, image.channels};

        for (const Box& box : infer_tile(view, conf, iou_threshold)) {
            // 片内归一化坐标 -> 原图像素 -> 裁掉落在镜像填充区的部分 -> 原图归一化
            const float px1 = box.x1 * size - tile.pad_left + tile.x;
            const float py1 = box.y1 * size - tile.pad_top + tile.y;
            const float px2 = box.x2 * size - tile.pad_left + tile.x;
            const float py2 = box.y2 * size - tile.pad_top + tile.y;

            const float cx1 = std::max(px1, 0.f), cy1 = std::max(py1, 0.f);
            const float cx2 = std::min(px2, static_cast<float>(image.width));
            const float cy2 = std::min(py2, static_cast<float>(image.height));
            if (cx2 <= cx1 || cy2 <= cy1) continue;   // 整个落在镜像区，丢弃

            Box mapped;
            mapped.class_id = box.class_id;
            mapped.score = box.score;
            mapped.x1 = cx1 / image.width;
            mapped.y1 = cy1 / image.height;
            mapped.x2 = cx2 / image.width;
            mapped.y2 = cy2 / image.height;
            collected.push_back(mapped);
        }
    }
    // 重叠区会把同一处缺陷检出两次；被切开的缺陷还会留下小残框，靠包含比例压掉
    return nms(std::move(collected), iou_threshold, true, 0.8f);
}

}  // namespace film