using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using Microsoft.ML;
using Microsoft.ML.Data;
using Microsoft.ML.Vision;

namespace VehicleModelTraining
{
    class Program
    {
        static void Main(string[] args)
        {
            // Caminho para os diretórios de treinamento (organizados por categoria)
            string datasetFolder = "C:\\Users\\USER\\Pictures\\Veiculos";
            string workspaceFolder = "C:\\Users\\USER\\Pictures\\WorkSpace";
            string outputModelPath = "C:\\Users\\USER\\Pictures\\neMmodel.zip";

            // Inicializar MLContext
            MLContext mlContext = new MLContext(seed: 1);

            var preprocessingPipeline =
                                     mlContext
                                            .Transforms
                                            .LoadRawImageBytes(
                                                outputColumnName: "ImageBytes",
                                                imageFolder: datasetFolder,
                                                inputColumnName: "ImagePath")
                                        .Append(
                                        mlContext
                                                .Transforms
                                                .Conversion
                                                .MapValueToKey(
                                                    inputColumnName: "Label",
                                                    outputColumnName: "LabelAsKey"
                                                 )
                                               );

            // Carregar dados - usamos um método mais direto aqui
            var images = LoadImagesFromDirectory(datasetFolder);
            Console.WriteLine($"Número total de imagens carregadas: {images.Count}");

            // Verificar se há imagens
            if (images.Count == 0)
            {
                Console.WriteLine("Nenhuma imagem foi carregada. Verifique o caminho e as extensões suportadas.");
                return;
            }

            // Criar o DataFrame
            IDataView imageData = mlContext.Data.LoadFromEnumerable(images);

            IDataView shuffledImageDataView = mlContext
                                    .Data
                                    .ShuffleRows(imageData, 0);


            Console.WriteLine("Pre processing images....");
            var timestamp = DateTime.Now;

            // Pre Process images and split into train/test/validation
            IDataView preProcessedImageDataView = preprocessingPipeline
                .Fit(shuffledImageDataView)
                .Transform(shuffledImageDataView);

            Console.WriteLine($"Image preprocessing done in {(DateTime.Now - timestamp).TotalSeconds} seconds");
            Console.WriteLine();

            var firstSplit = mlContext
                   .Data
                   .TrainTestSplit(data: preProcessedImageDataView,
                                   testFraction: 0.3,
                                   seed: 0);
            var trainSet = firstSplit.TrainSet;

            var secondSplit = mlContext
                                       .Data
                                       .TrainTestSplit(data: firstSplit.TestSet,
                                                       testFraction: 0.5, seed: 0);

            var validationSet = secondSplit.TrainSet;
            var testSet = secondSplit.TestSet;


            var classifierOptions = new ImageClassificationTrainer.Options()
            {
                FeatureColumnName = "ImageBytes",
                LabelColumnName = "LabelAsKey",
                Arch = ImageClassificationTrainer.Architecture.InceptionV3,
                //Arch = ImageClassificationTrainer.Architecture.MobilenetV2,
                //Arch = ImageClassificationTrainer.Architecture.ResnetV250,

                TestOnTrainSet = false,
                ValidationSet = validationSet,

                ReuseTrainSetBottleneckCachedValues = true,
                ReuseValidationSetBottleneckCachedValues = true,
                WorkspacePath = workspaceFolder,

                MetricsCallback = Console.WriteLine
            };

            var trainingPipeline = mlContext
                                    .MulticlassClassification
                                    .Trainers
                                    .ImageClassification(classifierOptions)
                                    .Append(mlContext
                                                .Transforms
                                                .Conversion
                                                .MapKeyToValue("PredictedLabel"));

            Console.WriteLine("Training model....");
            timestamp = DateTime.Now;

            var trainedModel = trainingPipeline.Fit(trainSet);

            Console.WriteLine($"Model training done in {(DateTime.Now - timestamp).TotalSeconds} seconds");
            Console.WriteLine();

            Console.WriteLine("Calculating metrics...");

            IDataView evaluationData = trainedModel.Transform(testSet);
            var metrics = mlContext
                            .MulticlassClassification
                            .Evaluate(evaluationData, "LabelAsKey");

            Console.WriteLine($"LogLoss: {metrics.LogLoss}");
            Console.WriteLine($"LogLossReduction: {metrics.LogLossReduction}");
            Console.WriteLine($"MicroAccuracy: {metrics.MicroAccuracy}");
            Console.WriteLine($"MacroAccuracy: {metrics.MacroAccuracy}");
            Console.WriteLine();
            Console.WriteLine($"{metrics.ConfusionMatrix.GetFormattedConfusionTable()}");

            Console.WriteLine();
            Console.WriteLine("Saving model");

            Directory.CreateDirectory("Model");
            mlContext.Model.Save(trainedModel, preProcessedImageDataView.Schema, outputModelPath);

            //Console.WriteLine();
            //// Exibir informações das categorias
            //var categoryCount = images.GroupBy(x => x.Label).Select(g => new { Category = g.Key, Count = g.Count() }).ToList();
            //foreach (var category in categoryCount)
            //{
            //    Console.WriteLine($"Categoria: {category.Category}, Contagem: {category.Count}");
            //}

            //// Dividir em dados de treinamento e teste
            //var dataSplit = mlContext.Data.TrainTestSplit(imageData, testFraction: 0.2);
            //var trainSet = dataSplit.TrainSet;
            //var testSet = dataSplit.TestSet;

            //// IMPORTANTE: Para o ImageClassificationTrainer, vamos usar a abordagem recomendada
            //var classifierOptions = new ImageClassificationTrainer.Options()
            //{
            //    FeatureColumnName = "Image", // Alterado para 'Image' em vez de 'Input'
            //    LabelColumnName = "LabelKey",
            //    ValidationSet = testSet,
            //    Arch = ImageClassificationTrainer.Architecture.ResnetV2101,
            //    Epoch = 10,
            //    BatchSize = 5,
            //    LearningRate = 0.01f,
            //    MetricsCallback = (metrics) => Console.WriteLine(metrics),
            //    WorkspacePath = Path.Combine(Environment.CurrentDirectory, "workspace"),
            //    //TrainDatasetPath = Path.Combine(Environment.CurrentDirectory, "trainFolder"),
            //    //TestDatasetPath = Path.Combine(Environment.CurrentDirectory, "testFolder"),
            //    //OutputTensorName = "softmax2_pre_activation", // Nome da camada de saída do Resnet
            //    ReuseValidationSetBottleneckCachedValues = true,
            //    ReuseTrainSetBottleneckCachedValues = true
            //};

            //try
            //{
            //    // Criar diretórios necessários
            //    EnsureDirectoryExists(classifierOptions.WorkspacePath);
            //    //EnsureDirectoryExists(classifierOptions.TrainDatasetPath);
            //    //EnsureDirectoryExists(classifierOptions.TestDatasetPath);

            //    // Construir pipeline com passos específicos para o ImageClassificationTrainer
            //    var pipeline = mlContext.Transforms.LoadImages(
            //            outputColumnName: "Image",
            //            imageFolder: "",
            //            inputColumnName: nameof(ImageData.ImagePath))
            //        .Append(mlContext.Transforms.Conversion.MapValueToKey(
            //            outputColumnName: "LabelKey",
            //            inputColumnName: nameof(ImageData.Label)))
            //        .Append(mlContext.MulticlassClassification.Trainers.ImageClassification(classifierOptions))
            //        .Append(mlContext.Transforms.Conversion.MapKeyToValue(
            //            outputColumnName: "PredictedLabel",
            //            inputColumnName: "PredictedLabel"));

            //    Console.WriteLine("\nIniciando o treinamento do modelo...");

            //    // Treinar o modelo
            //    ITransformer trainedModel = pipeline.Fit(trainSet);
            //    Console.WriteLine("Treinamento concluído com sucesso!");

            //    // Salvar o modelo
            //    mlContext.Model.Save(trainedModel, trainSet.Schema, outputModelPath);
            //    Console.WriteLine($"Modelo salvo em: {outputModelPath}");

            //    // Avaliar o modelo
            //    EvaluateModel(mlContext, testSet, trainedModel);

            //    // Opcional: Remover diretórios temporários após o uso
            //    //CleanupDirectory(classifierOptions.WorkspacePath);
            //    //CleanupDirectory(classifierOptions.TrainDatasetPath);
            //    //CleanupDirectory(classifierOptions.TestDatasetPath);
            //}
            //catch (Exception ex)
            //{
            //    Console.WriteLine($"Erro durante o treinamento: {ex.Message}");
            //    Console.WriteLine($"Detalhes: {ex.StackTrace}");

            //    if (ex.InnerException != null)
            //    {
            //        Console.WriteLine($"Inner Exception: {ex.InnerException.Message}");
            //        Console.WriteLine($"Inner Stack Trace: {ex.InnerException.StackTrace}");
            //    }
            //}
        }

        private static void EnsureDirectoryExists(string path)
        {
            if (!Directory.Exists(path))
            {
                Directory.CreateDirectory(path);
                Console.WriteLine($"Diretório criado: {path}");
            }
        }

        private static void CleanupDirectory(string path)
        {
            if (Directory.Exists(path))
            {
                try
                {
                    Directory.Delete(path, recursive: true);
                    Console.WriteLine($"Diretório removido: {path}");
                }
                catch (Exception ex)
                {
                    Console.WriteLine($"Erro ao remover diretório {path}: {ex.Message}");
                }
            }
        }

        // Avaliar o desempenho do modelo nos dados de teste
        private static void EvaluateModel(MLContext mlContext, IDataView testData, ITransformer trainedModel)
        {
            IDataView predictions = trainedModel.Transform(testData);
            var metrics = mlContext.MulticlassClassification.Evaluate(predictions, labelColumnName: "LabelKey", predictedLabelColumnName: "PredictedLabel");

            Console.WriteLine($"Acurácia micro: {metrics.MicroAccuracy:0.###}");
            Console.WriteLine($"Acurácia macro: {metrics.MacroAccuracy:0.###}");
            Console.WriteLine($"Pontuação de perda de log: {metrics.LogLoss:#.###}");

            // Exibir matriz de confusão
            Console.WriteLine("Matriz de confusão:");
            Console.WriteLine(metrics.ConfusionMatrix.GetFormattedConfusionTable());
        }

        // Carregar imagens do diretório
        private static List<ImageData> LoadImagesFromDirectory(string folder)
        {
            List<ImageData> images = new List<ImageData>();

            try
            {
                var files = Directory.GetFiles(folder, "*", searchOption: SearchOption.AllDirectories);
                var imageFiles = files.Where(file =>
                    Path.GetExtension(file).ToLower() == ".jpg"  ||
                    Path.GetExtension(file).ToLower() == ".jpeg" ||
                    Path.GetExtension(file).ToLower() == ".png"  ||
                    Path.GetExtension(file).ToLower() == ".webp").ToList();

                Console.WriteLine($"Encontrados {imageFiles.Count} arquivos de imagem em {folder}");

                foreach (var file in imageFiles)
                {
                    // Verificar se o arquivo existe
                    if (!File.Exists(file))
                    {
                        Console.WriteLine($"Arquivo não encontrado: {file}");
                        continue;
                    }

                    // Determinar a categoria com base no nome da pasta pai
                    var label = Directory.GetParent(file).Name;
                    Console.WriteLine($"Processando imagem: {file}, Label: {label}");

                    images.Add(new ImageData()
                    {
                        ImagePath = file,
                        Label = label
                    });
                }
            }
            catch (Exception ex)
            {
                Console.WriteLine($"Erro ao carregar imagens: {ex.Message}");
            }

            return images;
        }

        // Classe para representar dados de imagem
        public class ImageData
        {
            public string ImagePath { get; set; }
            public string Label { get; set; }
            public byte[] ImageBytes { get; set; }
        }
    }
}