using Hawkeye.VisionBuilder.Workflow.Datatypes;
using Hawkeye.VisionBuilder.Workflow.Operations.Attributes;
using ILGPU.Runtime.Cuda;
using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;
using Microsoft.Scripting.Runtime;
using OpenCvSharp;
namespace Hawkeye.VisionBuilder.Workflow.Operations.AI;

[Category("AI")]
public class OnnxModelOperation:BaseOperation
{

    private FilePath _modelFilePath = new FilePath(){Format = "ONNX file(*.onnx)|*.onnx" };

    
    public FilePath ModelFilePath
    {
        get => _modelFilePath;
        set
        {
            _modelFilePath = value;
            _initialized = false;
        }
    }

    private bool _initialized = false;
    private InferenceSession _AIsession;
    private string _inputName;
    private int[] _dimensions;

    protected override void InterpretInternal(Context context)
    {
        if (!CheckImageExists(context)) return;
        if (!CheckColorful(context)) return;

        if (!_initialized)
        {
            if (!File.Exists(ModelFilePath?.Path))
            {
                this.SetError($"Model file \"{ModelFilePath?.Path}\" not found");
                return;
            }

            var opts = new SessionOptions();
            _AIsession = new InferenceSession(ModelFilePath?.Path, 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)}]");
            }

            _initialized = true;
        }

        var currentImage = context.ActiveImage;

        var rightColor = currentImage.ImageData.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.ImageData.Width, currentImage.ImageData.Height));


        context.ActiveImage = new HawkeyeImage()
        {
            ImageData = outputMat
        };

    }

    public override void Save(Dictionary<string, object> dict)
    {
        base.Save(dict);
        dict[nameof(ModelFilePath)] = ModelFilePath?.Path;
    }

    public override void Load(Dictionary<string, object> dict)
    {
        base.Load(dict);
        if (dict.ContainsKey(nameof(ModelFilePath)))
            ModelFilePath.Path = dict[nameof(ModelFilePath)].ToString();
    }
}