using Hawkeye.VisionBuilder.Workflow.Datatypes;
using Hawkeye.VisionBuilder.Workflow.Operations.Attributes;
using OpenCvSharp;
using System.Runtime.InteropServices;
using System.Threading.Tasks;
using Hawkeye.VisionBuilder.Workflow.DataTransfer;

namespace Hawkeye.VisionBuilder.Workflow.Operations.AI;

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

    private readonly string _modelName;

    public MultichannelAI()
    {

        _cSharpDataTransferMQRPC = PythonModelProxyRPC.GetInterface();
        CanHaveProcessingError = false;
    }

    private readonly CSharpDataTransferMQRPC _cSharpDataTransferMQRPC;
    private int _imageWidth;
    private int _imageHeight;



    

    private bool _initialized = false;
    private FilePath _modelFilePath = new FilePath();


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

    

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

        if (!_initialized)
        {
            if (!File.Exists(ModelFilePath?.Path))
            {
                this.SetError("Model file not found");
                return;
            }
            _cSharpDataTransferMQRPC.TransferData<object>(new MethodCall("load_model", ModelFilePath.Path, _modelName));
            _initialized = true;
        }



        _cSharpDataTransferMQRPC.TransferData<object>(new MethodCall("activate_model", _modelName));

        var sizeEncoded = _cSharpDataTransferMQRPC.TransferData<object[]>(new MethodCall("get_accept_size"));
        var size = sizeEncoded.Select(Convert.ToInt32).ToArray();
        _imageWidth = size[1];
        _imageHeight = size[2];

        var outputSize = _cSharpDataTransferMQRPC.TransferData<object[]>(new MethodCall("get_output_size"));
        var outputSizeArray = outputSize.Select(Convert.ToInt32).ToArray();

        int classes;
        if (outputSizeArray.Length < 3)
        {
            classes = 1;
        }
        else
        {
            classes = outputSizeArray[2];
        }



        var currentImage = context.ActiveImage;

        var rightColor = currentImage.ImageData.CvtColor(ColorConversionCodes.BGR2RGB);
        var resized = rightColor.Resize(new Size(_imageHeight, _imageWidth));
        int dataSize = resized.Rows * resized.Cols * resized.ElemSize();
        var byteArray = new byte[dataSize];
        Marshal.Copy(resized.Data, byteArray, 0, dataSize);

        Mat[] channels = new Mat[classes];
        for (int i = 0; i < classes; i++)
        {
            channels[i] = new Mat(_imageHeight, _imageWidth, MatType.CV_8UC1);
        }

        var resultEncoded = _cSharpDataTransferMQRPC.TransferData<object[]>(new MethodCall("predict_raw", byteArray));

        var result = resultEncoded.Select(Convert.ToByte).ToArray();

        var chunked = result.Chunk(classes).ToArray();

        int pixel=0;
        foreach (var chunk in chunked)
        {
            for (int i = 0; i < classes; i++)
            {
                var data = (byte*)channels[i].Data.ToPointer();
                var coordX = (int)(pixel / _imageWidth);
                var coordY = pixel % _imageWidth;
                data[coordX * _imageWidth + coordY] = chunk[i];
            }
            pixel++;
        }

        //Cv2.ImShow("0", channels[0]);
        //Cv2.ImShow("1", channels[1]);

        using var res = new Mat();
        Cv2.Merge(channels,res);

        var resResized = res.Resize(new Size(currentImage.ImageData.Width, currentImage.ImageData.Height));
        context.ActiveImage = new HawkeyeImage() {ImageData = resResized };
        Result = true;
    }

    public override void Save(BinaryWriter bw)
    {
        base.Save(bw);
        bw.Write(_modelFilePath.Path);
        bw.Write(_modelFilePath.Format);
        
        
    }

    public override void Load(BinaryReader br)
    {
        base.Load(br);
        _modelFilePath.Path = br.ReadString();
        _modelFilePath.Format = br.ReadString();
        
        
    }
}