namespace Microsoft.ML.AutoML.CodeGen
{
    public enum EstimatorType
    {
        LightGbmBinary,
        LightGbmMulti,
        LightGbmRegression,
        FastForestBinary,
        FastForestOva,
        FastForestRegression,
        FastTreeBinary,
        FastTreeOva,
        FastTreeRegression,
        FastTreeTweedieRegression,
        LbfgsLogisticRegressionBinary,
        LbfgsLogisticRegressionOva,
        LbfgsMaximumEntropyMulti,
        LbfgsPoissonRegressionRegression,
        SdcaLogisticRegressionBinary,
        SdcaLogisticRegressionOva,
        SdcaMaximumEntropyMulti,
        SdcaRegression,
        MatrixFactorization,
        ImageClassificationMulti,
        TextClassificationMulti,
        SentenceSimilarityRegression,
        ObjectDetectionMulti,
        QuestionAnsweringMulti,
        NamedEntityRecognitionMulti,
        ForecastBySsa,
        Concatenate,
        Naive,
        OneHotEncoding,
        OneHotHashEncoding,
        LoadRawImageBytes,
        MapKeyToValue,
        ReplaceMissingValues,
        NormalizeMinMax,
        FeaturizeText,
        NormalizeText,
        ConvertType,
        MapValueToKey,
        ApplyOnnxModel,
        ResizeImages,
        ExtractPixels,
        LoadImages,
        DnnFeaturizerImage,
        Unknown,
    }
 
    public static class EstimatorTypeExtension
    {
        public static bool IsTrainer(this EstimatorType estimatorType)
        {
            switch(estimatorType)
            {
                case EstimatorType.LightGbmBinary:
                case EstimatorType.LightGbmMulti:
                case EstimatorType.LightGbmRegression:
                case EstimatorType.FastForestBinary:
                case EstimatorType.FastForestOva:
                case EstimatorType.FastForestRegression:
                case EstimatorType.FastTreeBinary:
                case EstimatorType.FastTreeOva:
                case EstimatorType.FastTreeRegression:
                case EstimatorType.FastTreeTweedieRegression:
                case EstimatorType.LbfgsLogisticRegressionBinary:
                case EstimatorType.LbfgsLogisticRegressionOva:
                case EstimatorType.LbfgsMaximumEntropyMulti:
                case EstimatorType.LbfgsPoissonRegressionRegression:
                case EstimatorType.SdcaLogisticRegressionBinary:
                case EstimatorType.SdcaLogisticRegressionOva:
                case EstimatorType.SdcaMaximumEntropyMulti:
                case EstimatorType.SdcaRegression:
                case EstimatorType.MatrixFactorization:
                case EstimatorType.ImageClassificationMulti:
                case EstimatorType.TextClassificationMulti:
                case EstimatorType.SentenceSimilarityRegression:
                case EstimatorType.ObjectDetectionMulti:
                case EstimatorType.QuestionAnsweringMulti:
                case EstimatorType.NamedEntityRecognitionMulti:
                case EstimatorType.ForecastBySsa:
                    return true;
                default:
                    return false;
            }
        }
    }
}