using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using CNTK;
using CrossCutting.Generic;
using NistUtilities.OpenCvSharp;
using OpenCvSharp;

namespace CognitiveToolkit.OpenCvSharp
{
    public class NumberInterpreter : INumberInterpreter<Mat, Rect>
    {
        private readonly DeviceDescriptor _device;
        private readonly Function _model;

        public NumberInterpreter(string modelFilename)
        {
            _device = DeviceDescriptor.CPUDevice; 
            _model = Function.Load(modelFilename, _device);
        }

        public (int Value, float Confidence) Interpret(Mat image)
        {
            var imageData = new MnistAdapter(image).To28X28Floats();
            var inputs = CreateInput(imageData);
            var outputs = new Dictionary<Variable, Value> { { _model.Output, null } };
            _model.Evaluate(inputs, outputs, _device);
            var value = outputs.FirstOrDefault().Value;
            var dense = value.GetDenseData<float>(_model.Output);
            var first = dense.First();
            if (first == null)
            {
                return (-1, 0);
            }
            var softmax = HelperFunctions.SoftMax(first.ToArray());
            return (Array.IndexOf(softmax, softmax.Max()), softmax.Max());
        }

        public (int Value, float Confidence)[] Interpret(Mat image, Rect area)
        {
            var imageData = new MnistAdapter(image.SubMat(area.Top, area.Bottom, area.Left, area.Right)).To28X28Floats();
            var inputs = CreateInput2(imageData);
            var outputs = new Dictionary<Variable, Value> { { _model.Output, null } };
            _model.Evaluate(inputs, outputs, _device);
            var value = outputs.FirstOrDefault().Value;
            var dense = value.GetDenseData<float>(_model.Output);
            var first = dense.FirstOrDefault();
            if (first == null)
            {
                return new(int Value, float Confidence)[] { };
            }
            var softmax = HelperFunctions.SoftMax(first.ToArray());
            return softmax.Select((qaz, wsx) => (wsx, qaz)).ToArray();
        }

        public IEnumerable<(int Value, float Confidence)> InterpretAll(Mat image, Rect area)
        {
            var imageData = new MnistAdapter(image.SubMat(area.Top, area.Bottom, area.Left, area.Right)).To28X28Floats();
            var inputs = CreateInput2(imageData);
            var outputs = new Dictionary<Variable, Value> { { _model.Output, null } };
            _model.Evaluate(inputs, outputs, _device);
            var value = outputs.FirstOrDefault().Value;
            var dense = value.GetDenseData<float>(_model.Output);
            return dense.First()?.Select((qaz, wsx) => (wsx, qaz)).ToArray();
        }

        private IDictionary<Variable, Value> CreateInput2(IEnumerable<float> data)
        {
            var ret = new Dictionary<Variable, Value>();
            var variable = _model.Inputs.First(i => i.IsInput); 
            var value = Value.CreateBatch(variable.Shape, data, _device, true);
            ret[variable] = value;
            return ret;
        }

        private IDictionary<Variable, Value> CreateInput(IEnumerable<float> data)
        {
            var ret = new Dictionary<Variable, Value>();
            var variable = _model.Inputs.First(i => i.IsInput); 
            var value = Value.CreateBatch(variable.Shape, data, _device, true);
            ret[variable] = value;
            return ret;
        }
    }
}
