#include <iostream>
#include <vector>
#include <string>
#include <map>
#include <cmath>

/**
 * @brief A simple Naive Bayes Classifier component.
 * Inspired by Chapter 2 of the "Building AI" course (Elements of AI).
 */

class NaiveBayes {
public:
    void train(const std::vector<std::string>& texts, const std::vector<int>& labels) {
        for (size_t i = 0; i < texts.size(); ++i) {
            int label = labels[i];
            class_counts[label]++;
            total_samples++;
            
            // Simple word tokenization (splitting by space)
            std::string word;
            for (char c : texts[i]) {
                if (c == ' ') {
                    word_counts[label][word]++;
                    word;
                } else {
                    word += c;
                }
            }
            if (!word.empty()) word_counts[label][word]++;
        }
    }

    int predict(const std::string& text) {
        double best_prob = -1e18;
        int best_label = -1;

        for (auto const& [label, count] : class_counts) {
            double log_prob = std::log((double)count / total_samples);
            
            std::string word;
            for (char c : text) {
                if (c == ' ') {
                    log_prob += calculate_word_log_prob(label, word);
                    word;
                } else {
                    word += c;
                }
            }
            if (!word.empty()) log_prob += calculate_word_log_prob(label, word);

            if (log_prob > best_prob) {
                best_prob = log_prob;
                best_label = label;
            }
        }
        return best_label;
    }

private:
    std::map<int, int> class_counts;
    std::map<int, std::map<std::string, int>> word_counts;
    int total_samples = 0;

    double calculate_word_log_prob(int label, const std::string& word) {
        // Laplace smoothing
        int count = word_counts[label][word];
        int total_words_in_class = 0;
        for (auto const& [w, c] : word_counts[label]) total_words_in_class += c;
        
        return std::log((double)(count + 1) / (total_words_in_class + 1000)); // Assuming vocab size 1000
    }
};

int main() {
    std::cout << "--- Naive Bayes AI Component ---" << std::endl;
    
    NaiveBayes nb;
    nb.train({"good great awesome", "bad terrible awful"}, {1, 0});
    
    std::string test = "great awesome";
    int prediction = nb.predict(test);
    
    std::cout << "Text: \"" << test << "\"" << std::endl;
    std::cout << "Prediction: " << (prediction == 1 ? "Positive" : "Negative") << std::endl;

    return 0;
}