use anyhow::{Context, Result};
use image::{imageops::FilterType, DynamicImage, GenericImageView};
use std::fmt;
use tract_onnx::prelude::tract_ndarray::{prelude::*, s, Array2, Array3, Array4, Axis};
use tract_onnx::prelude::{Tensor};


use crate::models::det::session::DetSession;

#[derive(Debug, Clone, Copy)]
pub struct DetectionResult {
    pub x1: i32,
    pub y1: i32,
    pub x2: i32,
    pub y2: i32,
    pub score: f32,
    pub class_id: u32,
}

impl fmt::Display for DetectionResult {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // 结构体只管自己这一行怎么显示，不用管外部的索引 [i]
        write!(
            f,
            "x1={}, y1={}, x2={}, y2={}, 分数={:.4}, 类别ID={}",
            self.x1, self.y1, self.x2, self.y2, self.score, self.class_id
        )
    }
}

#[derive(Debug)]
pub struct Detector<'a> {
    pub(crate) session: &'a DetSession,
    #[allow(dead_code)]
    pub(crate) use_gpu: bool,
    #[allow(dead_code)]
    pub(crate) device_id: u8,
}

impl<'a> Detector<'a> {
    pub fn new(session: &'a DetSession) -> Self {
        Detector {
            session,
            use_gpu: false,
            device_id: 0,
        }
    }

    pub fn predict(&self, image: &DynamicImage) -> Result<Vec<DetectionResult>> {
        // Rust 中通常在调用层处理文件/PIL转换，这里直接进入核心逻辑
        self.get_bbox(image)
    }
    /// 2. preproc: 纯 Rust 实现 (替代 OpenCV)
    fn preproc(&self, image: &DynamicImage, input_size: (u32, u32)) -> Result<(Tensor, f32)> {
        let (target_h, target_w) = input_size;
        let (img_w, img_h) = image.dimensions();

        // 计算缩放比例 (Letterbox)
        let r = (target_h as f32 / img_h as f32).min(target_w as f32 / img_w as f32);
        let new_h = (img_h as f32 * r) as u32;
        let new_w = (img_w as f32 * r) as u32;

        // Resize 图像
        let resized = image.resize_exact(new_w, new_h, FilterType::Triangle);
        // 2. 关键：将 DynamicImage 显式转换为 RgbImage (Rgb<u8>)
        let resized_rgb = resized.to_rgb8();
        // 创建 114 灰度填充的背景
        let mut base_img =
            image::ImageBuffer::from_pixel(target_w, target_h, image::Rgb([114u8, 114, 114]));

        // 将 resize 后的图像覆盖到左上角 (类似于原始代码中的 padded_img[:h, :w])
        image::imageops::overlay(&mut base_img, &resized_rgb, 0, 0);

        // 优化：直接获取底层的扁平 raw buffer，比 enumerate_pixels() 快得多
        let raw_samples = base_img.as_flat_samples();
        let slice = raw_samples.as_slice();

        // 构造 NCHW Tensor
        let mut array = Array4::<f32>::zeros((1, 3, target_h as usize, target_w as usize));

        // 用连续的 stride 步长进行写入，提高 CPU 缓存利用率
        for y in 0..target_h as usize {
            for x in 0..target_w as usize {
                let idx = (y * target_w as usize + x) * 3;
                // BGR 赋值
                array[[0, 0, y, x]] = slice[idx + 2] as f32; // B
                array[[0, 1, y, x]] = slice[idx + 1] as f32; // G
                array[[0, 2, y, x]] = slice[idx] as f32; // R
            }
        }

        Ok((array.into(), r))
    }

    /// 3. demo_postprocess (逻辑与 Python 一致)
    fn demo_postprocess(&self, mut outputs: Array3<f32>, img_size: (i32, i32)) -> Array3<f32> {
        let strides = [8, 16, 32];

        // 遍历每一个 Batch（支持动态 Batch 推理）
        for mut batch in outputs.axis_iter_mut(Axis(0)) {
            let mut offset = 0;

            for &stride in &strides {
                // 计算当前特征图的尺寸
                let h = img_size.0 / stride;
                let w = img_size.1 / stride;
                let f_stride = stride as f32;

                for y in 0..h {
                    for x in 0..w {
                        // 计算当前格子在 25200 个锚点中的线性索引
                        let idx = offset + (y * w + x) as usize;
                        // 1. 还原中心点坐标 (cx, cy)
                        // 公式: (output + grid_offset) * stride
                        batch[[idx, 0]] = (batch[[idx, 0]] + x as f32) * f_stride;
                        batch[[idx, 1]] = (batch[[idx, 1]] + y as f32) * f_stride;

                        // 2. 还原宽高 (w, h)
                        // 公式: exp(output) * stride
                        batch[[idx, 2]] = batch[[idx, 2]].exp() * f_stride;
                        batch[[idx, 3]] = batch[[idx, 3]].exp() * f_stride;
                    }
                }
                // 移动到下一个步长的起始位置
                offset += (h * w) as usize;
            }
        }
        outputs
    }

    /// 4. nms
    fn nms(&self, boxes: &Array2<f32>, scores: &Array1<f32>, nms_thr: f32) -> Vec<usize> {
        let mut keep = Vec::new();
        let x1 = boxes.column(0);
        let y1 = boxes.column(1);
        let x2 = boxes.column(2);
        let y2 = boxes.column(3);
        // 在每一项前加上 &，并确保括号内的计算顺序
        // 注意：ndarray 的 View 运算需要 &view1 - &view2
        let areas = (&x2 - &x1 + 1.0) * (&y2 - &y1 + 1.0);

        // 初始排序索引
        let mut v: Vec<usize> = (0..scores.len()).collect();
        v.sort_unstable_by(|&i, &j| {
            scores[j]
                .partial_cmp(&scores[i])
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        // 我们不使用 v.remove(0)，而是直接通过索引池操作
        let mut active_indices = v;

        while !active_indices.is_empty() {
            // 取出当前池子中得分最高的框（即第一个元素）
            let i = active_indices[0];
            keep.push(i);

            // 如果池子里只剩一个了，直接结束
            if active_indices.len() == 1 {
                break;
            }

            // 5. 核心逻辑：使用 retain 一次性过滤掉：
            // (a) 当前框自己 (idx == i)
            // (b) 与当前框重叠度过高的框 (iou > nms_thr)
            active_indices.retain(|&idx| {
                // 如果是当前正在处理的框，不保留（因为它已经进入 keep 了）
                if idx == i {
                    return false;
                }

                // 计算 IoU
                let xx1 = x1[i].max(x1[idx]);
                let yy1 = y1[i].max(y1[idx]);
                let xx2 = x2[i].min(x2[idx]);
                let yy2 = y2[i].min(y2[idx]);

                let w = (xx2 - xx1 + 1.0).max(0.0);
                let h = (yy2 - yy1 + 1.0).max(0.0);
                let inter = w * h;

                let iou = inter / (areas[i] + areas[idx] - inter);

                // 只保留 IoU 小于阈值的框
                iou <= nms_thr
            });
        }

        keep
    }

    /// 5. multiclass_nms
    //multiclass_nms_class_agnostic
    pub fn multiclass_nms(
        &self,
        boxes: &Array2<f32>,  // [25200, 4] -> xyxy 格式
        scores: &Array2<f32>, // [25200, 80] -> 已经乘以 objectness 的得分
        nms_thr: f32,
        score_thr: f32,
    ) -> Vec<[f32; 6]> {
        let mut candidates = Vec::new();

        // 1. 筛选高分框 (单次遍历完成 Argmax 和 Threshold 过滤)
        for i in 0..scores.nrows() {
            let row = scores.row(i);

            // 找到当前行（即当前锚点）得分最高的类别
            let mut max_score = 0.0;
            let mut cls_id = 0;
            for (j, &s) in row.iter().enumerate() {
                if s > max_score {
                    max_score = s;
                    cls_id = j;
                }
            }

            // 仅保留超过阈值的候选框
            if max_score > score_thr {
                // 暂时存储索引和元数据，避免频繁创建大数组
                candidates.push((i, max_score, cls_id));
            }
        }

        if candidates.is_empty() {
            return vec![];
        }

        // 2. 准备 NMS 输入
        // 构造 NMS 需要的子集数组
        let mut b_subset = Array2::<f32>::zeros((candidates.len(), 4));
        let mut s_subset = Array1::<f32>::zeros(candidates.len());

        for (new_idx, &(orig_idx, score, _)) in candidates.iter().enumerate() {
            b_subset.row_mut(new_idx).assign(&boxes.row(orig_idx));
            s_subset[new_idx] = score;
        }

        // 3. 执行 NMS (返回保留下来的子集索引)
        let keep = self.nms(&b_subset, &s_subset, nms_thr);

        // 4. 组装最终结果 [x1, y1, x2, y2, score, class_id]
        keep.into_iter()
            .map(|k_idx| {
                let (orig_idx, score, cls_id) = candidates[k_idx];
                let b = boxes.row(orig_idx);
                [b[0], b[1], b[2], b[3], score, cls_id as f32]
            })
            .collect()
    }
    /// 6. get_bbox (完全解耦 OpenCV)
    pub fn get_bbox(&self, dynamic_img: &DynamicImage) -> Result<Vec<DetectionResult>> {
        // 使用 utils crate 解码
        // let dynamic_img = image::load_from_memory(image_bytes).context("Failed to decode utils")?;
        let (orig_w, orig_h) = dynamic_img.dimensions();

        let (input_tensor, ratio) = self.preproc(dynamic_img, (416, 416))?;

        // tract 推理
        // let outputs = self.session.session.run(tvec!(input_tensor.into()))?;
        let outputs = self.session.inference(input_tensor)?;
        // let output_array = outputs[0]
        let output_array = outputs
            .to_array_view::<f32>()?
            .to_owned()
            .into_dimensionality::<Ix3>()?;

        let predictions = self.demo_postprocess(output_array, (416, 416));
        let pred = predictions.slice(s![0, .., ..]);

        let boxes = pred.slice(s![.., 0..4]);
        let obj_conf = pred.slice(s![.., 4..5]);
        let cls_conf = pred.slice(s![.., 5..]);
        let obj_broadcast = obj_conf
            .broadcast(cls_conf.dim())
            .context("ndarray broadcasting failed for scores calculation")?;
        let scores = &obj_broadcast * &cls_conf;
        // let scores = &pred.slice(s![.., 4..5]) * &pred.slice(s![.., 5..]);

        let mut boxes_xyxy = Array2::<f32>::zeros(boxes.raw_dim());
        for i in 0..boxes.nrows() {
            boxes_xyxy[[i, 0]] = (boxes[[i, 0]] - boxes[[i, 2]] / 2.0) / ratio;
            boxes_xyxy[[i, 1]] = (boxes[[i, 1]] - boxes[[i, 3]] / 2.0) / ratio;
            boxes_xyxy[[i, 2]] = (boxes[[i, 0]] + boxes[[i, 2]] / 2.0) / ratio;
            boxes_xyxy[[i, 3]] = (boxes[[i, 1]] + boxes[[i, 3]] / 2.0) / ratio;
        }

        let detections = self.multiclass_nms(&boxes_xyxy, &scores, 0.45, 0.1);
        let final_results = detections
            .into_iter()
            .map(|d| DetectionResult {
                x1: (d[0] as i32).max(0).min(orig_w as i32),
                y1: (d[1] as i32).max(0).min(orig_h as i32),
                x2: (d[2] as i32).max(0).min(orig_w as i32),
                y2: (d[3] as i32).max(0).min(orig_h as i32),
                score: d[4],
                class_id: d[5] as u32,
            })
            .collect();
        Ok(final_results)
    }
}