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

namespace CognitiveToolkit
{
    public class SingleModelCharacterClassifier : ICharacterClassifier
    {
        private readonly DeviceDescriptor _device;
        private readonly Function _model;
        private readonly Variable _inputVariable;

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

        public CharacterClassificationResult Classify(float[] graytone28X28Data)
        {
            return ClassifyEx(graytone28X28Data).First();
        }

        public IEnumerable<CharacterClassificationResult> ClassifyEx(float[] graytone28X28Data)
        {
            var inputs = CreateInput(graytone28X28Data);
            var output = _model.Outputs.First(op => op.Shape.TotalSize == 29);
            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 CharacterClassificationResult
                {
                    Character = IntOutputToChar(i + 1),
                    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 char IntOutputToChar(int value)
        {
            if (value == 0)
            {
                return '?';
            }

            if (value < 27)
            {
                return (char)(value + 64);
            }

            if (value == 27)
            {
                return 'Æ';
            }

            if (value == 28)
            {
                return 'Ø';
            }

            return 'Å';
        }

    }
}

