diff --git a/LICENSE b/LICENSE
new file mode 100644
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,30 @@
+Copyright Author name here (c) 2018
+
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+    * Redistributions of source code must retain the above copyright
+      notice, this list of conditions and the following disclaimer.
+
+    * Redistributions in binary form must reproduce the above
+      copyright notice, this list of conditions and the following
+      disclaimer in the documentation and/or other materials provided
+      with the distribution.
+
+    * Neither the name of Author name here nor the names of other
+      contributors may be used to endorse or promote products derived
+      from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
diff --git a/README.md b/README.md
new file mode 100644
--- /dev/null
+++ b/README.md
@@ -0,0 +1,33 @@
+# haskmorph
+
+## Compiling
+
+First install [Stack](https://docs.haskellstack.org) somewhere on your `PATH`.  For example, for `~/.local/bin`:
+
+```
+wget https://get.haskellstack.org/stable/linux-x86_64.tar.gz -O -|tar xpfz - -C /tmp
+cp /tmp/stack-*/stack ~/.local/bin
+rm -rf /tmp/stack-*
+```
+
+Then, while in the directory of this README file, run:
+
+```
+stack build
+```
+
+The first time this runs will take a while, 10 or 15 minutes, as it builds an entire Haskell environment from scratch.  Subsequent compilations are very fast.
+
+## Running
+
+Invoke the program using Stack.  To see available sub-commands, run:
+
+```
+stack exec -- haskmorph -h
+```
+
+To see detailed help, run e.g.:
+
+```
+stack exec -- haskmorph train -h
+```
diff --git a/Setup.hs b/Setup.hs
new file mode 100644
--- /dev/null
+++ b/Setup.hs
@@ -0,0 +1,2 @@
+import Distribution.Simple
+main = defaultMain
diff --git a/app/Main.hs b/app/Main.hs
new file mode 100644
--- /dev/null
+++ b/app/Main.hs
@@ -0,0 +1,155 @@
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+
+module Main where
+
+import Prelude hiding (lookup)
+import Options.Generic (Generic, ParseRecord, Unwrapped, Wrapped, unwrapRecord, (:::), type (<?>)(..))
+import Control.Monad (join, liftM, foldM)
+import System.Random (getStdGen, mkStdGen)
+import Data.List (unfoldr, nub, mapAccumL, intercalate, sort, foldl1')
+import Data.Maybe (fromMaybe, catMaybes)
+import Data.Map (Map)
+import qualified Data.Map as Map
+import Data.Set (Set)
+import qualified Data.Set as Set
+import Control.Monad.IO.Class (liftIO)
+import Text.Printf (printf, PrintfArg(..), fmtPrecision, fmtChar, errorBadFormat, formatString, vFmt, IsChar)
+import Control.Monad.Loops
+import Control.Monad.Log
+import Control.Monad.State.Class (MonadState(get, put))
+import Control.Monad.Reader.Class
+import Control.Monad.Reader (ReaderT)
+import Control.Monad.IO.Class (MonadIO(liftIO))
+import Control.Monad.Reader
+import Control.Monad.State.Strict
+import Control.Monad.Random
+import Data.Vector (Vector)
+import qualified Data.Vector as Vector
+import Text.HaskSeg.Probability (Prob, LogProb, Probability(..), Categorical(..))
+import Text.HaskSeg.Types (Locations, Morph, Counts, Site, Location(..), Lookup, showLookup, showCounts, SamplingState(..), Params(..), Model(..))
+import Text.HaskSeg.Metrics (f1)
+import Text.HaskSeg.Utils (readState, readDataset, writeDataset, writeState, datasetToVocabulary, applySegmentation)
+import Text.HaskSeg.Location (randomFlip, createData, randomizeLocations, updateLocations, nonConflicting, wordsToSites, siteToWords, updateLocations', formatWord, showLexicon, initReverseLookup)
+import Text.HaskSeg.Lookup (cleanLookup, initializeLookups, computeUpdates)
+import Text.HaskSeg.Counts (cleanCounts, initializeCounts, updateCounts, addCounts, subtractCounts)
+import Text.HaskSeg.Logging (showFullState)
+import Text.HaskSeg.Model (combinations, oneWordProb, g, distribution, sampleSite, sample, applyModel, fromState)
+
+type ProbRep = LogProb
+
+--
+--   Command-line parsing
+--
+
+logLevels :: Map String Severity
+logLevels = Map.fromList [ ("debug", Debug)
+                         , ("info", Informational)
+                         , ("warn", Warning)
+                         , ("error", Error)
+                         ]
+
+data Parameters w = Train { inputFile :: w ::: Maybe String <?> "Input data file"
+                          , lineCount :: w ::: Maybe Int <?> "Number of lines to read (default: all)"
+                          , stateFile :: w ::: Maybe String <?> "Sampling state file"
+                          , iterations :: w ::: Int <?> "Number of sampling iterations"
+                          , alphaParam :: w ::: Maybe Double <?> "Per-decision concentration parameter (default: 0.1)"
+                          , sharpParam :: w ::: Maybe Double <?> "Probability to stop generating characters when drawing an unseen word (default: 0.5)"
+                          , etaParam :: w ::: Maybe Double <?> "Initial probability of each site being a boundary (default: 1.0)"
+                          , useSpaces :: w ::: Bool <?> "Make whitespace characters static borders (default: false)"
+                          , typeBased :: w ::: Bool <?> "Run over word types, rather than tokens (default: false)"
+                          , logLevel :: w ::: Maybe String <?> "Minimum log severity to display, one of [debug, info, warn, error] (default: info)"
+                          , randomSeed :: w ::: Maybe Int <?> "Set a deterministic random seed (default: use system RNG)"
+                          , minCount :: w ::: Maybe Int <?> "Only consider words with the given minimum frequency"                          
+                          }
+                  | Segment { inputFile :: w ::: Maybe String <?> "Input data file"
+                            , lineCount :: w ::: Maybe Int <?> "Number of lines to read (default: all)"
+                            , stateFile :: w ::: Maybe String <?> "Sampling state file"
+                            , segmentationFile :: w ::: Maybe String <?> "Output file for segmented text"
+                            , logLevel :: w ::: Maybe String <?> "Minimum log severity to display, one of [debug, info, warn, error] (default: info)"
+                            , randomSeed :: w ::: Maybe Int <?> "Set a deterministic random seed (default: use system RNG)"
+                            }
+  deriving (Generic)                              
+
+instance ParseRecord (Parameters Wrapped)
+deriving instance Show (Parameters Unwrapped)
+
+instance (MonadLog (WithSeverity String) m) => MonadLog (WithSeverity String) (RandT g m)
+
+-- --
+-- --   Top-level actions
+-- --
+-- -- | Train a model on given data
+trainModel :: (Categorical p, Show p, Probability p, MonadIO m, MonadLog (WithSeverity String) m) => Vector Char -> Double -> (Params p) -> Int -> StdGen -> m (SamplingState Char)
+trainModel seq eta params iterations gen = do
+  logInfo (printf "Initial random seed: %v" (show gen))
+  (locations, gold) <- createData params seq
+
+  let numChars = (length . nub . map (\x -> _value x) . Vector.toList) locations
+      charProb = fromDouble $ 1.0 / (fromIntegral numChars)
+      (locations', gen') = randomizeLocations eta locations gen
+      counts = initializeCounts locations'
+      (lookupS, lookupE) = initializeLookups locations'
+  --logInfo (show (lookupS, lookupE))
+  let rLookup = initReverseLookup lookupS lookupE
+  --logInfo (show rLookup)
+  let state = SamplingState counts locations' lookupS lookupE rLookup Set.empty
+      params' = params { _gold=gold, _charProb=charProb }
+  runReaderT (execStateT (evalRandT (forM_ [1..iterations] sample) gen') state) params'
+
+
+--
+--   Main entrypoint
+--
+
+main :: IO ()
+main = do
+  args <- unwrapRecord "Type-based sampling for segmentation model with Dirichlet process prior on words"
+  gen <- case randomSeed args of Nothing -> getStdGen
+                                 Just i -> return $ mkStdGen i
+  let level = logLevels Map.! (fromMaybe "info" (logLevel args))
+
+  runLoggingT (case args of
+                  Train{..} -> do
+                    ds <- liftIO $ readDataset inputFile lineCount
+                    let vocab = datasetToVocabulary ds
+                        seq = (Vector.fromList . intercalate " ") (Set.toList vocab)
+                    --seq <- liftIO $ (liftM (Vector.fromList . intercalate " " . (case lineCount of Nothing -> id; Just lc -> take lc) . lines) . readFile) inputFile
+                    let numChars = (length . nub . Vector.toList) seq
+                        charProb = fromDouble $ 1.0 / (fromIntegral numChars) :: ProbRep
+                        params = Params (fromDouble $ fromMaybe 0.1 alphaParam) (fromDouble $ fromMaybe 0.5 sharpParam) (fromDouble $ 1.0 - (fromMaybe 0.5 sharpParam)) useSpaces typeBased Set.empty charProb (fromMaybe 1 minCount)
+                    state <- trainModel seq (fromMaybe 1.0 etaParam) params iterations gen
+                    liftIO $ writeState stateFile params (_locations state) 
+                  Segment{..} -> do
+                    dataSet <- liftIO $ readDataset inputFile lineCount
+                    (params :: Params ProbRep, locs :: Locations Char) <- liftIO $ readState stateFile
+                    let cs = nub $ concat (map concat dataSet)
+                    model <- fromState (params, locs) (Just cs)
+                    dataSet' <- applyModel model dataSet
+                    liftIO $ writeDataset segmentationFile dataSet'
+
+                    --let vocab = datasetToVocabulary ds
+                    --seq = (Vector.fromList . intercalate " ") (Set.toList vocab)
+                    --liftIO $ print seq
+              )
+                    (\msg -> case msgSeverity msg <= level of
+                               True -> putStrLn (discardSeverity msg)
+                               False -> return ()
+                      )
+                    
+                    --(params :: (Params ProbRep), modelLocations :: Locations Char) <- liftIO (readState modelFile)
+                    --model <- createModel params modelLocations                    
+                    --let seg = applyModel params modelLocations vocab
+                    --liftIO $ print seg
+                    --liftIO $ writeFile segmentationFile (show seg)
+                    --ds' = applySegmentation seg ds
+                    --liftIO $ writeDataset segmentationFile ds'
diff --git a/haskseg.cabal b/haskseg.cabal
new file mode 100644
--- /dev/null
+++ b/haskseg.cabal
@@ -0,0 +1,84 @@
+-- This file has been generated from package.yaml by hpack version 0.28.2.
+--
+-- see: https://github.com/sol/hpack
+--
+-- hash: 81db90b4fd7cad4b8cd8b9bbdb1ea0ad0e959abc9287bc434a66289a2925b2da
+
+name:           haskseg
+version:        0.1.0.0
+synopsis:       Simple unsupervised segmentation model
+description:    Implementation of the non-parametric segmentation model described in "Type-based MCMC" (Liang, Jordan, and Klein, 2010).
+category:       NLP
+homepage:       https://github.com/githubuser/haskseg#readme
+author:         Tom Lippincott
+maintainer:     tom@cs.jhu.edu
+copyright:      2018 Tom Lippincott
+license:        BSD3
+license-file:   LICENSE
+build-type:     Simple
+cabal-version:  >= 1.10
+extra-source-files:
+    README.md
+
+library
+  exposed-modules:
+      Text.HaskSeg.Counts
+      Text.HaskSeg.DataSet
+      Text.HaskSeg.Location
+      Text.HaskSeg.Logging
+      Text.HaskSeg.Lookup
+      Text.HaskSeg.Metrics
+      Text.HaskSeg.Model
+      Text.HaskSeg.Probability
+      Text.HaskSeg.Types
+      Text.HaskSeg.Utils
+  other-modules:
+      Paths_haskseg
+  hs-source-dirs:
+      src
+  default-extensions: Strict StrictData FlexibleContexts RecordWildCards MultiParamTypeClasses FlexibleInstances OverloadedStrings ScopedTypeVariables
+  build-depends:
+      MonadRandom >=0.5.1.1
+    , ansi-terminal >=0.8.0.4
+    , array
+    , base >=4.7 && <5
+    , bytestring >=0.10.8.1
+    , containers >=0.5.10.2
+    , exact-combinatorics >=0.2.0.8
+    , logging-effect >=1.3.2
+    , monad-loops >=0.4.3
+    , mtl >=2.2.2
+    , optparse-generic >=1.2.2
+    , random >=1.1
+    , random-shuffle >=0.0.4
+    , text >=1.2.2
+    , vector >=0.12.0.1
+    , zlib >=0.6.1
+  default-language: Haskell2010
+
+executable haskseg
+  main-is: Main.hs
+  other-modules:
+      Paths_haskseg
+  hs-source-dirs:
+      app
+  default-extensions: Strict StrictData FlexibleContexts RecordWildCards MultiParamTypeClasses FlexibleInstances OverloadedStrings ScopedTypeVariables
+  build-depends:
+      MonadRandom >=0.5.1.1
+    , ansi-terminal >=0.8.0.4
+    , array
+    , base >=4.7 && <5
+    , bytestring >=0.10.8.1
+    , containers >=0.5.10.2
+    , exact-combinatorics >=0.2.0.8
+    , haskseg
+    , logging-effect >=1.3.2
+    , monad-loops >=0.4.3
+    , mtl >=2.2.2
+    , optparse-generic >=1.2.2
+    , random >=1.1
+    , random-shuffle >=0.0.4
+    , text >=1.2.2
+    , vector >=0.12.0.1
+    , zlib >=0.6.1
+  default-language: Haskell2010
diff --git a/src/Text/HaskSeg/Counts.hs b/src/Text/HaskSeg/Counts.hs
new file mode 100644
--- /dev/null
+++ b/src/Text/HaskSeg/Counts.hs
@@ -0,0 +1,56 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+
+module Text.HaskSeg.Counts (cleanCounts, initializeCounts, updateCounts, addCounts, subtractCounts) where
+
+import Control.Monad.Random
+import Data.Set (Set)
+import qualified Data.Set as Set
+import Data.Map (Map)
+import qualified Data.Map as Map
+import Data.Vector (Vector)
+import qualified Data.Vector as Vector
+import Text.Printf (printf, PrintfArg(..), fmtPrecision, fmtChar, errorBadFormat, formatString, vFmt, IsChar)
+import Control.Monad.Log
+import Control.Monad.State.Class (MonadState(get, put))
+import Control.Monad.Reader.Class
+import Control.Monad.Reader (ReaderT)
+import Control.Monad.IO.Class (MonadIO(liftIO))
+import Control.Monad.State.Strict
+import Data.Tuple (swap)
+import Data.List (unfoldr, nub, mapAccumL, intercalate, sort, foldl1')
+
+import Text.HaskSeg.Probability (Prob, LogProb, Probability(..), showDist, sampleCategorical)
+import Text.HaskSeg.Types (Locations, Morph, Counts, Site, Location(..), Lookup, showLookup, showCounts, SamplingState(..), Params(..))
+
+
+-- | Remove zero-elements from a count object
+cleanCounts :: Counts elem -> Counts elem
+cleanCounts = Map.filter (\x -> x /= 0)
+
+-- | Initialize word counts from scratch, given boundary assignments
+initializeCounts :: (Ord elem, Show elem) => Locations elem -> Counts elem
+initializeCounts ls = Map.fromListWith (+) (Vector.toList (Vector.map (\x -> (x, 1)) words'))
+  where
+    words = Vector.unfoldr (\xs -> case span (\x -> _morphFinal x == False) xs of
+                               ([], []) -> Nothing
+                               (xs', x:ys) -> Just (xs' ++ [x], ys)
+                           ) (Vector.toList ls)
+    words' = Vector.map (Vector.fromList . map _value) words
+
+-- | Use provided function to update counts for a word
+updateCounts :: (Ord elem) => (Int -> Int -> Int) -> Morph elem -> Int -> Counts elem -> Counts elem
+updateCounts f w n = Map.insertWith f w n
+
+-- | Convenience function for adding counts
+addCounts :: (Ord elem) => Morph elem -> Int -> Counts elem -> Counts elem
+addCounts = updateCounts (+)
+
+-- | Convenience function for subtracting counts
+subtractCounts :: (Ord elem) => Morph elem -> Int -> Counts elem -> Counts elem
+subtractCounts = updateCounts (flip (-))
diff --git a/src/Text/HaskSeg/DataSet.hs b/src/Text/HaskSeg/DataSet.hs
new file mode 100644
--- /dev/null
+++ b/src/Text/HaskSeg/DataSet.hs
@@ -0,0 +1,9 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+
+module Text.HaskSeg.DataSet () where
diff --git a/src/Text/HaskSeg/Location.hs b/src/Text/HaskSeg/Location.hs
new file mode 100644
--- /dev/null
+++ b/src/Text/HaskSeg/Location.hs
@@ -0,0 +1,203 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+
+module Text.HaskSeg.Location (randomFlip, createData, randomizeLocations, updateLocations, updateLocations', nonConflicting, wordsToSites, siteToWords, formatWord, showLexicon, initReverseLookup) where
+
+import Control.Monad.Random
+import Data.Set (Set)
+import qualified Data.Set as Set
+import qualified Data.Maybe as Maybe
+import Data.Map (Map)
+import qualified Data.Map as Map
+import Data.Vector (Vector)
+import qualified Data.Vector as Vector
+import Text.Printf (printf, PrintfArg(..), fmtPrecision, fmtChar, errorBadFormat, formatString, vFmt, IsChar)
+import Control.Monad.Log
+import Control.Monad.State.Class (MonadState(get, put))
+import Control.Monad.Reader.Class
+import Control.Monad.Reader (ReaderT)
+import Control.Monad.IO.Class (MonadIO(liftIO))
+import Control.Monad.State.Strict
+import Data.Tuple (swap)
+import Data.List (unfoldr, nub, mapAccumL, intercalate, sort, foldl1')
+import Text.HaskSeg.Probability (Prob, LogProb, Probability(..), showDist, sampleCategorical)
+import Text.HaskSeg.Types (Locations, Morph, Counts, Site, Location(..), Lookup, showLookup, showCounts, SamplingState(..), Params(..))
+import Debug.Trace (traceShowId)
+
+randomFlip p g = (v < p, g')
+  where
+    (v, g') = randomR (0.0, 1.0) g
+
+
+createData :: (Probability p, MonadLog (WithSeverity String) m) => (Params p) -> Vector Char -> m (Locations Char, Set Int)
+createData Params{..} cs' = do
+  let cs = Vector.toList cs'
+      ls = lines cs
+      wss = concat $ map words ls
+      wc = Map.fromListWith (\a b -> a + b) (zip wss $ repeat 1)
+      keep = Map.filter (>= _minCount) wc
+      ws = if _types == True then Map.keys keep else concat $ map words ls
+      bs = map length ws
+      bs' = (reverse . drop 1 . reverse . drop 1) $ scanl (+) (-1) bs
+      ws' = if _spaces == True then ws else [concat ws]
+      ws'' = Vector.concat [sequenceToLocations w | w <- ws']
+  logInfo (printf "Loaded data set of %d characters/%d words" (length cs) (length ws))
+  return $! (ws'', Set.fromList bs')
+
+
+formatWord :: [Location Char] -> String
+formatWord ls = printf "%s : %s" word (intercalate "@@" (reverse (morphs ls [])))
+  where
+    word = map _value ls
+    morphs [] acc = acc
+    morphs ls' acc = morphs rem (morph:acc)
+      where
+        (pref, r:rem) = span (\x -> _morphFinal x == False) ls'
+        morph = map _value (pref ++ [r]) 
+
+
+showLexicon :: Locations Char -> [String]
+showLexicon ls = go [] (Vector.toList ls)
+  where
+    go acc [] = acc
+    go acc ls' = go (word:acc) rem
+      where
+        (pref, r:rem) = span (\x -> _static x == False) ls'
+        word = formatWord (pref ++ [r]) 
+
+
+-- | Switch each potential morpheme boundary (i.e. intra-word indices) to True or False
+randomizeLocations :: Double -> Locations elem -> StdGen -> (Locations elem, StdGen)
+randomizeLocations p xs g = (Vector.fromList xs', g')
+  where
+    (g', bs) = mapAccumL (\g'' Location{..} -> if _static == True then (g'', True) else swap (randomFlip p g'' :: (Bool, StdGen))) g (Vector.toList xs)
+    xs' = [x { _morphFinal=b } | (x, b) <- zip (Vector.toList xs) bs]
+
+
+updateLocations' :: elem -> Locations elem -> Set Int -> Set Int -> Locations elem
+updateLocations' a ls pos neg = Vector.update ls updates
+  where
+    p = Location a True False
+    n = Location a False False
+    pos' = (Vector.map (\i -> (i, p)) . Vector.fromList . Set.toList) pos
+    neg' = (Vector.map (\i -> (i, n)) . Vector.fromList . Set.toList) neg
+    updates = pos' Vector.++ neg'
+
+
+updateLocations :: (MonadState (SamplingState elem) m) => elem -> Set Int -> Set Int -> m ()
+updateLocations a pos neg = do
+  --Vector.update ls updates
+  let p = Location a True False
+      n = Location a False False
+      pos' = (Vector.map (\i -> (i, p)) . Vector.fromList . Set.toList) pos
+      neg' = (Vector.map (\i -> (i, n)) . Vector.fromList . Set.toList) neg
+      updates = pos' Vector.++ neg'
+  modify' (\state -> state)
+    
+
+-- | Turn a sequence of values into a sequence of locations
+sequenceToLocations :: [elem] -> Locations elem
+sequenceToLocations xs = Vector.fromList $ nonFinal ++ [final]
+  where
+    xs' = init xs
+    nonFinal = map (\x -> Location x False False) xs'
+    x = last xs
+    final = Location x True True
+
+
+-- -- | Find the two words implied by a boundary at the given site
+-- siteToWords :: (Show elem, MonadLog (WithSeverity String) m) => Locations elem -> Int -> m (Morph elem, Morph elem)
+-- siteToWords ls s = do
+--   let (before, after) = Vector.splitAt (s + 1) ls
+--       (bPref, bRem) = Vector.break _morphFinal after
+--       (b', before') = Vector.splitAt 1 (Vector.reverse before)
+--       (aPref, aRem) = Vector.break _morphFinal before'
+--       b = case Vector.length bRem of 0 -> bPref
+--                                      _ -> bPref Vector.++ (Vector.fromList [Vector.head bRem])
+--       (before'', after'') = (Vector.map _value (Vector.reverse (b' Vector.++ aPref)), Vector.map _value b)
+--   return $! (before'', after'')
+
+-- initReverseLookup :: Locations elem -> Map Int (Morph elem, Morph elem)
+-- initReverseLookup ls = Map.fromList ls'
+--   where
+--     items = Vector.map _value ls
+--     starts = Vector.toList $ Vector.findIndices _morphFinal ls
+--     ends = (drop 1 starts) ++ [Vector.length ls]
+--     spans = zip starts ends
+    
+--     dummy = Vector.fromList []
+--     ls' = map (\i -> (i, (dummy, dummy))) [0..Vector.length ls]
+
+initReverseLookup :: (Eq elem) => Lookup elem -> Lookup elem -> Map Int (Morph elem, Morph elem)
+initReverseLookup s e = Map.fromList [(i, (Maybe.fromJust a, Maybe.fromJust b)) | (i, (a, b)) <- atBoundaries ++ atNonBoundaries, a /= Nothing && b /= Nothing]
+  where
+    e' = Map.fromList $ concat [[(v', k) | v' <- Set.toList v] | (k, v) <- Map.toList e]
+    s' = Map.fromList $ concat [[(v', k) | v' <- Set.toList v] | (k, v) <- Map.toList s]
+    indices = Map.keys s'
+    atBoundaries = [(i, (e' Map.!? (i), s' Map.!? i)) | i <- indices]
+    atNonBoundaries = concat $ [[(i + i', (Just $ Vector.slice 0 i' m, Just $ Vector.slice i' (Vector.length m - i') m)) | i' <- [1..Vector.length m - 1]] | (i, m) <- map (\i -> (i, s' Map.! i)) (Map.keys s')]
+
+
+-- | Find the two words implied by a boundary at the given site
+siteToWords' :: (Show elem, MonadLog (WithSeverity String) m, MonadState (SamplingState elem) m) => Int -> m (Morph elem, Morph elem)
+siteToWords' s = do
+  SamplingState{..} <- get
+  let (a, b) = _wordsLookup Map.! s
+  --(a', b') <- siteToWords' s
+  --logInfo (show ((a, b), (a', b')))
+  
+  return (a, b)
+
+
+-- | Find the two words implied by a boundary at the given site
+siteToWords :: (Show elem, MonadLog (WithSeverity String) m, MonadState (SamplingState elem) m) => Int -> m (Morph elem, Morph elem)
+siteToWords s = do
+  SamplingState{..} <- get
+  let ls = _locations
+  let (before, after) = Vector.splitAt (s + 1) ls
+      (bPref, bRem) = Vector.break _morphFinal after
+      (b', before') = Vector.splitAt 1 (Vector.reverse before)
+      (aPref, aRem) = Vector.break _morphFinal before'
+      b = case Vector.length bRem of 0 -> bPref
+                                     _ -> bPref Vector.++ (Vector.fromList [Vector.head bRem])
+      (before'', after'') = (Vector.map _value (Vector.reverse (b' Vector.++ aPref)), Vector.map _value b)
+  return $! (before'', after'')
+
+
+-- | For sites with matching type, return a subset that don't conflict
+nonConflicting :: (MonadLog (WithSeverity String) m) => (Int, (Int, Int)) -> Set (Int, (Int, Int)) -> Set (Int, (Int, Int)) -> m (Set Int, Set Int)
+nonConflicting piv@(pivi, (si1, si2)) a b = return $! (a'', b'')
+  where
+    reducer (ms, vs) (i, (s1, s2)) = (ms', vs')
+      where
+        affected = Set.fromList [s1..s2]
+        conflict = Set.size (ms `Set.intersection` affected) > 0
+        ms' = if conflict then ms else ms `Set.union` affected
+        vs' = if conflict then vs else i `Set.insert` vs
+    (mods, a') = Set.foldl' reducer (Set.fromList [si1..si2], Set.empty) a
+    (mods', b') = Set.foldl' reducer (mods, Set.empty) b
+    a'' = if piv `Set.member` a then pivi `Set.insert` a' else a'
+    b'' = if piv `Set.member` b then pivi `Set.insert` b' else b'
+
+  
+-- | For two words, return all compatible sites
+wordsToSites :: (Probability p, MonadState (SamplingState elem) m, MonadReader (Params p) m, MonadLog (WithSeverity String) m, Show elem, Ord elem, PrintfArg elem) => Int -> Lookup elem -> Lookup elem -> Morph elem -> Morph elem -> m (Set Int, Set Int)
+wordsToSites piv luS luE a b = do
+  let j = a Vector.++ b
+      jS = Vector.fromList $ map (\x -> x + (Vector.length a)) (Set.toList $ Map.findWithDefault Set.empty j luS)
+      aE = Map.findWithDefault Set.empty a luE
+      bS = Map.findWithDefault Set.empty b luS
+      splits' = Set.map (\i -> (i, (i - length a, i + length b))) $ Set.intersection aE bS
+      nonSplits' = Set.map (\i -> (i, (i - length a, i + length b))) $  (Set.fromList . Vector.toList) jS
+      piv' = (piv, (piv - length a, piv + length b))
+  (splits, nonSplits) <- nonConflicting piv' splits' nonSplits'
+  let nSplit = Set.size splits
+      nFull = Set.size nonSplits
+  --s <- showFullState Nothing Nothing
+  --if nSplit + nFull == 0 then (logDebug s) >> (logDebug $ show (luS, luE)) >> error "Found zero sites!" else return ()
+  return $! (nonSplits, splits)
diff --git a/src/Text/HaskSeg/Logging.hs b/src/Text/HaskSeg/Logging.hs
new file mode 100644
--- /dev/null
+++ b/src/Text/HaskSeg/Logging.hs
@@ -0,0 +1,93 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+
+module Text.HaskSeg.Logging (showFullState) where
+
+import Prelude hiding (lookup)
+import Options.Generic (Generic, ParseRecord, Unwrapped, Wrapped, unwrapRecord, (:::), type (<?>)(..))
+import Control.Monad (join, liftM, foldM)
+import System.IO (withFile, hPutStr, IOMode(..), readFile)
+import System.Random (getStdGen, mkStdGen)
+import Data.List (unfoldr, nub, mapAccumL, intercalate, sort, foldl1')
+import Data.Maybe (fromMaybe, catMaybes)
+import Data.Map (Map)
+import qualified Data.Map as Map
+import Data.Set (Set)
+import qualified Data.Set as Set
+import Control.Monad.IO.Class (liftIO)
+import Text.Printf (printf, PrintfArg(..), fmtPrecision, fmtChar, errorBadFormat, formatString, vFmt, IsChar)
+import Math.Combinatorics.Exact.Binomial (choose)
+import Control.Monad.Loops
+import Control.Monad.Log
+import Control.Monad.State.Class (MonadState(get, put))
+import Control.Monad.Reader.Class
+import Control.Monad.Reader (ReaderT)
+import Control.Monad.IO.Class (MonadIO(liftIO))
+import Data.Tuple (swap)
+import Control.Monad.Reader
+import Control.Monad.State.Strict
+import Control.Monad.Random
+import System.Random.Shuffle (shuffleM)
+import qualified Data.ByteString.Lazy as BS
+import qualified Data.ByteString.Char8 as BSC
+import Data.Vector (Vector)
+import qualified Data.Vector as Vector
+import qualified System.Console.ANSI as A
+import Text.HaskSeg.Probability (Prob, LogProb, Probability(..), showDist, sampleCategorical)
+import Text.HaskSeg.Types (Locations, Morph, Counts, Site, Location(..), Lookup, showLookup, showCounts, SamplingState(..), Params(..))
+import Text.HaskSeg.Metrics (f1)
+import Text.HaskSeg.Utils (readDataset, writeDataset) --, readVocabulary, writeVocabulary)
+import Text.HaskSeg.Location (randomFlip, createData, randomizeLocations, updateLocations, nonConflicting, wordsToSites, siteToWords, updateLocations')
+import Text.HaskSeg.Lookup (cleanLookup, initializeLookups, computeUpdates)
+import Text.HaskSeg.Counts (cleanCounts, initializeCounts, updateCounts, addCounts, subtractCounts)
+
+
+
+goldA = A.setSGRCode [A.SetColor A.Background A.Vivid A.Green]
+goldB = A.setSGRCode [A.SetColor A.Background A.Dull A.Green]
+goldAlts = [if i `mod` 2 == 0 then goldA else goldB | i <- [1..]]
+
+goldFormat = A.setSGRCode [A.SetColor A.Background A.Vivid A.Blue]
+staticFormat = A.setSGRCode [A.SetColor A.Background A.Vivid A.Yellow]
+sampleFormat = A.setSGRCode [A.SetColor A.Foreground A.Vivid A.Red]
+siteFormat = A.setSGRCode [A.SetUnderlining A.SingleUnderline]
+pivotFormat = A.setSGRCode [A.SetConsoleIntensity A.BoldIntensity, A.SetUnderlining A.SingleUnderline]
+
+sampleA = A.setSGRCode [A.SetColor A.Foreground A.Vivid A.Black]
+sampleB = A.setSGRCode [A.SetColor A.Foreground A.Vivid A.Red]
+sampleAlts = [if i `mod` 2 == 0 then sampleA else sampleB | i <- [1..]]
+
+reset = A.setSGRCode [A.Reset]
+
+
+showFullState :: (Probability p, IsChar elem, MonadState (SamplingState elem) m, MonadReader (Params p) m, PrintfArg elem) => Maybe Int -> Maybe (Set Int) -> m String
+showFullState mi ms = do
+  SamplingState{..} <- get
+  params@(Params{..}) <- ask  
+  let ls = (Vector.toList . Vector.indexed) _locations
+      renderChar ([], golds, samples) = Nothing
+      renderChar ((i, Location{..}):locs, golds, samples) = Just (formatting ++ (printf "%v" _value) ++ reset, (locs, golds', samples'))
+        where
+          g:gs = golds
+          s:ss = samples
+          isGold = i `Set.member` _gold
+          isSet = _morphFinal
+          isPivot = Just i == mi
+          isStatic = _static
+          isSite = i `Set.member` (fromMaybe Set.empty ms)
+          gf = if isGold then Just goldFormat else Nothing
+          sf = if isSet then Just sampleFormat else Nothing
+          pf = if isPivot then Just pivotFormat else Nothing          
+          ssf = if isSite then Just siteFormat else Nothing
+          stf = if isStatic then Just staticFormat else Nothing
+          formatting = (concat . catMaybes) [gf, sf, pf, ssf, stf]
+          golds' = if isGold then gs else g:gs
+          samples' = if isSet then ss else s:ss
+      toks = unfoldr renderChar (ls, goldAlts, sampleAlts)
+  --  return $ concat toks
+  return $! (intercalate "\n" [concat toks, printf "Starts: %v" (showLookup _startLookup), printf "Ends: %v" (showLookup _endLookup), printf "Counts: %v" (showCounts _counts)])
diff --git a/src/Text/HaskSeg/Lookup.hs b/src/Text/HaskSeg/Lookup.hs
new file mode 100644
--- /dev/null
+++ b/src/Text/HaskSeg/Lookup.hs
@@ -0,0 +1,56 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+
+module Text.HaskSeg.Lookup (cleanLookup, initializeLookups, computeUpdates) where
+
+import Control.Monad.Random
+import Data.Set (Set)
+import qualified Data.Set as Set
+import Data.Map (Map)
+import qualified Data.Map as Map
+import Data.Vector (Vector)
+import qualified Data.Vector as Vector
+import Text.Printf (printf, PrintfArg(..), fmtPrecision, fmtChar, errorBadFormat, formatString, vFmt, IsChar)
+import Control.Monad.Log
+import Control.Monad.State.Class (MonadState(get, put))
+import Control.Monad.Reader.Class
+import Control.Monad.Reader (ReaderT)
+import Control.Monad.IO.Class (MonadIO(liftIO))
+import Control.Monad.State.Strict
+import Data.Tuple (swap)
+import Data.List (unfoldr, nub, mapAccumL, intercalate, sort, foldl1')
+
+import Text.HaskSeg.Probability (Prob, LogProb, Probability(..), showDist, sampleCategorical)
+import Text.HaskSeg.Types (Locations, Morph, Counts, Site, Location(..), Lookup, showLookup, showCounts, SamplingState(..), Params(..))
+
+
+-- | Remove morphs with no associated locations
+cleanLookup :: Lookup elem -> Lookup elem
+cleanLookup = Map.filter (\x -> Set.size x /= 0)
+
+-- | Initialize word lookup from scratch, given sampling state
+initializeLookups :: (Ord a, Show a) => Locations a -> (Lookup a, Lookup a)
+initializeLookups ls = go ((Vector.toList . Vector.indexed) ls) Map.empty Map.empty []
+  where
+    go ((i, l):ls') mS mE w = case _morphFinal l of
+                                False -> go ls' mS mE w'
+                                True -> go ls' (Map.insertWith (Set.union) (Vector.fromList $ reverse w') (Set.singleton $ i - (length w) - 1) mS) (Map.insertWith (Set.union) (Vector.fromList $ reverse w') (Set.singleton $ i) mE) []
+      where
+        w' = _value l : w
+    go [] mS mE w = (mS, mE)
+
+-- | Compute the start and end lookup updates implied by setting the given sites to positive and negative, based on the two context-words
+computeUpdates :: (Ord elem, Show elem) => Set Int -> Set Int -> Morph elem -> Morph elem -> (Lookup elem, Lookup elem) -- , [(Int, (Morph elem, Morph elem))])
+computeUpdates pos neg a b = (sUp, eUp)
+  where
+    c = a Vector.++ b
+    aLocs = Set.map (\x -> (x - (Vector.length a), x)) pos
+    bLocs = Set.map (\x -> (x, x + (Vector.length b))) pos
+    cLocs = Set.map (\x -> (x - (Vector.length a), x + (Vector.length b))) neg
+    sUp = Map.fromListWith Set.union [(w, Set.map fst ls) | (w, ls) <- zip [a, b, c] [aLocs, bLocs, cLocs]]
+    eUp = Map.fromListWith Set.union [(w, Set.map snd ls) | (w, ls) <- zip [a, b, c] [aLocs, bLocs, cLocs]]
diff --git a/src/Text/HaskSeg/Metrics.hs b/src/Text/HaskSeg/Metrics.hs
new file mode 100644
--- /dev/null
+++ b/src/Text/HaskSeg/Metrics.hs
@@ -0,0 +1,43 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+
+module Text.HaskSeg.Metrics (precision, recall, f1) where
+
+import Data.Set (Set)
+import qualified Data.Set as Set
+
+precision :: Set Int -> Set Int -> Double
+precision guesses golds = tp / gs
+  where
+    tp = (fromIntegral . Set.size . Set.intersection guesses) golds
+    gs = (fromIntegral . Set.size) guesses
+
+recall :: Set Int -> Set Int -> Double
+recall guesses golds = tp / gs
+  where
+    tp = (fromIntegral . Set.size . Set.intersection guesses) golds
+    gs = (fromIntegral . Set.size) golds
+
+f1 :: Set Int -> Set Int -> Double
+f1 guesses golds = 2.0 * numer / denom
+  where
+    p = precision guesses golds
+    r = recall guesses golds
+    numer = p * r
+    denom = p + r
+
+-- evaluateBoundaries :: Locations a -> Maybe [Int] -> Double
+-- evaluateBoundaries guesses (Just golds) = (2.0 * precision * recall) / (precision + recall)
+--   where
+--     guesses' = Set.fromList $ Vector.toList $ Vector.findIndices (\x -> _morphFinal x) guesses
+--     golds' = Set.fromList golds
+--     trueCount = fromIntegral (Set.size golds')
+--     guessCount = fromIntegral (Set.size guesses')
+--     correctCount = fromIntegral (Set.size $ guesses' `Set.intersection` golds')
+--     precision = correctCount / guessCount
+--     recall = correctCount / trueCount
diff --git a/src/Text/HaskSeg/Model.hs b/src/Text/HaskSeg/Model.hs
new file mode 100644
--- /dev/null
+++ b/src/Text/HaskSeg/Model.hs
@@ -0,0 +1,299 @@
+module Text.HaskSeg.Model (applyModel, combinations, oneWordProb, g, distribution, sampleSite, sample, fromState) where
+
+import Data.List (unfoldr, nub, mapAccumL, intercalate, sort, foldl1', sortOn, maximumBy)
+import Data.Maybe (fromMaybe, catMaybes)
+import Data.Ord (comparing)
+import Data.Map (Map)
+import qualified Data.Map as Map
+import Data.Set (Set)
+import qualified Data.Set as Set
+import Text.Printf (printf, PrintfArg(..), fmtPrecision, fmtChar, errorBadFormat, formatString, vFmt, IsChar)
+import Math.Combinatorics.Exact.Binomial (choose)
+import Control.Monad.Loops
+import Control.Monad.Log
+import Control.Monad.State.Class (MonadState(get, put))
+import Control.Monad.Reader.Class
+import Control.Monad.Reader (ReaderT)
+import Control.Monad.Reader
+import Control.Monad.State.Strict
+import Control.Monad.Random
+import System.Random.Shuffle (shuffleM)
+import Data.Vector (Vector)
+import qualified Data.Vector as Vector
+import Text.HaskSeg.Probability (Prob, LogProb, Probability(..), showDist, sampleCategorical, Categorical)
+import Text.HaskSeg.Types (Locations, Morph, Counts, Site, Location(..), Lookup, showLookup, showCounts, SamplingState(..), Params(..), Vocabulary, Segmentation, Dataset, ReverseLookup)
+import Text.HaskSeg.Metrics (f1)
+import Text.HaskSeg.Location (randomFlip, createData, randomizeLocations, updateLocations, nonConflicting, wordsToSites, siteToWords, updateLocations', initReverseLookup)
+import Text.HaskSeg.Lookup (cleanLookup, initializeLookups, computeUpdates)
+import Text.HaskSeg.Counts (cleanCounts, initializeCounts, updateCounts, addCounts, subtractCounts)
+import Text.HaskSeg.Probability (Prob, LogProb, Probability(..), showDist, sampleCategorical)
+import Debug.Trace (traceShowId)
+import Control.Monad.ST
+import Data.STRef
+import Control.Monad
+import Data.Array.ST
+
+
+type Model p elem = Map (Vector elem) p
+
+
+fromState :: (MonadLog (WithSeverity String) m, Ord elem, Show elem, Probability p) => (Params p, Locations elem) -> Maybe [elem] -> m (Model p elem)
+fromState (p, ls) cs = do
+  let cts = initializeCounts ls
+      ups = case cs of Nothing -> Map.fromList []
+                       Just cs' -> Map.fromList [(Vector.fromList [c], 1) | c <- cs']
+      cts' = Map.unionWith (\a b -> a) cts ups
+      ps = map (\w -> (w, oneTimeOneWord cts p w)) (Map.keys cts')      
+  return $ Map.fromList ps
+
+    
+likelihood :: (Probability p, Categorical p, Show p, MonadIO m, MonadRandom m, (MonadReader (Params p)) m, MonadState (SamplingState Char) m, MonadLog (WithSeverity String) m) => m p
+likelihood = do
+  SamplingState{..} <- get
+  Params{..} <- ask
+  ps <- sequence $ map (\(w, n) -> oneWordProb _counts _charProb _stop _dontStop _alpha n w) (Map.toList _counts)
+  let p = foldl1' (*) ps
+  return $! p
+
+
+-- | Run one sampling iteration
+sample :: (Probability p, Categorical p, Show p, MonadIO m, MonadRandom m, (MonadReader (Params p)) m, MonadState (SamplingState Char) m, MonadLog (WithSeverity String) m) => Int -> m ()
+sample i = do
+  ll <- unwrap <$> likelihood
+  state <- get
+  params <- ask
+  logInfo (printf "\nIteration #%d" i)
+  let indices = Set.fromList [i | (l, i) <- zip ((Vector.toList (_locations state))) [0..], _static l == False]  
+  iterateUntilM (\s -> Set.size s == 0) sampleSite indices
+  state' <- get
+  put $ state' { _counts=cleanCounts (_counts state'), _startLookup=cleanLookup (_startLookup state'), _endLookup=cleanLookup (_endLookup state') }
+  ll' <- unwrap <$> likelihood
+  let guesses' = Set.fromList $ Vector.toList $ Vector.findIndices (\x -> _morphFinal x && (not $ _static x)) (_locations state')
+      guesses = Set.fromList $ Vector.toList $ Vector.findIndices (\x -> _morphFinal x && (not $ _static x)) (_locations state)
+      score = f1 guesses (_gold params)
+      score' = f1 guesses' (_gold params)
+  logInfo (printf "Log-likelihood old/new: %.3v/%.3v\tF-Score old/new: %.3f/%.3f" ll ll' score score')
+  return $! ()
+
+
+formatMorphs :: [Vector Char] -> [Vector Char]
+formatMorphs ms = Vector.toList ms'
+  where
+    suff = Vector.fromList "@@"
+    ms' = Vector.imap (\i m -> if i == length ms - 1 then m else Vector.concat [m, suff]) (Vector.fromList ms)
+
+
+mapAccumLM :: (Monad m) => (a -> b -> m (a, c)) -> a -> [b] -> m (a, [c])
+mapAccumLM = mapAccumLM' []
+
+
+mapAccumLM' :: (Monad m) => [c] -> (a -> b -> m (a, c)) -> a -> [b] -> m (a, [c])
+mapAccumLM' cs f acc [] = return (acc, reverse cs)
+mapAccumLM' cs f acc (b:bs) = do
+  (acc', c) <- f acc b  
+  mapAccumLM' (c:cs) f acc' bs
+
+
+applyModel :: (MonadLog (WithSeverity String) m, Probability p, Show p) => Model p Char -> Dataset -> m Dataset
+applyModel model dataSet = do
+  let uniqueWords = (map Vector.fromList . Set.toList . Set.fromList . concat) dataSet
+      segCache = Map.empty :: SegCache p
+  logInfo (printf "Segmenting %d words" (length uniqueWords))
+  (sc, segs) <- mapAccumLM (segment model) segCache uniqueWords
+  let segMap = Map.fromList segs
+  return $ map (map Vector.toList . concat . map (\w -> segMap Map.! (Vector.fromList w))) dataSet
+
+
+type Table p = Map (Int, Int) p
+type SegCache p = Map (Vector Char) p
+
+
+type DPState prob = (SegCache prob, Table prob, Table Int)
+
+
+traceBack :: (MonadLog (WithSeverity String) m) => Table Int -> Int -> Vector Char -> m [Vector Char]
+traceBack pathTable end token = return $ go pathTable end token []
+  where
+    go pt 0 t acc = acc    
+    go pt e t acc = go pt e' t' (s:acc)
+      where
+        e' = pt Map.! (0, e)
+        (t', s) = Vector.splitAt e' t
+
+
+printTable :: (Show p, Probability p) => Table p -> Int -> String
+printTable table size = unlines rows
+  where
+    rows = map unwords cells
+    cells = [[case table Map.!? (r, c)  of Nothing -> "         "
+                                           Just p -> printf "%.7f" (toDouble p)
+             | c <- [1..size + 1]] | r <- [0..size]]
+
+
+printPathTable :: Table Int -> Int -> String
+printPathTable table size = unlines rows
+  where
+    rows = map unwords cells
+    cells = [[case table Map.!? (r, c)  of Nothing -> "  "
+                                           Just p -> printf "%.2d" p
+             | c <- [1..size + 1]] | r <- [0..size]]
+
+
+fillTable :: (MonadLog (WithSeverity String) m, Probability p, Show p) => Model p Char -> Vector Char -> DPState p -> (Int, Int) -> m (DPState p, p)
+fillTable model token (cache, probTable, pathTable) (from, to) = do
+  --logInfo (printf "Considering span from %d to %d" from to)
+  let ct = to - from
+      gram = Vector.slice from ct token
+      --cachedSeg = cache Map.!? gram
+      --noSegProb = model Map.!? gram
+      pairs = [(i, Vector.slice (from + i) (to - (from + i)) token) | i <- [0..ct - 1]]
+  
+  --logInfo (printf "Substring '%s'" (Vector.toList gram))
+  --logInfo (show pairs)
+  let scores = [(i, (Map.findWithDefault (fromDouble 1.0) (from, from + i) probTable) * (Map.findWithDefault (fromDouble 0.0) g model)) | (i, g) <- pairs]
+      --best = (maximumBy (comparing id) . catMaybes) ([noSegProb, Just (fromDouble 0.0)] ++ [])
+      (bestI, best) = (maximumBy (comparing snd)) scores --  . catMaybes) ([noSegProb, Just (fromDouble 0.0)] ++ [])
+      --case Vector.length gram of 0 -> fromDouble 1.0
+      --                                 --_ -> (maximumBy (comparing id)) scores
+      --                                 _ -> 
+      cache' = Map.insert gram best cache
+      probTable' = Map.insert (from, to) best probTable
+      pathTable' = Map.insert (from, to) bestI pathTable
+  --logInfo (show scores)
+  --logInfo (printf "No seg prob for %v: %s" gram (show noSegProb))
+  --logInfo (printf "Cache size: %d" (Map.size cache'))
+  
+  --logInfo (printTable probTable' (Vector.length token - 1))
+  --logInfo (printPathTable pathTable' (Vector.length token - 1))
+  return ((cache', probTable', pathTable'), best)
+
+
+segment :: (MonadLog (WithSeverity String) m, Probability p, Show p) => Model p Char -> SegCache p -> Vector Char -> m (SegCache p, (Vector Char, [Vector Char]))
+segment model cache token = do
+  --logInfo (printf "Segmenting '%s'" (Vector.toList token))
+  let max = Vector.length token
+      order = concat [[(from, to) | from <- reverse [0..to - 1]] | to <- [1..max]]
+      probTable = Map.empty :: Table p
+      pathTable = Map.empty :: Table Int
+  --logInfo (printf "Sequence of spans to consider: %s" (show order))
+  ((cache', probTable', pathTable'), _) <- mapAccumLM (fillTable model token) (cache, probTable, pathTable) order
+  --logInfo (printPathTable pathTable' (Vector.length token - 1))
+
+
+  toks <- traceBack pathTable' max token
+  
+  return (cache', (token, formatMorphs toks))
+
+
+splits :: Model p elem -> Vector elem -> [(Vector elem, Vector elem)]
+splits m w = [Vector.splitAt i w | i <- [1..Vector.length w]]
+
+
+segProb :: (Probability p, Ord elem) => Model p elem -> [Vector elem] -> p
+segProb m ws = product $ map (\w -> Map.findWithDefault (fromDouble 0.0) w m) ws --fromDouble 1.0
+
+
+combinations :: (MonadLog (WithSeverity String) m, Show p, Probability p) => Int -> m (Vector p)
+combinations n = do
+  return $ Vector.generate (n + 1) (fromDouble . fromIntegral . (n `choose`))
+
+
+-- | Compute the log-probability of generating the given word n times, based on counts
+oneWordProb :: (Show p, MonadLog (WithSeverity String) m, Probability p, Show elem, Ord elem) => Counts elem -> p -> p -> p -> p -> Int -> Morph elem -> m p
+oneWordProb counts charProb stopProb dontStopProb alpha n word = do
+  let mu = ((dontStopProb * charProb) ^ (length word)) * stopProb
+      total = fromIntegral $ sum $ Map.elems counts
+      count = fromIntegral $ Map.findWithDefault 0 word counts
+      numer = ((alpha * mu) + count)
+      denom = (alpha + total)
+  return $! ((numer ^ n) / (denom ^ n))
+
+
+oneTimeOneWord :: (Probability p, Ord elem) => Counts elem -> Params p -> Vector elem -> p
+oneTimeOneWord counts Params{..} word = p
+  where
+    mu = ((_dontStop * _charProb) ^ (Vector.length word)) * _stop
+    total = fromIntegral $ sum $ Map.elems counts
+    count = fromIntegral $ Map.findWithDefault 0 word counts
+    numer = ((_alpha * mu) + count)
+    denom = (_alpha + total)
+    p = numer / denom
+
+
+-- | Compute the log-probability of setting a single set of m sites, out of n, to positive
+g :: (Show p, MonadLog (WithSeverity String) m, Ord elem, Show elem, Probability p) => Counts elem -> p -> p -> p -> Morph elem -> Morph elem -> p -> Int -> Int -> m p
+g counts charProb stopProb dontStopProb before after alpha n m = do
+  beforeProb <- oneWordProb counts charProb stopProb dontStopProb alpha m before
+  afterProb <- oneWordProb counts charProb stopProb dontStopProb alpha m after
+  let posProb = beforeProb * afterProb
+  negProb <- oneWordProb counts charProb stopProb dontStopProb alpha (n - m) (before Vector.++ after)
+  return $! posProb * negProb
+
+
+-- | Compute the log-categorical distribution of possible number of sites to set to positive:
+--     P(m) = (n choose m) * g(m)
+distribution :: (Show p, MonadLog (WithSeverity String) m, Probability p, Show elem, Ord elem, Show p) => Counts elem -> p -> p -> p -> Morph elem -> Morph elem -> p -> Int -> m (Vector p)
+distribution counts charProb stopProb dontStopProb before after alpha n = do
+  gs <- (liftM Vector.fromList . sequence) [g counts charProb stopProb dontStopProb before after alpha n m | m <- [0..n]]
+  combs <- combinations n
+  let unScaled = Vector.map (\(x, y) -> x * y) (Vector.zip combs gs)
+  return $! unScaled
+
+
+-- | Randomly sample a site from those currently available, and then block-sample all compatible sites, returning the updated list of available sites
+sampleSite :: (Probability p, Categorical p, Show p, MonadIO m, MonadLog (WithSeverity String) m, MonadRandom m, MonadState (SamplingState Char) m, MonadReader (Params p) m) => Set Int -> m (Set Int)
+sampleSite ix = do
+  params@(Params{..}) <- ask
+  state@(SamplingState{..}) <- get
+  logDebug ('\n':(printf "%v" params))
+  logDebug (printf "%v" params)
+  i <- uniform ix
+  (a, b) <- siteToWords i
+  let c = a Vector.++ b
+  (fullSites', splitSites') <- wordsToSites i _startLookup _endLookup a b
+  let fullSites = Set.intersection fullSites' ix
+      splitSites = Set.intersection splitSites' ix      
+      sites = Set.union fullSites splitSites
+      nSplit = Set.size splitSites
+      nFull = Set.size fullSites
+      cs' = (subtractCounts c nFull . subtractCounts a nSplit . subtractCounts b nSplit) _counts
+  d <- distribution cs' _charProb _stop _dontStop a b _alpha (Set.size sites)
+  numPos <- sampleCategorical d
+  put state{ _counts=cleanCounts cs' }
+  logDebug (printf "Pivot: %d" i)
+  logDebug (printf "Morphs: left=%v, right=%v" (show a) (show b))
+  logDebug (printf "Matching, non-conflicting positive sites: [%v]" splitSites)
+  logDebug (printf "Matching, non-conflicting negative sites: [%v]" fullSites)  
+  logDebug (printf "Distribution: [%v]" (showDist d))
+  logDebug (printf "Chose positive count: %d" numPos)
+  sites' <- shuffleM (Set.toList sites)
+  let (pos, neg) = splitAt numPos sites'
+      pos' = Set.fromList pos
+      neg' = Set.fromList neg
+      nPos = length pos
+      nNeg = length neg
+      cs'' = (addCounts c nNeg . addCounts a nPos . addCounts b nPos) cs'
+      cs''' = Map.fromList $ [(k, v) | (k, v) <- Map.toList cs'', v /= 0]
+      locations' = updateLocations' (_value (_locations Vector.! i)) _locations pos' neg'
+      (upS, upE) = computeUpdates splitSites fullSites a b
+      luS' = Map.unionWith (Set.\\) _startLookup upS
+      luE' = Map.unionWith (Set.\\) _endLookup upE      
+      (upS', upE') = computeUpdates pos' neg' a b
+      luS = cleanLookup $ Map.unionWith Set.union luS' upS'
+      luE = cleanLookup $ Map.unionWith Set.union luE' upE'
+      --wordsLookup' = 
+      ix' = ix Set.\\ sites
+      --wordsLookup' = initReverseLookup luS luE
+      --wordsLookup' = updateReverseLookup _wordsLookup pos' neg' a b
+  put $ SamplingState cs''' locations' luS luE _wordsLookup ix'
+  return $! ix Set.\\ sites
+
+
+updateReverseLookup :: (Show elem) => ReverseLookup elem -> Set Int -> Set Int -> Vector elem -> Vector elem -> ReverseLookup elem
+updateReverseLookup rlu pos neg a b = rlu
+  where
+    --updates = error (show (pos, neg, a, b))
+    --negPrefUpdates = []
+    --negSuffUpdates = []
+    --posUpdates = []
+    --updates = Map.fromList (posUpdates ++ negPrefUpdates ++ negSuffUpdates)
diff --git a/src/Text/HaskSeg/Probability.hs b/src/Text/HaskSeg/Probability.hs
new file mode 100644
--- /dev/null
+++ b/src/Text/HaskSeg/Probability.hs
@@ -0,0 +1,68 @@
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+
+module Text.HaskSeg.Probability (Prob, LogProb, showDist, Probability(..), sampleCategorical, Categorical) where
+
+import Data.List (unfoldr, nub, mapAccumL, intercalate, sort)
+import Data.Vector (Vector)
+import qualified Data.Vector as Vector
+import Control.Monad.Random
+import Text.Printf (printf, PrintfArg(..), fmtPrecision, fmtChar, errorBadFormat, formatString, vFmt, IsChar)
+
+newtype Prob = Prob Double deriving (Show, Read, Eq, Ord, Num)
+type Dist = Vector Prob
+newtype LogProb = LogProb Double deriving (Show, Read, Eq, Ord)
+type LogDist = Vector LogProb
+
+showDist :: (Probability p, Show p) => Vector p -> String
+showDist ps = intercalate ", " $ (map (\v -> printf "%.8f" v :: String) . map (/ total)) ps'
+  where
+    ps' = map toDouble (Vector.toList ps)
+    total = sum ps'
+
+instance Num LogProb where
+  (LogProb a) + (LogProb b) = LogProb (l + (logBase 2 v))
+    where
+      (l, s) = if a > b then (a, b) else (b, a)
+      d = s - l
+      v = 1 + (2 ** d)
+  (LogProb a) * (LogProb b) = LogProb (a + b)
+  negate = undefined
+  abs = undefined
+  signum = undefined
+  fromInteger i = fromDouble (fromIntegral i)
+
+instance Fractional Prob where
+  recip (Prob a) = Prob (1.0 / a)
+  fromRational a = undefined
+  
+instance Fractional LogProb where  
+  recip (LogProb a) = LogProb (-a)
+  fromRational a = undefined
+  
+class (Ord p, Num p, Fractional p) => Probability p where
+  fromDouble :: Double -> p
+  toDouble :: p -> Double
+  unwrap :: p -> Double
+  
+class (Probability p) => Categorical p where
+  sampleCategorical :: (MonadRandom m) => Vector p -> m Int
+  sampleCategorical xps = do
+    let sums = Vector.scanl (+) (fromDouble 0.0 :: p) xps
+        maxP = toDouble $ Vector.last sums
+    v <- getRandomR (0.0, maxP)
+    let v' = fromDouble v :: p
+    return (Vector.length (Vector.takeWhile (\x -> x < v') sums) - 1)
+
+instance Probability LogProb where
+  fromDouble p = LogProb (logBase 2 p)
+  toDouble (LogProb lp) = 2 ** lp
+  unwrap (LogProb lp) = lp
+  
+instance Probability Prob where
+  fromDouble p = Prob p
+  toDouble (Prob p) = p
+  unwrap (Prob p) = p
+
+instance Categorical LogProb
+instance Categorical Prob
diff --git a/src/Text/HaskSeg/Types.hs b/src/Text/HaskSeg/Types.hs
new file mode 100644
--- /dev/null
+++ b/src/Text/HaskSeg/Types.hs
@@ -0,0 +1,80 @@
+module Text.HaskSeg.Types (Locations, Morph, Counts, Site, Location(..), Lookup, showLookup, showCounts, SamplingState(..), Params(..), Model, Token, Sentence, Dataset, Filename, Vocabulary, Segmentation, ReverseLookup) where
+
+import Data.List (unfoldr, nub, mapAccumL, intercalate, sort)
+import Data.Map (Map)
+import qualified Data.Map as Map
+import Data.Set (Set)
+import qualified Data.Set as Set
+import Text.Printf (printf, PrintfArg(..), fmtPrecision, fmtChar, errorBadFormat, formatString, vFmt, IsChar)
+import Data.Vector (Vector)
+import qualified Data.Vector as Vector
+import Data.Foldable (toList)
+import Text.HaskSeg.Probability (Probability)
+
+type Token = String
+type Sentence = [Token]
+type Dataset = [Sentence]
+type Filename = String
+type Vocabulary = Set Token
+type Segmentation = Map Token [Token]
+
+
+
+type Locations elem = Vector (Location elem)
+type Morph elem = Vector elem
+type Counts elem = Map (Morph elem) Int
+type Site = Int
+
+type Model elem p = Map [elem] p
+
+data Location elem = Location { _value :: !elem
+                              , _morphFinal :: !Bool
+                              , _static :: !Bool
+                              } deriving (Show, Read)
+
+-- | A "start" lookup points to the boundary *before* the first item, an "end" lookup points to the boundary *of* the last item
+type Lookup elem = Map (Morph elem) (Set Int)
+
+type ReverseLookup elem = Map Int (Morph elem, Morph elem)
+
+showLookup :: (PrintfArg elem, IsChar elem) => Lookup elem -> String
+showLookup lu = intercalate ", " [printf "\"%v\"=[%v]" (toList k) v | (k, v) <- Map.toList lu]
+
+showCounts :: (PrintfArg elem, IsChar elem) => Counts elem -> String
+showCounts cs = intercalate ", " [printf "\"%v\"=%d" (toList k) v | (k, v) <- Map.toList cs]
+
+-- | A coherent state of boundary assignments, counts, and word start/end lookups
+data SamplingState elem = SamplingState { _counts :: !(Counts elem)
+                                        , _locations :: !(Locations elem)
+                                        , _startLookup :: !(Lookup elem)
+                                        , _endLookup :: !(Lookup elem)
+                                        , _wordsLookup :: !(ReverseLookup elem)
+                                        , _toSample :: !(Set Int)
+                                        } deriving (Show, Read)
+
+instance Show elem => PrintfArg (SamplingState elem) where
+  formatArg SamplingState{..} fmt | fmtChar (vFmt 'P' fmt) == 'P' = formatString (printf "SamplingState" :: String) (fmt { fmtChar = 's', fmtPrecision = Nothing })
+  formatArg _ fmt = errorBadFormat $ fmtChar fmt
+  
+-- | Parameters that are set at training time
+data Params p = Params { _alpha :: !p
+                       , _stop :: !p
+                       , _dontStop :: !p
+                       , _spaces :: !Bool
+                       , _types :: !Bool
+                       , _gold :: !(Set Int)
+                       , _charProb :: !p
+                       , _minCount :: !Int
+                       } deriving (Show, Read)
+
+instance Show p => PrintfArg (Params p) where
+  formatArg Params{..} fmt | fmtChar (vFmt 'P' fmt) == 'P' = formatString (printf "Params: alpha=%v, stopProb=%v, dontStop=%v, uniformCharProb=%v" (show _alpha) (show _stop) (show _dontStop) (show _charProb) :: String) (fmt { fmtChar = 's', fmtPrecision = Nothing })
+  formatArg _ fmt = errorBadFormat $ fmtChar fmt
+
+instance PrintfArg (Set Int) where
+  formatArg is fmt | fmtChar (vFmt 'P' fmt) == 'P' = formatString (intercalate ", " ((map show . Set.toList) is)) (fmt { fmtChar = 's', fmtPrecision = Nothing })
+  formatArg _ fmt = errorBadFormat $ fmtChar fmt
+
+instance (Show elem) => PrintfArg (Vector elem) where
+  formatArg is fmt | fmtChar (vFmt 'P' fmt) == 'P' = formatString (intercalate ", " ((map show . Vector.toList) is)) (fmt { fmtChar = 's', fmtPrecision = Nothing })
+  formatArg _ fmt = errorBadFormat $ fmtChar fmt
diff --git a/src/Text/HaskSeg/Utils.hs b/src/Text/HaskSeg/Utils.hs
new file mode 100644
--- /dev/null
+++ b/src/Text/HaskSeg/Utils.hs
@@ -0,0 +1,87 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE ExplicitNamespaces #-}
+
+module Text.HaskSeg.Utils (readDataset, writeDataset, writeState, readState, datasetToVocabulary, applySegmentation) where
+
+import Prelude hiding (lookup, getContents, readFile, strip, lines, writeFile, words)
+import System.IO (withFile, IOMode(..), stdin, stderr, openFile, stdout, hClose, Handle(..))
+import Data.Text (Text, strip, lines, stripPrefix, splitOn, pack, unpack, words)
+import Data.Text.IO (getContents, readFile, hGetContents, hPutStr, writeFile, hPutStrLn)
+--import qualified Data.ByteString.Lazy as BS
+import qualified Data.ByteString.Lazy.Char8 as BS
+import Data.Text (Text)
+import qualified Data.Text.Lazy as T
+import qualified Data.Text.Lazy.IO as T
+import qualified Data.Text.Lazy.Encoding as T
+import Control.Monad (join, liftM, foldM)
+import Data.Set (Set)
+import qualified Data.Set as Set
+import Data.Map (Map)
+import qualified Data.Map as Map
+import Data.List (nub)
+
+import Codec.Compression.GZip (compress, decompress)
+import Text.HaskSeg.Types (Locations, Morph, Counts, Site, Location(..), Lookup, showLookup, showCounts, SamplingState(..), Params(..), Model, Token, Sentence, Dataset)
+import Text.HaskSeg.Probability (Probability)
+
+--type Token = String
+--type Sentence = [Token]
+--type Dataset = [Sentence]
+type Filename = String
+type Vocabulary = Set Token
+type Segmentation = Map Token [Token]
+
+
+readFileOrStdin :: Maybe String -> IO Text
+readFileOrStdin (Just f) = case suf of "gz" -> (liftM (pack . BS.unpack . decompress . BS.pack . unpack) . readFile) f
+                                       _ -> readFile f
+  where
+    suf = (reverse . take 2 . reverse) f
+readFileOrStdin Nothing = getContents
+
+
+writeFileOrStdout :: Maybe String -> Text -> IO ()
+writeFileOrStdout (Just f) s = case suf of "gz" -> writeFile f ((pack . BS.unpack . compress . BS.pack . unpack) s)
+                                           _ -> writeFile f s
+  where
+    suf = (reverse . take 2 . reverse) f
+writeFileOrStdout Nothing s = hPutStr stdout s
+
+
+readDataset :: Maybe Filename -> Maybe Int -> IO Dataset
+readDataset (Just f) n = do
+  bs <- readFile f
+  let ls = (map words . (case n of Nothing -> id; Just i -> take i) . lines) bs
+  --let ls = (map words . (case n of Nothing -> id; Just i -> take i) . lines . T.unpack . T.decodeUtf8) bs
+  return $ map (map unpack) ls
+
+datasetToVocabulary :: Dataset -> Vocabulary
+datasetToVocabulary ss = Set.fromList $ nub ws
+  where
+    ws = concat ss
+
+writeDataset :: Maybe Filename -> Dataset -> IO ()
+writeDataset (Just f) cs = BS.writeFile f bs
+  where
+    bs = (T.encodeUtf8 . T.pack . unlines . map unwords) cs
+
+applySegmentation :: Segmentation -> Dataset -> Dataset
+applySegmentation seg ds = map (concat . (map (\w -> Map.findWithDefault [[c] | c <- w] w seg))) ds
+
+
+--readVocabulary :: Filename -> IO Dataset
+--readVocabulary f = undefined
+
+--writeVocabulary :: Filename -> Dataset -> IO ()
+--writeVocabulary f d = undefined
+
+writeState :: (Show a, Show p) => Maybe Filename -> Params p -> Locations a -> IO ()
+writeState (Just f) p l = BS.writeFile f ((compress . T.encodeUtf8 . T.pack . show) $ (p, l))
+
+readState :: (Read a, Read p) => Maybe Filename -> IO (Params p, Locations a)
+readState (Just f) = (liftM (read . T.unpack . T.decodeUtf8 . decompress) . BS.readFile) f