#include <thread>
#include <fstream>
#include <unordered_map>
#include <iostream>
#include <vector>
#include <mutex>

template<class Callback>
struct WordCounter {
    Callback callback;

    void operator()(std::string filename) {
        std::ifstream myFile(filename);
        if (myFile) {
            countWords(myFile);
        }
        else {
            std::cerr << "Unable to open " + filename << '\n';
        }
    }
    void countWords(std::ifstream& filestream) {
        std::unordered_map<std::string, int> wordCount;
        std::string word;
        while (!filestream.eof() && !filestream.fail()) {
            filestream >> word;
            wordCount[word] += 1;
        }
        callback(wordCount);
    }
};

template<class Callback>
WordCounter<Callback> makeWordCounter(Callback const& func) {
    return WordCounter<Callback>{func};
}

using std::unordered_map;
using std::string;
using std::vector;

unordered_map<string, int> countWords(vector<string> const& filenames) {
    vector<std::thread> threads;
    threads.reserve(filenames.size());

    std::mutex map_lock;
    std::unordered_map<std::string, int> totalWordCount;

    // Define the callback function
    // This operation is basically free
    // Internally, it just copies a reference to the mutex and a reference
    // to the totalWordCount
    auto callback = [&](unordered_map<string, int> const& partial_count) {
        map_lock.lock();
        for(auto count : partial_count) {
            totalWordCount[count.first] += count.second;
        }
        map_lock.unlock();
    };

    // Create a new thread for each file
    for(auto& file : filenames) {
        threads.push_back(std::thread(makeWordCounter(callback), file));
    }

    // Wait until all threads have finished
    for(auto& thread : threads) {
        thread.join();
    }

    return totalWordCount;
}