diff --git a/07_final_assignment/baboonSimulation.R b/07_final_assignment/baboonSimulation.R
index e6bfa45..9959aee 100644
--- a/07_final_assignment/baboonSimulation.R
+++ b/07_final_assignment/baboonSimulation.R
@@ -83,10 +83,10 @@
 
 PresentTrials <- function(trialCount, present, data){
     trials <- data.frame()
-    for ( curTrial in 1:trialCount ) {
-        curBlock <- ( (curTrial - 1) %/% BLOCK_SIZE) + 1
-        curTrialInBlock <- ( (curTrial - 1) %% BLOCK_SIZE ) + 1
-        isNewBlock <- curTrialInBlock == 1
+    for ( curTrialNum in 1:trialCount ) {
+        curBlock <- ( (curTrialNum - 1) %/% BLOCK_SIZE) + 1
+        curTrialNumInBlock <- ( (curTrialNum - 1) %% BLOCK_SIZE ) + 1
+        isNewBlock <- curTrialNumInBlock == 1
 
         if ( isNewBlock ){
             isNotFirstBlock <- curBlock > 1
@@ -102,16 +102,16 @@
             block <- CreateBlock(data)
             block$Block <- curBlock
             block$TrialInBlock <- 1:nrow(block)
-            # TODO: Add words to learned
         }
 
-        curStim <- block$Stimulus[curTrialInBlock]
-        curStimIsWord <- present(curStim)
+        curTrial <- block[curTrialNumInBlock,]
+        curStim <- curTrial$Stimulus
+        curStimIsWord <- present(curStim, ifelse(curTrial$StimulusType == "Nonword", "Nonword", "Word"))
 
-        block$Trial[curTrialInBlock] <- curTrial
-        block$Response[curTrialInBlock] <- curStimIsWord
+        block$Trial[curTrialNumInBlock] <- curTrialNum
+        block$Response[curTrialNumInBlock] <- curStimIsWord
 
-        trials <- rbind(trials, block[curTrialInBlock,])
+        trials <- rbind(trials, block[curTrialNumInBlock,])
     }
     rownames(trials) <- 1:trialCount
     result <- list(Trials=trials,Stimuli=data)
@@ -127,8 +127,8 @@
     count <- BLOCK_SIZE * 2.5
 
     result <- PresentTrials(count,
-                            function(word){
-                                "Word"
+                            function(cue, outcome){
+                                return(outcome)
                             },
                             data)
     trials <- result$Trials
@@ -137,16 +137,44 @@
     expect_equal(length(stimuli$StimulusType[stimuli$StimulusType == "LearnedWord"]), 102)
 })
 
+# The original colSums fails if matrix filters to 1 row and gets implicit castet to a vector
+ColSums <- function(x) { 
+    if(is.matrix(x)){
+       return(colSums(x))
+    } else {
+        return(x)
+    }
+}
+
+test_that("ColSums",{
+    m2x2 <- matrix(1,2,2)
+    colnames(m2x2) <- c("a","b")
+
+    expect_equal(ColSums(m2x2), c(a=2,b=2))
+    expect_equal(ColSums(m2x2[1,]), c(a=1,b=1))
+})
+
 MakeMonkey <- function(cueset, outcomeset=c("Word", "Nonword"),
                        alpha=sqrt(0.001), beta=sqrt(0.001), lambda=1.0){
     weights <- matrix(0, length(cueset), length(outcomeset))
     rownames(weights) <- cueset
     colnames(weights) <- outcomeset
 
+    rate <- function(cue){
+        cues <- unlist(strsplit(cue, "_"))
+        cueWeights <- weights[cues, outcomeset]
+        totalActivation <- ColSums(cueWeights)
+
+        maxActivated <- names(totalActivation[totalActivation == max(totalActivation)])
+        return(RandomPick(maxActivated,1))
+    }
+    
     learner <- function(stim, resp){
         cues <- unlist(strsplit(stim, "_"))
 
-        totalActivation <- colSums(weights[cues, outcomeset])
+        cueWeights <- weights[cues, outcomeset]
+        totalActivation <- ColSums(cueWeights)
+       
         for (j in 1:length(cues)) {
             if (resp == "Nonword") {
                 yesType <- "Nonword"
@@ -162,8 +190,7 @@
             weights[cues[j],noType] <<-  weights[cues[j],noType] +
                 alpha * beta * (0 - totalActivation[noType])
         }
-        print(max(totalActivation))
-        # TODO: RETURN GUESS!!!
+        return(rate(stim))
     }
 
     give_weights <- function(){
@@ -172,7 +199,8 @@
 
     list(
         Weights=give_weights,
-        Learner=learner
+        Learner=learner,
+        Rate=rate
     )
 }
 
@@ -189,8 +217,7 @@
     monkey$Learner("a_b","Word")
     monkey$Learner("b_c","Nonword")
 
-    print(monkey$Weights())
-    print(str(monkey))
-    # TODO: MAKE TESTCASES!!
-})
+    expect_equal(monkey$Rate("a"), "Word")
+    expect_equal(monkey$Rate("c"), "Nonword")
+ })
 # EOF