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

namespace CognitiveToolkit
{
    public class CheckboxClassifier : ICheckboxInterpreter<Mat>
    {
        private readonly DeviceDescriptor _device;
        private readonly Function _model;
        private readonly Variable _inputVariable;

        public CheckboxClassifier(string modelFilename)
        {
            _device = DeviceDescriptor.CPUDevice; 
            _model = Function.Load(modelFilename, _device);
            _inputVariable = _model.Inputs.First(i => i.IsInput);
        }

        private IEnumerable<ClassificationResult<CheckboxState>> ClassifyEx(float[] graytone28X28Data)
        {
            var inputs = CreateInput(graytone28X28Data);
            var output = _model.Outputs.First(op => op.Shape.TotalSize == 3);
            var outputs = new Dictionary<Variable, Value> { { output, null } };
            _model.Evaluate(inputs, outputs, _device);
            var value = outputs.FirstOrDefault().Value;
            var dense = value.GetDenseData<float>(output);
            var first = dense.First();
            var softmax = HelperFunctions.SoftMax(first.ToArray());
            return softmax.Select((f, i) =>
                new ClassificationResult<CheckboxState>
                {
                    Result = IntOutputToCheckboxState(i),
                    Confidence = f
                })
                .OrderByDescending(entry => entry.Confidence);
        }

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

        private static CheckboxState IntOutputToCheckboxState(int value)
        {
            switch (value)
            {
                case 0:
                    return CheckboxState.NotChecked;
                case 1:
                    return CheckboxState.Checked;
                default:
                    return CheckboxState.Filled;
            }

        }

        public ClassificationResult<CheckboxState> Interpret(Mat image)
        {
            return InterpretAll(image).FirstOrDefault();
        }

        public IEnumerable<ClassificationResult<CheckboxState>> InterpretAll(Mat image)
        {
            var adapter = new EmnistAdapter(image);
            return ClassifyEx(adapter.To28X28Floats());
        }
    }
}

