// 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.Extensions.ObjectPool;
using Microsoft.ML;
 
namespace Microsoft.Extensions.ML
{
    /// <summary>
    /// <see cref="PooledObjectPolicy{T}"/> for <see cref="PredictionEngine{TData, TPrediction}"/>
    /// which is responsible for creating pooled objects, and when to return objects to the pool.
    /// </summary>
    internal class PredictionEnginePoolPolicy<TData, TPrediction>
        : PooledObjectPolicy<PredictionEngine<TData, TPrediction>>
        where TData : class
        where TPrediction : class, new()
    {
        private readonly MLContext _mlContext;
        private readonly ITransformer _model;
 
        /// <summary>
        /// Initializes a new instance of <see cref="PredictionEnginePoolPolicy{TData, TPrediction}"/>.
        /// </summary>
        /// <param name="mlContext">
        /// <see cref="MLContext"/> used to load the model.
        /// </param>
        /// <param name="model">The transformer to use for prediction.</param>
        public PredictionEnginePoolPolicy(MLContext mlContext, ITransformer model)
        {
            _mlContext = mlContext;
            _model = model;
        }
 
        /// <inheritdoc />
        public override PredictionEngine<TData, TPrediction> Create() =>
            _mlContext.Model.CreatePredictionEngine<TData, TPrediction>(
                _model,
                new PredictionEngineOptions { OwnsTransformer = false });
 
        /// <inheritdoc />
        public override bool Return(PredictionEngine<TData, TPrediction> obj) => true;
    }
}