         (doseq [f (remove #(.isDirectory %) 
                           (file-seq (io/file lang-dir)))]
           (train (slurp f) lang-key))))))
-  
-(defn- log-prob [token language]
-  (Math/log
-   (/ (double (inc (get (get @cat-map language {}) token 0)))
-      (+ (get @global-map token 0) (reduce + (vals @global-map))))))
 
-;; TODO probability of category
+;; TODO make abstractions like dictsize, tokencat instead of get-in..
+(defn unigram-probability [token language]
+  (let [a (get-in @model [:lang language :n1 token] 0)
+        b (get-in @model [:n1 token] 0)] ;; to aviod div/0
+    (/ (double (inc a))
+       (+ b (get-in @model [:n1 :size]))))) ;; dict size
+          
+
+(defn- probability [token language]
+  (let [trigram (get-in @model [:lang language :n3 token] 0)]
+    (cond (> trigram 0) 
+          (/ (double trigram) (get-in @model [:n3 token]))
+          :else ;; backoff strategy
+          (let [bigram (get-in @model [:lang language :n2 (butlast token)] 0)]
+            (cond (> bigram 0)
+                  (* (/ (double bigram) (get-in @model [:n2 (butlast token)]))
+                     (unigram-probability (last token) language))
+                  :else 
+                  (reduce *' (map #(unigram-probability % language) token)))))))
+
+(defn- lang-probability [lang]
+  (let [total (get-in @model [:n1 :size])]
+    (if (zero? total) 
+      (throw (IllegalArgumentException. "Model is empty"))
+      (/ (double (get-in @model [:lang lang :n1 :size] 0))
+         total))))
+
 (defn- score [tokens language]
-  (reduce +' (map #(log-prob % language) tokens))) 
+  (reduce +'
+          ;;(Math/log (lang-probability language))
+          (map #(Math/log (probability % language)) tokens)))
 
 (defn- scores [tokens languages]
   (map #(vec [% (score tokens %)]) languages))

  (let [tokens (t/tokenize source)]
    (do (fill-ngrams :n3 language (trigrams (wrap tokens)))
        (fill-ngrams :n2 language (bigrams (wrap tokens)))
        (fill-ngrams :n1 language tokens))))

(defn train-dir [dir]
  (let [langs (->> dir


(defn classify [source]
  (-> source
      (t/tokenize) ;; at least 1
      (wrap)
      (trigrams)
      (scores (languages))
      (best)))

(defn languages []
  (keys (get @model :lang)))
