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

namespace Hawkeye.VisionBuilder.Workflow.Operations.AI;

[IgnoreOperation]
public abstract class ColorAIOperation:BaseOperation
{
    private readonly string _modelName;

    public ColorAIOperation(string model_name)
    {
        _modelName = model_name;
        _cSharpDataTransferMQRPC = PythonModelProxyHTTP.GetInterface();
        CanHaveProcessingError = false;
    }

    private readonly CSharpDataTransferHTTP _cSharpDataTransferMQRPC;
    private int _imageWidth;
    private int _imageHeight;

    private bool _initialized = false;
    private FilePath _modelFilePath;

    
    public virtual string FilterClasses { get; set; }

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

    public bool IsRaw { get; set; }

    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 not found");
                return;
            }

            _cSharpDataTransferMQRPC.LoadModel(ModelFilePath.Path, _modelName);
            _initialized = true;
        }


        _cSharpDataTransferMQRPC.ActivateModel(_modelName);

        var sizeEncoded = _cSharpDataTransferMQRPC.GetAcceptSize();
        
        var size = sizeEncoded.Select(Convert.ToInt32).ToArray();
        _imageWidth = size[1];
        _imageHeight = size[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);

        var allowed = FilterClasses.Split(',', StringSplitOptions.RemoveEmptyEntries).Select(x => Convert.ToInt32(x)).ToArray();

        var allowedFlags = new byte[128];
        foreach (var i in allowed)
        {
            allowedFlags[i] = 1;
        }


        byte[] result;
        if (IsRaw)
        {
            result = _cSharpDataTransferMQRPC.PredictRaw(byteArray);
        }
        else
        {
            result = _cSharpDataTransferMQRPC.Predict(byteArray);
        }


        byte[] resultColor;

        if (IsRaw)
        {
            resultColor = result.Chunk(25).Select(x =>
            {
                return x.Select((x, i) => allowedFlags[i] == 1 ? (byte) (x) : (byte) 0).Max();
            }).ToArray();
        }
        else
        {
            resultColor = result.Select(x => allowedFlags[x]==1? (byte)(x * 255.0) : (byte)0).ToArray();
        }
        

        // black and white mask
        var mask = new Mat(_imageHeight, _imageWidth, MatType.CV_8UC1);
        mask.SetArray(resultColor);

        var maskResized = mask.Resize(new Size(currentImage.ImageData.Width, currentImage.ImageData.Height));

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


        Result = true;
    }
}