using System.Reflection;
using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;
using OpenCvSharp;

namespace CandyboxPlugin.Detectors
{
    public class CandyDetector
    {
        
        private const string MODEL_NAME = "model_transparent_candy_resnet.onnx";
        private readonly InferenceSession _AIsession;
        private readonly string _inputName;
        private readonly int[] _dimensions;
        public CandyDetector()
        {
            var opts = new SessionOptions();
            // get this dll path
            var path = Assembly.GetCallingAssembly().Location;
            var dir = System.IO.Path.GetDirectoryName(path);
            var modelPath = Path.Combine(dir, "models\\" + MODEL_NAME);
            _AIsession = new InferenceSession(modelPath, opts);
            _inputName = _AIsession.InputMetadata.First().Key;
            _dimensions = _AIsession.InputMetadata.First().Value.Dimensions;
            foreach (var input in _AIsession.InputMetadata)
            {
                Console.WriteLine($"Name: {input.Key}");
                Console.WriteLine($"  Type: {input.Value.ElementType}");
                Console.WriteLine($"  Shape: [{string.Join(", ", input.Value.Dimensions)}]");
            }
        }


        public Mat Eval(Mat bmpTest)
        {
            var currentImage = bmpTest;

            var rightColor = currentImage.CvtColor(ColorConversionCodes.BGR2RGB);

            // convert to float32
            var floatImage = new Mat();
            rightColor.ConvertTo(floatImage, MatType.CV_32FC3, 1.0 / 255);
            // to array

            var width = _dimensions[1];
            var height = _dimensions[2];

            // check if need resize
            if (floatImage.Width != width || floatImage.Height != height)
                Cv2.Resize(floatImage, floatImage, new OpenCvSharp.Size(width, height));


            float[] data = new float[3 * width * height];
            int idx = 0;
            var rows = floatImage.Rows;
            var cols = floatImage.Cols;
            for (int y = 0; y < rows; y++)
            {
                for (int x = 0; x < cols; x++)
                {
                    Vec3f pixel = floatImage.At<Vec3f>(y, x);
                    for (int c = 0; c < 3; c++) // channels
                    {
                        data[idx++] = pixel[c]; // width*height*channel
                    }
                }
            }

            var inputTensor = new DenseTensor<float>(
                data,
                new int[] { 1, floatImage.Width, floatImage.Height, 3 }
            );

            var inputs = new List<NamedOnnxValue> {
                NamedOnnxValue.CreateFromTensor("input", inputTensor)
            };
            var sw = System.Diagnostics.Stopwatch.StartNew();
            using var results = _AIsession.Run(inputs);
            sw.Stop();
            Console.WriteLine($"Inference time: {sw.ElapsedMilliseconds} ms");
            var output = results.First().AsEnumerable<float>().ToArray();

            // convert output to Mat and normalize to 0-255
            var outputMat = new Mat(new OpenCvSharp.Size(width, height), MatType.CV_32FC1);
            outputMat.SetArray<float>(output);

            outputMat.ConvertTo(outputMat, MatType.CV_8UC1, 255);

            // resize back to original size
            Cv2.Resize(outputMat, outputMat, new OpenCvSharp.Size(currentImage.Width, currentImage.Height));
            return outputMat;
        }
    }
}