// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.
 
using Microsoft.ML.TorchSharp.Utils;
using TorchSharp;
using TorchSharp.Modules;
 
namespace Microsoft.ML.TorchSharp.NasBert.Models
{
    internal sealed class PredictionHead : BaseHead, torch.nn.IModule<torch.Tensor, torch.Tensor>
    {
        [System.Diagnostics.CodeAnalysis.SuppressMessage("Naming", "MSML_PrivateFieldName:Private field name not in: _camelCase format", Justification = "Has to match TorchSharp model.")]
        private readonly Sequential Classifier;
        private bool _disposedValue;
 
        public PredictionHead(int inputDim, int numClasses, double dropoutRate)
            : base(nameof(PredictionHead))
        {
            var dropoutLayer = torch.nn.Dropout(dropoutRate);
            var dense = torch.nn.Linear(inputDim, numClasses);
 
            ModelUtils.InitXavierUniform(dense.weight);
            ModelUtils.InitZeros(dense.bias);
 
            Classifier = torch.nn.Sequential(
                ("dropout1", dropoutLayer),
                ("dense", dense)
            );
 
            RegisterComponents();
        }
 
        [System.Diagnostics.CodeAnalysis.SuppressMessage("Naming", "MSML_GeneralName:This name should be PascalCased", Justification = "Need to match TorchSharp")]
        public torch.Tensor call(torch.Tensor features)
        {
            // TODO: try whitening-like techniques
            // take <s> token (equiv. to [CLS])
            using var x = features[torch.TensorIndex.Colon, torch.TensorIndex.Single(0), torch.TensorIndex.Colon];
            return Classifier.forward(x);
        }
 
        protected override void Dispose(bool disposing)
        {
            if (!_disposedValue)
            {
                if (disposing)
                {
                    Classifier.Dispose();
                    _disposedValue = true;
                }
            }
 
            base.Dispose(disposing);
        }
    }
}