import Foundation

// Extension to split an array into chunks
extension Array {
    func chunks(_ size: Int) -> [[Element]] {
        return stride(from: 0, to: count, by: size).map {
            Array(self[$0..<Swift.min($0 + size, count)])
        }
    }
}

func generateRandomWords(size: Int) -> [String] {
    let words = ["apple", "banana", "cherry", "date", "elderberry", "fig", "grape", "honeydew", "kiwi", "lemon", "mango", "nectarine", "orange", "peach", "quince", "raspberry", "strawberry", "tangerine", "watermelon"]

    var wordCounts: [String: Int] = [:]
    var randomWords: [String] = []

    // Generate a random word count for each word
    for word in words {
        let count = Int.random(in: 1...50)
        wordCounts[word] = count
    }

    // Create an array of random words
    for (word, count) in wordCounts {
        for _ in 0..<count {
            randomWords.append(word)
        }
    }

    // Return shuffled array of words
    return Array(repeating: randomWords, count: size/randomWords.count + 1).flatMap { $0 }.prefix(size).shuffled()
}

func getWordCountConcurrently(words: [String]) -> [String: Int] {
    let queue = DispatchQueue(label: "io.greedydev.wordcount", attributes: .concurrent)
    let group = DispatchGroup()

    // Use GCD to divide the work betweek multiple threads
    let chunkSize = words.count / ProcessInfo.processInfo.activeProcessorCount
    let chunks = words.chunks(chunkSize)

    var wordCounts: [String: Int] = [:]

    // Use map-reduce to count the occurrences of each word in each chunk and add them up
    for chunk in chunks {
        group.enter()
        queue.async {
            let count = chunk.reduce(into: [:]) { counts, word in
                counts[word, default: 0] += 1
            }
            queue.async(flags: .barrier) {
                for (word, occurrences) in count {
                    wordCounts[word, default: 0] += occurrences
                }
                group.leave()
            }
        }
    }

    // Wait for all the work to complete
    group.wait()

    return wordCounts
}


func getWordCountSync(words: [String]) -> [String: Int] {
    let wordCounts = words
        .map {
            ($0, 1)
        }
        .reduce(into: [String: Int]()) { wordCounts, tuple in
            let (word, count) = tuple
            wordCounts[word, default: 0] += count
        }

    return wordCounts
}

// Function to measure execution time
func measureTime(_ block: () -> Void) -> TimeInterval {
    let startTime = Date()
    block()
    let endTime = Date()
    return endTime.timeIntervalSince(startTime)
}

// Input array of words
let randomWords = generateRandomWords(size: 10_000_000)
print(randomWords.count)

var concurrentWordCount = [String: Int]()
let concurrentExecTime = measureTime {
    concurrentWordCount = getWordCountConcurrently(words: randomWords)
}

print("\(concurrentWordCount)")
print("Concurrent: \(concurrentExecTime)\n")

var syncWordCount = [String: Int]()
let syncExecTime = measureTime {
    syncWordCount = getWordCountSync(words: randomWords)
}
print("\(syncWordCount)")
print("Sync: \(syncExecTime)")