using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using CNTK;
using OpenCvSharp;
using OpenCvSharp.Extensions;

namespace OCR
{
    internal class ConvolutionNeuralNet
    {
        private static readonly List<string> Decoder = new List<string> { "0123456789", "abcdefghijklmnopqrstuvwxyz", "ABCDEFGHIJKLMNOPQRSTUVWXYZ" };


        /// <summary>
        /// Trains convolution neural network for given number of classes and saves the model
        /// </summary>
        /// <param name="numClasses">number of classes for classification</param>
        /// <param name="classifierName">name of the classifier, possible values: nunbers, upper, lower</param>
        /// <param name="validate">should the existing model be validated</param>
        /// <param name="forceRetrain">should the existing model be retrained</param>
        public void TrainAndSave(int numClasses, string classifierName, bool validate = true, bool forceRetrain = false)
        {
            var inputDim = NDShape.CreateNDShape(new[] { 32, 32, 3});
            var modelFile = classifierName + ".model";

            Console.WriteLine($"========== Model: {classifierName.ToUpper()} ==========");

            // If a model already exists and not set to force retrain, validate the model and return.
            if (File.Exists(modelFile) && !forceRetrain)
            {
                Console.WriteLine($"Model file {modelFile} exists.");
                if (validate)
                {
                    var minibatchSourceExistModel = CreateMinibatchSource($"./test/{classifierName}.txt", numClasses);
                    Function model = Function.Load(modelFile, DeviceDescriptor.GPUDevice(0));
                    Console.WriteLine("========== Validation ==========");
                    var errorOnValidation = ValidateModelWithMinibatchSource(minibatchSourceExistModel, model, classifierName);
                    Console.WriteLine($" Validation error: {errorOnValidation}");
                    Console.WriteLine("");
                } else {
                    Console.WriteLine("Model validation skipped.");
                }
                return;
            }

            //create input
            var input = CNTKLib.InputVariable(inputDim, DataType.Float, "images");
            var scaledInput = CNTKLib.ElementTimes(Constant.Scalar(1/255f, DeviceDescriptor.GPUDevice(0)), input);

            //create network
            Function neuralNetwork = CreateConvolutionalNeuralNetwork(scaledInput, numClasses, classifierName);

            var labels = CNTKLib.InputVariable(new[] { numClasses }, DataType.Float, "labels");
            var lossFunction = CNTKLib.CrossEntropyWithSoftmax(neuralNetwork.Output, labels, "lossFunction");
            var evaluationFunction = CNTKLib.ClassificationError(neuralNetwork.Output, labels, "classificationError");

            //learning rate and training parameters
            var learningRateSchedule = new TrainingParameterScheduleDouble(0.01);
            var momentumSchedule = new TrainingParameterScheduleDouble(0.9);
            IList<Learner> parameterLearners = new List<Learner> {Learner.MomentumSGDLearner(neuralNetwork.Parameters(), learningRateSchedule, momentumSchedule, true)};

            //create trainer
            Trainer trainer = Trainer.CreateTrainer(neuralNetwork, lossFunction, evaluationFunction, parameterLearners);
            // create minibatch source for training, currently chars74k dataset (filtered) + real data 
            MinibatchSource trainingMinibatch = CreateMinibatchSource($"./training/{classifierName}.txt", numClasses);
            // create minibatch source for validation, chars segmented from real data
            MinibatchSource validatingMinibatch = CreateMinibatchSource($"./test/{classifierName}.txt", numClasses);

            var imageStream = trainingMinibatch.StreamInfo("images");
            var labelStream = trainingMinibatch.StreamInfo("labels");

            int infoFrequency = 100;
            int minibatchIdx = 0;
            uint minibatchSize = 100;
            int consecZeroEvaluation = 0;
            Console.WriteLine("========== Training started ==========");
            while (true)
            {
                // get minibatch
                var nextMinibatch = trainingMinibatch.GetNextMinibatch(minibatchSize, DeviceDescriptor.GPUDevice(0));
                var arguments = new Dictionary<Variable, MinibatchData> {
                    { input, nextMinibatch[imageStream] },
                    { labels, nextMinibatch[labelStream] }
                };
                // train
                trainer.TrainMinibatch(arguments, DeviceDescriptor.GPUDevice(0));
                trainer.SummarizeTrainingProgress();
                float trainLossValue = (float)trainer.PreviousMinibatchLossAverage();
                float evaluationValue = (float)trainer.PreviousMinibatchEvaluationAverage();
                minibatchIdx++;

                if (minibatchIdx % infoFrequency == 0 && trainer.PreviousMinibatchSampleCount() > 0)
                    Console.WriteLine(
                        $"Minibatch: {minibatchIdx} CrossEntropyLoss = {trainLossValue}, EvaluationCriterion = {evaluationValue}");
                consecZeroEvaluation = evaluationValue == 0.0f ? consecZeroEvaluation + 1 : 0;

                //validate at the end of epoch
                if (nextMinibatch.Values.Any(a => a.sweepEnd)) {
                    float errorRate = ValidateModelWithMinibatchSource(validatingMinibatch, neuralNetwork, classifierName, -1, false);
                    Console.WriteLine($"Validation error rate: {errorRate:0.000}.");

                    // check training end condition
                    if (errorRate < 0.01f || minibatchIdx > 10000 || consecZeroEvaluation > 30) {
                        neuralNetwork.Save(modelFile);
                        Console.WriteLine("===== Response on training data =====");
                        var errorOnTraining = ValidateModelWithMinibatchSource(trainingMinibatch, neuralNetwork, classifierName);
                        Console.WriteLine("===== Response on validating data =====");
                        var errorOnValidation = ValidateModelWithMinibatchSource(validatingMinibatch, neuralNetwork, classifierName);
                        Console.WriteLine("========== Training finished ==========");
                        Console.WriteLine($"Training data error: {errorOnTraining}");
                        Console.WriteLine($"Validation data error: {errorOnValidation}");
                        Console.WriteLine("");
                        return;
                    }
                }
            }
        }


        /// <summary>
        /// Method for evaluation of single character using trained model of given type
        /// </summary>
        /// <param name="image">image to be evaluated, needs to be in shape 32*32</param>
        /// <param name="model">model for evaluation</param>
        /// <param name="modelType">type of used model, possible values: Numbers, Lower, Upper</param>
        /// <returns>classification result for given image</returns>
        public ClassificationResult EvaluateSingleCharacter(Mat image, Function model, ModelType modelType)
        {
            var testImage = image.Clone().ToBitmap();
            var features = testImage.ParallelExtractCHW();
            var modelOut = CNTKLib.Softmax(model, "softmax");
            var inputVar = modelOut.Arguments.Single();
            var labelOutput = modelOut.Outputs.Single(o => o.Name == "softmax");
            var inputDataMap = new Dictionary<Variable, Value> {
                { inputVar, Value.CreateBatch(inputVar.Shape, features, DeviceDescriptor.GPUDevice(0)) }
            };
            var outputDataMap = new Dictionary<Variable, Value> {
                { labelOutput, null }
            };
            modelOut.Evaluate(inputDataMap, outputDataMap, DeviceDescriptor.GPUDevice(0));
            // Get evaluate result as dense output
            var outputData = outputDataMap[labelOutput].GetDenseData<float>(labelOutput)[0];
            ClassificationResult result = new ClassificationResult {
                Label = outputData.IndexOf(outputData.Max()),
                Probability = outputData.Max(),
                DecodedLabel = Decoder[(int) modelType][outputData.IndexOf(outputData.Max())].ToString()
            };
            return result;
        }
        

        /// <summary>
        /// Creates convolution neural network model for given number of output classes
        /// </summary>
        /// <param name="input">input variable</param>
        /// <param name="outputDim">number of output classes</param>
        /// <param name="classifierName">name of the classifier, possible: numbers, lower, upper</param>
        /// <returns>function representing the convolution neural network model</returns>
        private Function CreateConvolutionalNeuralNetwork(Variable input, int outputDim, string classifierName)
        {
            //first layer - convolution, batch normalization, ReLU, pooling
            var kernel1 = new Parameter(NDShape.CreateNDShape(new[] {5, 5, 3, 6}), DataType.Float, CNTKLib.GlorotNormalInitializer());
            var convolutionLayer1 = CNTKLib.Convolution(kernel1, input, NDShape.CreateNDShape(new []{1,1,3}), BoolVector.Repeat(true, 1), BoolVector.Repeat(true, 1));
            var batchNormalLayer1 = CNTKLib.BatchNormalization(convolutionLayer1,
                new Parameter(new[] {NDShape.InferredDimension}, 1, DeviceDescriptor.GPUDevice(0), ""),
                new Parameter(new[] {NDShape.InferredDimension}, 0, DeviceDescriptor.GPUDevice(0), ""),
                new Constant(new[] {NDShape.InferredDimension}, 0.0f, DeviceDescriptor.GPUDevice(0)),
                new Constant(new[] { NDShape.InferredDimension }, 0.0f, DeviceDescriptor.GPUDevice(0)),
                Constant.Scalar(0.0f, DeviceDescriptor.GPUDevice(0)), true, 5000, 0.0, 1e-5,false);
            var reluLayer1 = CNTKLib.ReLU(batchNormalLayer1);
            var poolingLayer1 = CNTKLib.Pooling(reluLayer1, PoolingType.Max, NDShape.CreateNDShape(new[] { 2, 2}), NDShape.CreateNDShape(new[] { 2, 2 }));
            //second layer - convolution, batch normalization, ReLU, pooling
            var kernel2 = new Parameter(NDShape.CreateNDShape(new[] { 5, 5, 6, 16 }), DataType.Float, CNTKLib.GlorotNormalInitializer());
            var convolutionLayer2 = CNTKLib.Convolution(kernel2, poolingLayer1, NDShape.CreateNDShape(new[] { 1, 1, 6 }), BoolVector.Repeat(true, 1), BoolVector.Repeat(true, 1));
            var batchNormalLayer2 = CNTKLib.BatchNormalization(convolutionLayer2,
                new Parameter(new[] { NDShape.InferredDimension }, 1, DeviceDescriptor.GPUDevice(0), ""),
                new Parameter(new[] { NDShape.InferredDimension }, 0, DeviceDescriptor.GPUDevice(0), ""),
                new Constant(new[] { NDShape.InferredDimension }, 0.0f, DeviceDescriptor.GPUDevice(0)),
                new Constant(new[] { NDShape.InferredDimension }, 0.0f, DeviceDescriptor.GPUDevice(0)),
                Constant.Scalar(0.0f, DeviceDescriptor.GPUDevice(0)), true, 5000, 0.0, 1e-5, false);
            var reluLayer2 = CNTKLib.ReLU(batchNormalLayer2);
            var poolingLayer2 = CNTKLib.Pooling(reluLayer2, PoolingType.Max, NDShape.CreateNDShape(new[] { 2, 2 }), NDShape.CreateNDShape(new[] { 2, 2 }));
            //third layer - convolution, batch normalization, ReLU
            var kernel3 = new Parameter(NDShape.CreateNDShape(new[] { 5, 5, 16, 120 }), DataType.Float, CNTKLib.GlorotNormalInitializer());
            Function convolutionLayer3 = CNTKLib.Convolution(kernel3, poolingLayer2, NDShape.CreateNDShape(new[] { 1, 1, 16 }), BoolVector.Repeat(true,1), BoolVector.Repeat(false,1));
            var batchNormalLayer3 = CNTKLib.BatchNormalization(convolutionLayer3,
                new Parameter(new[] { NDShape.InferredDimension }, 1, DeviceDescriptor.GPUDevice(0), ""),
                new Parameter(new[] { NDShape.InferredDimension }, 0, DeviceDescriptor.GPUDevice(0), ""),
                new Constant(new[] { NDShape.InferredDimension }, 0.0f, DeviceDescriptor.GPUDevice(0)),
                new Constant(new[] { NDShape.InferredDimension }, 0.0f, DeviceDescriptor.GPUDevice(0)),
                Constant.Scalar(0.0f, DeviceDescriptor.GPUDevice(0)), true, 5000, 0.0, 1e-5, false);
            var reluLayer3 = CNTKLib.ReLU(batchNormalLayer3);
            //reshape for fully connected layer
            int dimensions = reluLayer3.Output.Shape.Dimensions.Aggregate((d1, d2) => d1 * d2);
            reluLayer3 = CNTKLib.Reshape(reluLayer3, new NDShape(1, dimensions));
            //fully connected layer
            var timesParam = new Parameter(NDShape.CreateNDShape(new[] {outputDim, reluLayer3.Output.Shape[0]}),
                DataType.Float, 
                CNTKLib.GlorotUniformInitializer(),
                DeviceDescriptor.GPUDevice(0),
                "timesParam");
            Function timesFuncion = CNTKLib.Times(timesParam, reluLayer3, "times"); // weights
            var plusParam = new Parameter(new[] { outputDim }, 0.0f, DeviceDescriptor.GPUDevice(0), "plusParam"); // biases
            return CNTKLib.Plus(plusParam, timesFuncion, classifierName);
        }


        /// <summary>
        /// Creates minibatch source of images from text file including direct paths and classes for each file
        /// </summary>
        /// <param name="filePath">path to text file</param>
        /// <param name="numClasses">number of classes in minibatch</param>
        /// <returns>Minibatch Source</returns>
        private MinibatchSource CreateMinibatchSource(string filePath, int numClasses)
        {
            List<CNTKDictionary> transforms = new List<CNTKDictionary> { CNTKLib.ReaderScale(32, 32, 3) };
            var cntkDictionary = CNTKLib.ImageDeserializer(filePath, "labels", (uint)numClasses, "images", transforms);
            MinibatchSourceConfig configuration = new MinibatchSourceConfig(new List<CNTKDictionary> {cntkDictionary});
            return CNTKLib.CreateCompositeMinibatchSource(configuration);
        }


        /// <summary>
        /// Validates the model with minibatch
        /// </summary>
        /// <param name="source">minibatch source for validation</param>
        /// <param name="model">model to be validated</param>
        /// <param name="outputName">name of the output layer, should be </param>
        /// <param name="maxCount">maximum count of samples used for validation, set to negative value for validation on all samples in source</param>
        /// <param name="notify">should be notified about the validation cycles</param>
        /// <returns>error rate during validation</returns>
        private float ValidateModelWithMinibatchSource(MinibatchSource source, Function model, string outputName, int maxCount = -1, bool notify = true)
        {
            var imageInput = model.Arguments[0];
            var labelOutput = model.Outputs.Single(o => o.Name == outputName);
            var imageStreamInfo = source.StreamInfo("images");
            var labelStreamInfo = source.StreamInfo("labels");
            int batchSize = 500;
            int miscountTotal = 0, totalCount = 0;
            while (true)
            {
                var minibatchData = source.GetNextMinibatch((uint) batchSize, DeviceDescriptor.GPUDevice(0));
                if (minibatchData == null || minibatchData.Count == 0)
                    break;
                totalCount += (int) minibatchData[labelStreamInfo].numberOfSamples;

                var labelData = minibatchData[labelStreamInfo].data.GetDenseData<float>(labelOutput);
                var expectedLabels = labelData.Select(l => l.IndexOf(l.Max())).ToList();
                var inputDataMap = new Dictionary<Variable, Value> {
                    { imageInput, minibatchData[imageStreamInfo].data }
                };
                var outputDataMap = new Dictionary<Variable, Value> {
                    { labelOutput, null }
                };
                model.Evaluate(inputDataMap, outputDataMap, DeviceDescriptor.GPUDevice(0));
                var outputData = outputDataMap[labelOutput].GetDenseData<float>(labelOutput);
                var actualLabels = outputData.Select(l => l.IndexOf(l.Max())).ToList();

                int misMatches = actualLabels.Zip(expectedLabels, (a, b) => a.Equals(b) ? 0 : 1).Sum();
                miscountTotal += misMatches;
                if (notify)
                {
                    Console.WriteLine($"Validating Model: Total Samples = {totalCount}, MisclassifyCount = {miscountTotal}");
                }

                bool condition = maxCount < 0 ? minibatchData.Values.Any(a => a.sweepEnd) : totalCount > maxCount;

                if (condition)
                    break;
            }
            float errorRate = 1f * miscountTotal / totalCount;
            return errorRate;
        }
    }
}