typealias TaggedToken = (String, String?) // Can’t add tuples to an array without typealias. Compiler bug... Sigh.

func tag(text: String, scheme: String) -> [TaggedToken] {
    let options: NSLinguisticTaggerOptions = .OmitWhitespace | .OmitPunctuation | .OmitOther
    let tagger = NSLinguisticTagger(tagSchemes: NSLinguisticTagger.availableTagSchemesForLanguage("en"),
        options: Int(options.rawValue))
    tagger.string = text

    var tokens: [TaggedToken] = []

    // Using NSLinguisticTagger
    tagger.enumerateTagsInRange(NSMakeRange(0, count(text)), scheme:scheme, options: options) { tag, tokenRange, _, _ in
        let token = (text as NSString).substringWithRange(tokenRange)
        tokens.append((token, tag))
    }
    return tokens
}

func partOfSpeech(text: String) -> [TaggedToken] {
    return tag(text, NSLinguisticTagSchemeLexicalClass)
}

// returns I as a pronoun, went as a verb, to as a preposition, etc
partOfSpeech("I went to the store")

func lemmatize(text: String) -> [TaggedToken] {
    return tag(text, NSLinguisticTagSchemeLemma)
}

// turns "went" into "go"
lemmatize("I went to the store")

func language(text: String) -> [TaggedToken] {
    return tag(text, NSLinguisticTagSchemeLanguage)
}

language("Ik ben Ayaka") // tag is nl for Nederlands
language("Ich bin Ayaka") // tag is de for Deutsch
language("私の名前は彩花です") // tag is ja for Japanese

// MARK: - Training

    public func trainWithText(text: String, category: Category) {
        trainWithTokens(tokenizer(text), category: category)
    }

    public func trainWithTokens(tokens: [String], category: Category) {
        let tokens = Set(tokens)
        for token in tokens {
            incrementToken(token, category: category)
        }
        incrementCategory(category)
        trainingCount++
    }

    // MARK: - Classifying

    public func classifyText(text: String) -> Category? {
        return classifyTokens(tokenizer(text))
    }

    public func classifyTokens(tokens: [String]) -> Category? {
        // Compute argmax_cat [log(P(C=cat)) + sum_token(log(P(W=token|C=cat)))]
        var maxCategory: Category?
        var maxCategoryScore = -Double.infinity
        for (category, _) in categoryOccurrences {
            let pCategory = P(category)
            let score = tokens.reduce(log(pCategory)) { (total, token) in
                // P(W=token|C=cat) = P(C=cat, W=token) / P(C=cat)
                total + log((P(category, token) + smoothingParameter) / (pCategory + smoothingParameter * Double(tokenCount)))
            }
            if score > maxCategoryScore {
                maxCategory = category
                maxCategoryScore = score
            }
        }
        return maxCategory
    }

nbc.trainWithText("spammy spam spam", category: "spam")
nbc.trainWithText("spam has a lot of sodium and cholesterol", category: "spam")

nbc.trainWithText("nom nom ham", category: "ham")
nbc.trainWithText("please put the ham and eggs in the fridge", category: "ham")

let nbc = NaiveBayesClassifier { (text: String) -> [String] in
    return text.componentsSeparatedByString(" ")
}

nbc.classifyText("sodium and cholesterol") // "spam"
nbc.classifyText("spam and eggs") // "spam"
nbc.classifyText("do you like spam?") // "spam"

nbc.classifyText("use the eggs in the fridge") // "ham"
nbc.classifyText("ham and eggs") // "ham"
nbc.classifyText("do you like ham?") // "spam", INCORRECT
nbc.classifyText("do you eat egg?") // "spam", INCORRECT

// partOfSpeech
let nbc = NaiveBayesClassifier { (text: String) -> [String] in
    return partOfSpeech(text).map { (token, _) in
        return token
    }
}

nbc.classifyText("do you like ham?") // changed to "ham"
nbc.classifyText("do you eat egg?") // "spam"

// lemmatize
let nbc = NaiveBayesClassifier { (text: String) -> [String] in
    return lemmatize(text).map { (token, tag) in
        return tag ?? token
    }
}

nbc.classifyText("do you eat egg?") // now is "ham"!