using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Drawing;
using System.Drawing.Imaging;
using System.Linq;

namespace YoloTest
{
    internal class YoloV6Detector : IDisposable
    {
        private InferenceSession _session;
        public string[] _classNames;  // 添加此字段

        public YoloV6Detector(string modelPath, string[] classNames = null, bool useGpu = true)
        {
            var options = new SessionOptions
            {
                InterOpNumThreads = 1,  // 减少线程切换开销
                IntraOpNumThreads = Environment.ProcessorCount,
                GraphOptimizationLevel = GraphOptimizationLevel.ORT_ENABLE_ALL
            };
            options.AddSessionConfigEntry("session.intra_op.allow_spinning", "1");
            options.AddSessionConfigEntry("session.inter_op.allow_spinning", "1");

            if (useGpu)
            {
                try
                {
                    //options.AppendExecutionProvider("OpenVINO");
                    //options.AppendExecutionProvider_DML(0);
                    options.AppendExecutionProvider_CUDA(0);
                    Console.WriteLine("GPU推理运行");
                }
                catch(Exception e)
                {
                    Console.WriteLine($"{e.Message}");
                    options.AppendExecutionProvider_CPU(0);
                    Console.WriteLine("CPU推理运行");
                }
            }



            _session = new InferenceSession(modelPath, options);

            _classNames = classNames ?? new string[]
            {
                "OK", "NG"
            };
        }

        public YoloV6Detector(string modelPath, string[] classNames = null)
        {
            var options = new SessionOptions();
            _session = new InferenceSession(modelPath, options);

            // 初始化类别名称
            _classNames = classNames ?? new string[]
            {
                "OK", "NG"
                // 根据你的模型实际类别修改
            };
        }

        Stopwatch OnceRunTime = new Stopwatch();
        public List<Detection> Detect(Image inputImage)
        {
            var inputTensor = PreprocessImage(inputImage);
            var inputs = new List<NamedOnnxValue>
            {
                NamedOnnxValue.CreateFromTensor("images", inputTensor)
            };
            OnceRunTime.Restart();
            using (var results = _session.Run(inputs))
            {
                OnceRunTime.Stop();
                Console.WriteLine("单次运行耗时{0}",OnceRunTime.ElapsedMilliseconds);
                return Postprocess(results, inputImage.Width, inputImage.Height);
            }
        }

        private Tensor<float> PreprocessImage(Image image)
        {
            const int inputWidth = 640;
            const int inputHeight = 640;
            const int channels = 3;

            using (var resizedImage = new Bitmap(image, inputWidth, inputHeight))
            {
                var tensorData = new float[1 * channels * inputHeight * inputWidth];

                // 使用 LockBits 直接访问内存，比 GetPixel 快 10-50 倍
                var bitmapData = resizedImage.LockBits(
                    new Rectangle(0, 0, inputWidth, inputHeight),
                    ImageLockMode.ReadOnly,
                    PixelFormat.Format24bppRgb);

                try
                {
                    unsafe
                    {
                        byte* ptr = (byte*)bitmapData.Scan0;
                        int stride = bitmapData.Stride;

                        for (int y = 0; y < inputHeight; y++)
                        {
                            for (int x = 0; x < inputWidth; x++)
                            {
                                byte* pixel = ptr + y * stride + x * 3;
                                int idx = y * inputWidth + x;
                                tensorData[idx] = pixel[2] / 255.0f;  // R
                                tensorData[inputHeight * inputWidth + idx] = pixel[1] / 255.0f;  // G
                                tensorData[2 * inputHeight * inputWidth + idx] = pixel[0] / 255.0f;  // B
                            }
                        }
                    }
                }
                finally
                {
                    resizedImage.UnlockBits(bitmapData);
                }

                return new DenseTensor<float>(tensorData, new[] { 1, channels, inputHeight, inputWidth });
            }
        }

        //private Tensor<float> PreprocessImage(Image image)
        //{
        //    const int inputWidth = 640;
        //    const int inputHeight = 640;
        //    const int channels = 3;

        //    using (var resizedImage = new Bitmap(image, inputWidth, inputHeight))
        //    {
        //        var tensorData = new float[1 * channels * inputHeight * inputWidth];

        //        for (int y = 0; y < inputHeight; y++)
        //        {
        //            for (int x = 0; x < inputWidth; x++)
        //            {
        //                var pixel = resizedImage.GetPixel(x, y);
        //                tensorData[y * inputWidth + x] = pixel.R / 255.0f;
        //                tensorData[inputHeight * inputWidth + y * inputWidth + x] = pixel.G / 255.0f;
        //                tensorData[2 * inputHeight * inputWidth + y * inputWidth + x] = pixel.B / 255.0f;
        //            }
        //        }

        //        return new DenseTensor<float>(tensorData, new[] { 1, channels, inputHeight, inputWidth });
        //    }
        //}

        private List<Detection> Postprocess(IEnumerable<NamedOnnxValue> outputs, int imgWidth, int imgHeight)
        {
            var detections = new List<Detection>();

            // 1. 获取模型输出
            var output = outputs.FirstOrDefault();
            if (output == null) return detections;

            var outputTensor = output.AsTensor<float>();
            var outputShape = outputTensor.Dimensions;

            // YOLOv6 输出格式：[batch, num_anchors, 85] 或 [batch, num_anchors, 4 + 1 + num_classes]
            int numClasses = _classNames.Length;
            int numAnchors = outputShape[1];
            int numValuesPerAnchor = outputShape[2]; // 通常为 4 + 1 + numClasses

            // 2. 解析输出张量
            var rawDetections = new List<RawDetection>();
            const float confidenceThreshold = 0.4f;
            const float nmsThreshold = 0.2f;

            for (int i = 0; i < numAnchors; i++)
            {
                // 获取目标置信度 (第 5 个值，索引从 4 开始)
                float objectness = outputTensor[0, i, 4];

                if (objectness < confidenceThreshold) continue;

                // 获取各类别概率并计算最大置信度
                float maxClassProb = 0;
                int maxClassId = 0;

                for (int c = 0; c < numClasses; c++)
                {
                    float classProb = outputTensor[0, i, 5 + c];
                    if (classProb > maxClassProb)
                    {
                        maxClassProb = classProb;
                        maxClassId = c;
                    }
                }

                float confidence = objectness * maxClassProb;
                if (confidence < confidenceThreshold) continue;

                // 获取边界框坐标 (cx, cy, w, h)
                float cx = outputTensor[0, i, 0];
                float cy = outputTensor[0, i, 1];
                float w = outputTensor[0, i, 2];
                float h = outputTensor[0, i, 3];

                // 转换为左上角坐标
                float x = cx - w / 2;
                float y = cy - h / 2;

                rawDetections.Add(new RawDetection
                {
                    X = x,
                    Y = y,
                    Width = w,
                    Height = h,
                    ClassId = maxClassId,
                    Confidence = confidence
                });
            }

            // 3. 执行非极大值抑制 (NMS)
            var nmsDetections = ApplyNMS(rawDetections, nmsThreshold);

            // 4. 将坐标从模型尺寸转换回原始图像尺寸
            float scaleX = (float)imgWidth / 640;
            float scaleY = (float)imgHeight / 640;

            foreach (var det in nmsDetections)
            {
                detections.Add(new Detection(
                    det.X * scaleX,
                    det.Y * scaleY,
                    det.Width * scaleX,
                    det.Height * scaleY,
                    det.ClassId,
                    det.Confidence
                ));
            }

            return detections;
        }

        // 辅助类：原始检测结果
        private class RawDetection
        {
            public float X { get; set; }
            public float Y { get; set; }
            public float Width { get; set; }
            public float Height { get; set; }
            public int ClassId { get; set; }
            public float Confidence { get; set; }
        }

        // 非极大值抑制 (NMS)
        private List<RawDetection> ApplyNMS(List<RawDetection> detections, float iouThreshold)
        {
            if (detections.Count == 0) return detections;

            // 按置信度降序排序
            var sorted = detections.OrderByDescending(d => d.Confidence).ToList();
            var results = new List<RawDetection>();

            while (sorted.Count > 0)
            {
                var best = sorted[0];
                results.Add(best);
                sorted.RemoveAt(0);

                // 移除与当前最佳检测框 IoU 过高的框
                sorted = sorted.Where(d =>
                {
                    if (d.ClassId != best.ClassId) return true;
                    return CalculateIoU(best, d) < iouThreshold;
                }).ToList();
            }

            return results;
        }

        // 计算 IoU (交并比)
        private float CalculateIoU(RawDetection a, RawDetection b)
        {
            // 计算交集
            float x1 = Math.Max(a.X, b.X);
            float y1 = Math.Max(a.Y, b.Y);
            float x2 = Math.Min(a.X + a.Width, b.X + b.Width);
            float y2 = Math.Min(a.Y + a.Height, b.Y + b.Height);

            float intersection = Math.Max(0, x2 - x1) * Math.Max(0, y2 - y1);

            // 计算并集
            float areaA = a.Width * a.Height;
            float areaB = b.Width * b.Height;
            float union = areaA + areaB - intersection;

            return union > 0 ? intersection / union : 0;
        }

        public void Dispose()
        {
            _session?.Dispose();
        }
    }

    // Detection 类定义
    public class Detection
    {
        public float X { get; set; }
        public float Y { get; set; }
        public float Width { get; set; }
        public float Height { get; set; }
        public int ClassId { get; set; }
        public float Confidence { get; set; }

        public Detection(float x, float y, float width, float height, int classId, float confidence)
        {
            X = x;
            Y = y;
            Width = width;
            Height = height;
            ClassId = classId;
            Confidence = confidence;
        }
    }
}