    /*
        This program demonstrates a minimal example of a Very Small Language Model (VSLM)
        using dlib's deep learning tools. It includes two modes:

        1) --train  : Train a small Transformer-based language model on a character-based
                      corpus extracted from "slm_data.h" (named shakespeare_text).

        2) --generate: Generate new text from a trained model, given an initial prompt
                       extracted from "slm_data.h" (named shakespeare_prompt).

        The "slm_dels.h" header is expected to provide a comprehensive Transformer
        definition with the following key elements:
          - A configurable transformer_config
          - The use of classification_head to output a single token
          - The network_type<true> or network_type<false> for training vs inference
          - The typical dlib constructs (input<matrix<int>>, etc.)

        Character-level tokenization is used here. Each character is directly transformed
        into an integer token. The model attempts to learn the sequence of characters in
        shakespeare_text. Then you can ask the model to generate new text from a short
        prompt.

        This model is intentionally kept small (few neurons/parameters) to ensure
        simplicity and efficiency. As a result, it may not generalize well to unseen
        patterns or concepts. However, it effectively illustrates the principle of
        attention and the ability to perfectly memorize and reproduce sequences from
        the training data. This makes it a useful educational tool for understanding
        the mechanics of Transformer models, even if it lacks the capacity for
        sophisticated language understanding.
    */

    #include <iostream>
    #include <string>
    #include <vector>
    #include <algorithm>
    #include <cmath>
    #include <random>
    #include <dlib/data_io.h>
    #include <dlib/cmd_line_parser.h>
    #include <dlib/misc_api.h>

    // Include Transformer definitions
    #include "slm_defs.h"

    // This header "slm_data.h" is assumed to contain:
    //   const std::string shakespeare_text;
    //   const std::string shakespeare_prompt;
    #include "slm_data.h"

    // ----------------------------------------------------------------------------------------

    using namespace std;
    using namespace dlib;

    // We treat each character as a token ID in [0..255].
    const int MAX_TOKEN_ID = 255;
    const int PAD_TOKEN = 256; // an extra "pad" token if needed

    // For simplicity, we assume each line from shakespeare_text is appended, ignoring them.
    std::vector<int> char_based_tokenize(const std::string& text)
    {
        std::vector<int> tokens;
        tokens.reserve(text.size());
        for (const int c : text)
        {
            tokens.push_back(std::min(c, MAX_TOKEN_ID));
        }
        return tokens;
    }

    // Function to shuffle samples and labels in sync
    void shuffle_samples_and_labels(std::vector<matrix<int, 0, 1>>& samples, std::vector<unsigned long>& labels) {
        std::vector<size_t> indices(samples.size());
        std::iota(indices.begin(), indices.end(), 0); // Fill with 0, 1, 2, ..., N-1
        std::shuffle(indices.begin(), indices.end(), std::default_random_engine{});

        // Create temporary vectors to hold shuffled data
        std::vector<matrix<int, 0, 1>> shuffled_samples(samples.size());
        std::vector<unsigned long> shuffled_labels(labels.size());

        // Apply the shuffle
        for (size_t i = 0; i < indices.size(); ++i)
        {
            shuffled_samples[i] = samples[indices[i]];
            shuffled_labels[i] = labels[indices[i]];
        }

        // Replace the original data with shuffled data
        samples = std::move(shuffled_samples);
        labels = std::move(shuffled_labels);
    }

    // ----------------------------------------------------------------------------------------

    int main(int argc, char** argv)
    {
        try
        {
            command_line_parser parser;
            parser.add_option("train", "Train a small transformer on the built-in Shakespeare text");
            parser.add_option("generate", "Generate text from a previously trained model (needs shakespeare_prompt)");
            parser.add_option("learning-rate", "Set the learning rate for training (default: 1e-4)", 1);
            parser.add_option("batch-size", "Set the mini-batch size for training (default: 64)", 1);
            parser.add_option("generation-length", "Set the length of generated text (default: 400)", 1);
            parser.add_option("alpha", "Set the weight decay for Adam optimizer (default: 0.004)", 1);
            parser.add_option("beta1", "Set the first moment coefficient (default: 0.9)", 1);
            parser.add_option("beta2", "Set the second moment coefficient (default: 0.999)", 1);
            parser.add_option("max-samples", "Set the maximum number of training samples (default: 50000)", 1);
            parser.add_option("shuffle", "Shuffle training sequences and labels before training (default: false)");
            parser.parse(argc, argv);

            if (parser.number_of_arguments() == 0 && !parser.option("train") && !parser.option("generate"))
            {
                parser.print_options();
                return 0;
            }

            // Default values
            const double learning_rate = get_option(parser, "learning-rate", 1e-4);
            const long batch_size = get_option(parser, "batch-size", 64);
            const int generation_length = get_option(parser, "generation-length", 400);
            const double alpha = get_option(parser, "alpha", 0.004);       // Initial learning rate for Adam
            const double beta1 = get_option(parser, "beta1", 0.9);         // Decay rate for the first moment estimate
            const double beta2 = get_option(parser, "beta2", 0.999);       // Decay rate for the second moment estimate
            const size_t max_samples = get_option(parser, "max-samples",50000); // Default maximum number of training samples

            // We define a minimal config for demonstration
            const long vocab_size = MAX_TOKEN_ID + 1 + 1;   // 256 for chars + 1 pad token
            const long num_layers = 3;
            const long num_heads = 4;
            const long embedding_dim = 64;
            const long max_seq_len = 80;   // a small sequence length for the example
            const bool use_squeezing = false;

            using my_transformer_cfg = transformer::transformer_config<
                vocab_size,
                num_layers,
                num_heads,
                embedding_dim,
                max_seq_len,
                use_squeezing,
                gelu,
                dropout_10
            >;

            // For GPU usage (if any), set gpus = {0} for a single GPU, etc.
            std::vector<int> gpus{ 0 };

            // The model file to store or load
            const std::string model_file = "shakespeare_lm_char_model.dat";

            // ----------------------------------------------------------------------------------------
            // Train mode
            // ----------------------------------------------------------------------------------------
            if (parser.option("train"))
            {
                cout << "=== TRAIN MODE ===\n";

                // 1) Prepare training data (simple approach)
                // We will store characters from shakespeare_text into a vector
                // and then produce training samples of length (max_seq_len+1),
                // where the last token is the label to predict from the preceding max_seq_len.
                auto full_tokens = char_based_tokenize(shakespeare_text);
                if (full_tokens.empty())
                {
                    cerr << "ERROR: The Shakespeare text is empty. Please provide a valid training text.\n";
                    return 0;
                }

                // Calculate the maximum number of sequences
                size_t max_sequences = (full_tokens.size() > (size_t)max_seq_len + 1)
                    ? (full_tokens.size() - ((size_t)max_seq_len + 1))
                    : 0;

                // Display the size of the training text and the number of sequences
                cout << "Training text size: " << full_tokens.size() << " characters\n";
                cout << "Maximum number of sequences: " << max_sequences << "\n";

                // Check if the text is too short
                if (max_sequences == 0)
                {
                    cerr << "ERROR: The Shakespeare text is too short for training. It must contain at least "
                        << (max_seq_len + 1) << " characters.\n";
                    return 0;
                }

                std::vector<matrix<int, 0, 1>> samples;
                std::vector<unsigned long> labels;

                // Let's create a training set of about (N) samples from the text
                // Each sample: [x0, x1, ..., x_(max_seq_len-1)] -> y
                // We'll store them in "samples" and "labels".
                const size_t N = (max_sequences < max_samples) ? max_sequences : max_samples;
                for (size_t start = 0; start < N; ++start)
                {
                    matrix<int, 0, 1> seq(max_seq_len, 1);
                    for (long t = 0; t < max_seq_len; ++t)
                        seq(t, 0) = full_tokens[start + t];
                    samples.push_back(seq);
                    labels.push_back(full_tokens[start + max_seq_len]);
                }

                // Shuffle samples and labels if the --shuffle option is enabled
                if (parser.option("shuffle"))
                {
                    cout << "Shuffling training sequences and labels...\n";
                    shuffle_samples_and_labels(samples, labels);
                }

                // 3) Construct the network in training mode
                using net_type = my_transformer_cfg::network_type<true>;
                net_type net;
                if (file_exists(model_file))
                    deserialize(model_file) >> net;

                // 4) Create dnn_trainer
                dnn_trainer<net_type, adam> trainer(net, adam(alpha, beta1, beta2), gpus);
                trainer.set_learning_rate(learning_rate);
                trainer.set_min_learning_rate(1e-6);
                trainer.set_mini_batch_size(batch_size);
                trainer.set_iterations_without_progress_threshold(15000);
                trainer.set_max_num_epochs(400);
                trainer.be_verbose();

                // 5) Train
                trainer.train(samples, labels);

                // 6) Evaluate quickly on the training set
                auto predicted = net(samples);
                size_t correct = 0;
                for (size_t i = 0; i < labels.size(); ++i)
                    if (predicted[i] == labels[i])
                        correct++;
                double accuracy = (double)correct / labels.size();
                cout << "Training accuracy (on this sample set): " << accuracy << "\n";

                // 7) Save the model
                net.clean();
                serialize(model_file) << net;
                cout << "Model saved to " << model_file << "\n";
            }

            // ----------------------------------------------------------------------------------------
            // Generate mode
            // ----------------------------------------------------------------------------------------
            if (parser.option("generate"))
            {
                cout << "=== GENERATE MODE ===\n";
                // 1) Load the trained model
                using net_infer = my_transformer_cfg::network_type<false>;
                net_infer net;
                if (file_exists(model_file))
                {
                    deserialize(model_file) >> net;
                    cout << "Loaded model from " << model_file << "\n";
                }
                else
                {
                    cerr << "Error: model file not found. Please run --train first.\n";
                    return 0;
                }
                cout << my_transformer_cfg::model_info::describe() << endl;
                cout << "Model parameters: " << count_parameters(net) << endl << endl;

                // 2) Get the prompt from the included slm_data.h
                std::string prompt_text = shakespeare_prompt;
                if (prompt_text.empty())
                {
                    cerr << "No prompt found in slm_data.h.\n";
                    return 0;
                }
                // If prompt is longer than max_seq_len, we keep only the first window
                if (prompt_text.size() > (size_t)max_seq_len)
                    prompt_text.erase(prompt_text.begin() + max_seq_len, prompt_text.end());

                // Convert prompt to a token sequence
                const auto prompt_tokens = char_based_tokenize(prompt_text);

                // Put into a dlib matrix
                matrix<int, 0, 1> input_seq(max_seq_len, 1);
                // Fill with pad if prompt is shorter than max_seq_len
                for (long i = 0; i < max_seq_len; ++i)
                {
                    if ((size_t)i < prompt_tokens.size())
                        input_seq(i, 0) = prompt_tokens[i];
                    else
                        input_seq(i, 0) = PAD_TOKEN;
                }

                cout << "\nInitial prompt:\n" << prompt_text << " (...)\n\n\nGenerated text:\n" << prompt_text;

                // 3) Generate new text
                // We'll predict one character at a time, then shift the window
                for (int i = 0; i < generation_length; ++i)
                {
                    const int next_char = net(input_seq); // single inference

                    // Print the generated character
                    cout << static_cast<char>(std::min(next_char, MAX_TOKEN_ID)) << flush;

                    // Shift left by 1
                    for (long i = 0; i < max_seq_len - 1; ++i)
                        input_seq(i, 0) = input_seq(i + 1, 0);
                    input_seq(max_seq_len - 1, 0) = std::min(next_char, MAX_TOKEN_ID);
                }

                cout << "\n\n(end of generation)\n";
            }

            return 0;
        }
        catch (exception& e)
        {
            cerr << "Exception thrown: " << e.what() << endl;
            return 1;
        }
    }

    /*
     * This program demonstrates the training of a language model on about 15k sequences.
     * The training process produces a data file of approximately 32MB on disk.
     *
     * - Transformer model configuration:
     *    + vocabulary size: 257
     *    + layers: 3
     *    + attention heads: 4
     *    + embedding dimension: 64
     *    + max sequence length: 80
     * - Number of parameters: 8,247,496
     *
     * The training cab be done using the following command line:
     * >./slm_basic_train_ex --train --shuffle
     *
     * After this phase, the model achieves perfect prediction accuracy (i.e acc=1).
     * The generation option produces text that is very close to the original training data,
     * as illustrated by the example below:
     * > Generated text:
     * > QUEEN ELIZABETH:
     * > But thou didst kill my children.
     * >
     * > KING RICHARD III:
     * > But in your daughter's womb I bury them:
     * > Where in that nest of spicery they shall breed
     * > Selves of themselves, to your recomforture.
     * >
     * > QUEEN ELIZABETH:
     * > Shall I go win my daughter to thy will?
     * >
     * > KING RICHARD III:
     * > And be a happy mother by the deed.
     * >
     * > QUEEN ELIZABETH:
     * > I go. Write to me very shortly.
     * > And you shall understand from me her mind.
     */