main.cpp

//
// Created by MattFor on 12.04.2025.
//

#define DEV_MODE 0

#include <ranges>
#include <vector>
#include <string>
#include <random>
#include <cstdlib>
#include <utility>
#include <iostream>
#include <algorithm>
#include <filesystem>
#include <functional>
#include <unordered_map>

#include "Utilities.h"
#include "NeuralNetwork.h"
#include "DataTransformations.h"

// --- Main Program Loop ---
int main()
{
    try // Top-level error handling
    {
        std::cout << "--- Neural Network by MattFor ---\n" << '\n';

        // --- Dataset Processors Map ---
        const std::unordered_map<std::string, DataProcessingFunc> dataset_processors = {
            {
                "car.data",
                process_car
            },
            {
                "mnist",
                process_mnist
            }
            // Add more data here
        };

        // --- Persistent State ---
        std::vector<std::vector<double>> loaded_train_inputs,  loaded_test_inputs;
        std::vector<std::vector<double>> loaded_train_targets, loaded_test_targets;
        std::vector<std::size_t>         loaded_layer_sizes;
        auto                             loaded_hidden_activation = ActivationType::RELU;
        auto                             loaded_output_activation = ActivationType::SIGMOID;
        std::string                      current_dataset_name;
        bool                             dataset_loaded       = false;
        bool                             architecture_defined = false;

        bool running       = true;
        int  current_cycle = 0;

        while (running)
        {
            std::cout << "\n===================== " << "Cycle #" << current_cycle << " =====================\n" << '\n';

            // --- Per-Iteration State ---
            bool                             weights_were_loaded = false; // Track if weights were loaded in this iteration
            std::size_t                      epochs              = 100;
            double                           learning_rate       = 0.01;
            constexpr double                 train_split_ratio   = 0.8;
            std::size_t                      num_targets         = 0;
            std::size_t                      num_attributes      = 0;
            std::vector<std::size_t>         layer_sizes;
            std::vector<std::vector<double>> train_inputs,  test_inputs;
            std::vector<std::vector<double>> train_targets, test_targets;
            auto                             hidden_activation_choice = ActivationType::RELU;
            auto                             output_activation_choice = ActivationType::SIGMOID;

            // --- Dataset Selection ---
            std::cout << "Available data:" << '\n';
            for (const auto& name : std::ranges::views::keys(dataset_processors))
            {
                std::cout << " - " << name << '\n';
            }
            std::cout << " - Randomly Generated Data" << '\n';
            std::cout << '\n';

            // --- Data Loading and Processing ---
            if (auto dataset_choice = get_input<std::string>("Enter dataset name", "Randomly Generated Data"); dataset_loaded && dataset_choice == current_dataset_name)
            {
                std::cout << "\nReusing previously loaded dataset: " << current_dataset_name << '\n';
                train_inputs             = loaded_train_inputs;
                train_targets            = loaded_train_targets;
                test_inputs              = loaded_test_inputs;
                test_targets             = loaded_test_targets;
                num_attributes           = loaded_train_inputs[0].size();
                num_targets              = loaded_train_targets[0].size();
                layer_sizes              = loaded_layer_sizes;
                hidden_activation_choice = loaded_hidden_activation;
                output_activation_choice = loaded_output_activation;
            }
            else
            {
                // --- Processing based on choice ---
                if (const auto processor_iter = dataset_processors.find(dataset_choice); processor_iter != dataset_processors.end()) // Existing dataset chosen
                {
                    // No pre-flight existence check here: the dataset key ("car.data", "mnist")
                    // is not a path, and each loader resolves and reports its own files.
                    std::cout << "\nProcessing dataset: " << dataset_choice << "..." << '\n';

                    DataPair all_data;
                    execute_timed("Loading and processing data...", "Data processed!", [&]
                    {
                        all_data = processor_iter->second();
                    });

                    auto& [all_inputs, all_targets] = all_data;
                    if (all_inputs.empty() || all_targets.empty() || all_inputs.size() != all_targets.size())
                    {
                        throw std::runtime_error("Loaded dataset is empty or mismatched.");
                    }

                    num_attributes = all_inputs[0].size();
                    num_targets    = all_targets[0].size();
                    std::cout << "\nDataset loaded: " << all_inputs.size() << " samples." << '\n';
                    std::cout << "Input attributes per sample: " << num_attributes << '\n';
                    std::cout << "Target outputs per sample: " << num_targets << '\n';

                    execute_timed("\nSplitting data...", "Data split!", [&]
                    {
                        auto [train_data_pair, test_data_pair] = split_data(all_inputs, all_targets, train_split_ratio);
                        train_inputs                           = std::move(train_data_pair.first);
                        train_targets                          = std::move(train_data_pair.second);
                        test_inputs                            = std::move(test_data_pair.first);
                        test_targets                           = std::move(test_data_pair.second);
                    });

                    std::cout << "Training samples: " << train_inputs.size() << '\n';
                    std::cout << "Testing samples: " << test_inputs.size() << "\n\n";

                    if constexpr (DEV_MODE)
                    {
                        if (!train_inputs.empty())
                        {
                            display_sample_data("Training Data Sample", train_inputs, train_targets, 5);
                        }
                        else
                        {
                            std::cout << "No training data generated to display." << '\n';
                        }
                    }

                    // Store loaded dataset
                    loaded_train_inputs  = train_inputs;
                    loaded_train_targets = train_targets;
                    loaded_test_inputs   = test_inputs;
                    loaded_test_targets  = test_targets;
                    current_dataset_name = dataset_choice;
                    dataset_loaded       = true;
                    architecture_defined = false; // Reset architecture as dataset might have changed
                }
                else if (dataset_choice == "Randomly Generated Data")
                {
                    std::cout << "\nGenerating random data..." << '\n';
                    num_attributes               = get_input<std::size_t>("Enter number of input attributes", 10);
                    num_targets                  = get_input<std::size_t>("Enter number of target classes (output neurons)", 4);
                    const auto num_train_samples = get_input<std::size_t>("Enter number of training samples", 10000);
                    const auto num_test_samples  = get_input<std::size_t>("Enter number of test samples", 2000);

                    // Define Network Topology (for random data, define it before generating)
                    std::size_t hidden1_size = std::max(static_cast<size_t>(10), ( num_attributes + num_targets ) / 2);
                    loaded_layer_sizes       = {
                        num_attributes,
                        hidden1_size,
                        num_targets
                    };
                    architecture_defined = true;
                    layer_sizes          = loaded_layer_sizes; // Set for this iteration as well

                    std::random_device rd;
                    std::mt19937       rng(rd());

                    execute_timed("Generating training data...", "Training data generated!", generate_data_network, std::ref(train_inputs), std::ref(train_targets), num_train_samples, num_attributes, num_targets, std::ref(rng));

                    execute_timed("Generating test data...", "Test data generated!", generate_data_network, std::ref(test_inputs), std::ref(test_targets), num_test_samples, num_attributes, num_targets, std::ref(rng));

                    if constexpr (DEV_MODE)
                    {
                        if (!train_inputs.empty())
                        {
                            display_sample_data("Generated Training Data Sample", train_inputs, train_targets, 5);
                        }
                        else
                        {
                            std::cout << "No training data generated to display." << '\n';
                        }
                    }

                    // Store generated dataset
                    loaded_train_inputs  = train_inputs;
                    loaded_train_targets = train_targets;
                    loaded_test_inputs   = test_inputs;
                    loaded_test_targets  = test_targets;
                    current_dataset_name = dataset_choice;
                    dataset_loaded       = true;
                }
                else // Invalid choice
                {
                    throw std::runtime_error("Invalid dataset choice: '" + dataset_choice + "'");
                }
            }

            // --- Network Architecture Selection ---
            if (!architecture_defined || layer_sizes.empty())
            {
                auto complexity = NetworkComplexity::MEDIUM;

                auto complexity_choice = get_input<std::string>("Do you want to set the network complexity yourself (y), use a predefined level (p), or automatically determine it (a)?", "a");

                if (train_inputs[0].size() < 10)
                {
                    complexity = NetworkComplexity::SIMPLE;
                }
                else if (train_inputs[0].size() < 150)
                {
                    complexity = NetworkComplexity::MEDIUM;
                }
                else
                {
                    complexity = NetworkComplexity::COMPLEX;
                }

                // Generate the layer sizes
                execute_timed("\nGenerating network topology...", "Topology generated!", [&]
                {
                    if (complexity_choice == "y")
                    {
                        std::cout << "Enter the layer sizes separated by spaces (f.e " << num_attributes << " 128 " << num_targets << ") [FIRST AND LAST LAYER MUST ADHERE TO THE EXAMPLE]:\n";

                        std::string layer_sizes_str;
                        std::getline(std::cin, layer_sizes_str);
                        std::stringstream ss(layer_sizes_str);
                        std::size_t       layer_size;
                        layer_sizes.clear();

                        while (ss >> layer_size)
                        {
                            layer_sizes.push_back(layer_size);
                        }

                        if (layer_sizes.empty())
                        {
                            std::cerr << "Error: No layer sizes provided. Falling back to automatic complexity.\n";
                            complexity_choice = "a";
                        }
                        else if (layer_sizes.front() != num_attributes)
                        {
                            std::cerr << "Error: The first layer size must be equal to the number of attributes (" << num_attributes << "). Falling back to automatic complexity.\n";
                            complexity_choice = "a";
                        }
                        else if (layer_sizes.back() != num_targets)
                        {
                            std::cerr << "Error: The last layer size must be equal to the number of targets (" << num_targets << "). Falling back to automatic complexity.\n";
                            complexity_choice = "a";
                        }
                        else
                        {
                            std::cout << "User-defined layer sizes accepted.\n";
                            loaded_layer_sizes   = layer_sizes;
                            architecture_defined = true;
                            return;
                        }
                    }

                    if (complexity_choice == "p" || complexity_choice == "a")
                    {
                        int complexity_level_input;

                        if (complexity_choice == "p")
                        {
                            complexity_level_input = get_input<int>("Enter the complexity level (0: Simple, 1: Medium, 2: Complex):", 1);
                        }
                        else // complexity_choice == "a"
                        {
                            if (train_inputs[0].size() < 10)
                            {
                                complexity_level_input = 0;
                            }
                            else if (train_inputs[0].size() < 150)
                            {
                                complexity_level_input = 1;
                            }
                            else
                            {
                                complexity_level_input = 2;
                            }

                            std::cout << "Automatically determined complexity level: " << complexity_level_input << '\n';
                        }

                        switch (complexity_level_input)
                        {
                            case 0:
                            {
                                complexity = NetworkComplexity::SIMPLE;
                            }
                            break;

                            case 1:
                            {
                                complexity = NetworkComplexity::MEDIUM;
                            }
                            break;

                            case 2:
                            {
                                complexity = NetworkComplexity::COMPLEX;
                            }
                            break;

                            default:
                            {
                                std::cerr << "Warning: Invalid complexity level. Using Medium.\n";
                                complexity = NetworkComplexity::MEDIUM;
                            }
                            break;
                        }

                        loaded_layer_sizes = NeuralNetwork::generate_layer_topology(num_attributes, num_targets, train_inputs.size(), complexity);

                        layer_sizes          = loaded_layer_sizes;
                        architecture_defined = true;
                    }
                    else
                    {
                        std::cerr << "Warning: Invalid choice for complexity. Falling back to automatic complexity.\n";

                        if (train_inputs[0].size() < 10)
                        {
                            complexity = NetworkComplexity::SIMPLE;
                        }
                        else if (train_inputs[0].size() < 150)
                        {
                            complexity = NetworkComplexity::MEDIUM;
                        }
                        else
                        {
                            complexity = NetworkComplexity::COMPLEX;
                        }

                        loaded_layer_sizes   = NeuralNetwork::generate_layer_topology(num_attributes, num_targets, train_inputs.size(), complexity);
                        layer_sizes          = loaded_layer_sizes;
                        architecture_defined = true;
                    }
                });

                std::cout << "Generated layer sizes: ";
                for (size_t i = 0; i < layer_sizes.size(); i++)
                {
                    std::cout << ( i > 0 ? " -> " : "" ) << layer_sizes[i];
                }
                std::cout << '\n';

                // --- Activation Function Selection ---
                std::cout << "\n--- Activation Function Selection ---" << '\n';
                std::cout << "Choose activation function for HIDDEN layers:" << '\n';
                std::cout << "1: ReLU\n  2: LeakyReLU (alpha=" << LRELU_ALPHA << ")\n  3: Tanh\n  4: Sigmoid\n";
                const int hidden_choice  = get_input<int>("Enter choice for hidden layers", 1);
                loaded_hidden_activation = int_to_activation(hidden_choice);
                hidden_activation_choice = loaded_hidden_activation;

                std::cout << "\nChoose activation function for OUTPUT layer:" << '\n';
                std::cout << "[Usually Sigmoid for binary/multi-label, Tanh for regression [-1,1]]" << '\n';
                std::cout << "1: ReLU\n  2: LeakyReLU (alpha=" << LRELU_ALPHA << ")\n  3: Tanh\n  4: Sigmoid\n";
                const int output_choice  = get_input<int>("Enter choice for output layer", 4);
                loaded_output_activation = int_to_activation(output_choice);
                output_activation_choice = loaded_output_activation;
            }
            else
            {
                std::cout << "\nReusing previously defined network architecture." << '\n';
                layer_sizes              = loaded_layer_sizes;
                hidden_activation_choice = loaded_hidden_activation;
                output_activation_choice = loaded_output_activation;
            }

            // --- Common Steps: Network Creation, Activation, Load/Train, Evaluate, Save ---
            if (train_inputs.empty() || test_inputs.empty())
            {
                throw std::runtime_error("No training or testing data available to proceed.");
            }

            // --- Network Creation ---
            // Create the network structure based on selected data and activations
            NeuralNetwork network = execute_timed("Creating Neural Network Structure...", "Neural Network structure created!", [&]
            {
                return NeuralNetwork(layer_sizes, hidden_activation_choice, output_activation_choice);
            });

            network.display_topology();

            // --- Load Weights or Train ---
            std::cout << "\n--- Action ---" << '\n';
            const auto action_choice = get_input<std::string>("Do you want to [T]rain the network, [L]oad weights, or [C]ontinue training from loaded weights? (T/L/C)", "T");

            const bool wants_load     = action_choice == "L" || action_choice == "l" || action_choice == "C" || action_choice == "c";
            const bool wants_continue = action_choice == "C" || action_choice == "c";

            auto offer_to_save_weights = [&](const NeuralNetwork& trained_network)
            {
                std::cout << "\n--- Save Weights ---" << '\n';
                if (get_input<std::string>("Do you want to save the trained weights? (y/n)", "n") != "y")
                {
                    return;
                }

                std::stringstream ss;
                ss << "network_weights";

                // Include the dataset name in the filename
                if (!current_dataset_name.empty())
                {
                    // Sanitise the dataset name to be filename-friendly
                    std::string sanitized_dataset_name = current_dataset_name;
                    std::ranges::replace(sanitized_dataset_name, '.', '_');
                    std::ranges::replace(sanitized_dataset_name, ' ', '_');
                    ss << "_" << sanitized_dataset_name;
                }
                else
                {
                    ss << "_unknown_dataset"; // Fallback if the dataset name is not available
                }

                for (const size_t size : trained_network.get_layer_sizes())
                {
                    ss << "_" << size;
                }
                ss << ".bin";

                const auto save_filename = get_input<std::string>("Enter filename to save weights to", ss.str());

                execute_timed(std::format("Saving weights to '{}'...", save_filename), "Weights save operation complete.", [&]
                {
                    trained_network.save_weights(save_filename);
                });
            };

            if (wants_load)
            {
                auto load_filename = get_input<std::string>("Enter filename to load weights from", "network_weights.bin");

                if (const std::filesystem::path load_path = resolve_project_path(std::filesystem::path("nn_data") / load_filename); std::filesystem::exists(load_path))
                {
                    bool load_success = false;
                    execute_timed(std::format("Attempting to load weights from '{}'...", load_filename), "Weight loading attempt finished.", // Message neutral as success checked below
                                  [&]
                                  {
                                      load_success = network.load_weights(load_filename);
                                  });

                    if (load_success)
                    {
                        std::cout << "Weights loaded successfully. Network state restored." << '\n';
                        weights_were_loaded = true; // Set flag

                        if (auto display_topology_prompt = get_input<std::string>("Do you want to display loaded topology? (y/n)", "n"); display_topology_prompt == "y" || display_topology_prompt == "Y")
                        {
                            network.display_topology();
                        }

                        if (wants_continue)
                        {
                            std::cout << "\n--- Continue Network Training ---" << '\n';
                            epochs        = get_input<std::size_t>("Enter number of additional training epochs", epochs);
                            learning_rate = get_input<double>("Enter learning rate for continued training", learning_rate);

                            if (train_inputs.empty() || train_targets.empty())
                            {
                                std::cerr << "Cannot continue training: No training data available." << '\n';
                            }
                            else
                            {
                                execute_timed("", "Network Training finished!", [&]
                                {
                                    network.train(train_inputs, train_targets, epochs, learning_rate);
                                });

                                offer_to_save_weights(network);
                            }
                        }
                    }
                    else
                    {
                        std::cerr << "Failed to load weights. Check file and network configuration." << '\n';
                        std::cout << "Proceeding with training a new network." << '\n';
                        // weights_were_loaded remains false
                    }
                }
                else
                {
                    std::cerr << "Weight file '" << load_filename << "' not found under nn_data/. " << "Proceeding with training." << '\n';
                    // weights_were_loaded remains false
                }
            }

            // --- Network Training  ---
            if (!weights_were_loaded)
            {
                std::cout << "\n--- Starting Network Training ---" << '\n';
                if (train_inputs.empty() || train_targets.empty())
                {
                    std::cerr << "Cannot train: No training data available." << '\n';
                }
                else
                {
                    // --- Get training parameters ---
                    epochs        = get_input<std::size_t>("Enter number of training epochs", epochs);
                    learning_rate = get_input<double>("Enter learning rate", learning_rate);

                    execute_timed("", "Network Training finished!", [&]
                    {
                        network.train(train_inputs, train_targets, epochs, learning_rate);
                    });

                    offer_to_save_weights(network);
                }
            }
            else if (!wants_continue)
            {
                std::cout << "\nSkipping training because weights were loaded." << '\n';
            }

            // --- Network Evaluation (Always perform if test data exists) ---
            std::cout << "\n--- Network Evaluation ---" << '\n';
            if (!test_inputs.empty())
            {
                execute_timed("Starting Network Evaluation...", "Network Evaluation finished!", [&]
                {
                    evaluate_network(network, test_inputs, test_targets);
                });
            }
            else
            {
                std::cout << "Skipping evaluation: No test data available." << '\n';
            }

            // --- Loop Continuation ---
            std::cout << '\n';
            if (auto exit_prompt = get_input<std::string>("Do you want to run another cycle? (y/n)", "y"); exit_prompt != "y" && exit_prompt != "Y")
            {
                running = false;
            }

            current_cycle++;
        } // End while(running) loop
    }
    catch (const std::filesystem::filesystem_error& e)
    {
        std::cerr << "\n--- FILESYSTEM ERROR ---\n" << e.what() << "\nCode: " << e.code() << "\nPath1: " << e.path1() << "\nPath2: " << e.path2() << "\n----------------------" << '\n';
        std::cin.get();
        return EXIT_FAILURE;
    }
    catch (const std::exception& e)
    {
        std::cerr << "\n--- ERROR ---\n" << e.what() << "\n-------------" << '\n';
        std::cin.get();
        return EXIT_FAILURE;
    }
    catch (...)
    {
        std::cerr << "\n--- UNKNOWN ERROR ---\nAn unexpected error occurred." << '\n';
        std::cin.get();
        return EXIT_FAILURE;
    }

    // --- Successful Execution Exit ---
    std::cout << "\nExecution finished successfully." << '\n';
    std::cout << "Press ENTER to exit..." << '\n';
    std::cin.get();

    return EXIT_SUCCESS;
}