diff --git a/README.md b/README.md
--- a/README.md
+++ b/README.md
@@ -1,4 +1,4 @@
-# haskmorph
+# haskseg
 
 ## Compiling
 
@@ -23,11 +23,25 @@
 Invoke the program using Stack.  To see available sub-commands, run:
 
 ```
-stack exec -- haskmorph -h
+stack exec -- haskseg -h
 ```
 
 To see detailed help, run e.g.:
 
 ```
-stack exec -- haskmorph train -h
+stack exec -- haskseg train -h
 ```
+
+To train on the included data set from Goldwater, Griffiths, and Johnson 2009, run:
+
+```
+time zcat data/br-old-gold.txt.gz | perl -pe '$_=~s/ /\<GOLD\>/g;' | stack exec -- haskseg train --stateFile model.gz --iterations 3 --goldString "<GOLD>"
+```
+
+Note here the spaces are being replaced by a special string, to indicate boundaries to calculate F-score on (not necessary, but a nice way to track progress).  The model seems to converge after three iterations on this data set.  This takes about three minutes and achieves F-score of 0.61, somewhat higher than reported in Liang, Jordan and Klein 2010 (why is unclear).  To use the trained model to segment text, such as the training data set, run:
+
+```
+time zcat data/br-old-gold.txt.gz | perl -pe '$_=~s/ //g;' | stack exec -- haskseg segment --stateFile model.gz
+```
+
+Note that for this stage, spaces are simply removed, otherwise it treats whitespace as a static boundary (which may be what you want in other circumstances!).  The model takes about 4 seconds to segment the 9790 "words", about 2500 per second, though this stage is embarassingly parallel.  The output is in BPE format, i.e. with "@@" indicating the end of an internal morph.
diff --git a/Setup.hs b/Setup.hs
deleted file mode 100644
--- a/Setup.hs
+++ /dev/null
@@ -1,2 +0,0 @@
-import Distribution.Simple
-main = defaultMain
diff --git a/app/Main.hs b/app/Main.hs
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -13,9 +13,11 @@
 module Main where
 
 import Prelude hiding (lookup)
+import Data.Text (pack, unpack)
 import Options.Generic (Generic, ParseRecord, Unwrapped, Wrapped, unwrapRecord, (:::), type (<?>)(..))
 import Control.Monad (join, liftM, foldM)
 import System.Random (getStdGen, mkStdGen)
+import System.IO (stderr, hPutStrLn)
 import Data.List (unfoldr, nub, mapAccumL, intercalate, sort, foldl1')
 import Data.Maybe (fromMaybe, catMaybes)
 import Data.Map (Map)
@@ -36,10 +38,10 @@
 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.Types (Locations, Morph, Counts, Site, Location(..), Lookup, showLookup, showCounts, SamplingState(..), Params(..), Model(..), Character)
 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.Utils (readState, readDataset, writeDataset, writeState, datasetToVocabulary, applySegmentation, writeFileOrStdout, readFileOrStdin)
+import Text.HaskSeg.Location (randomFlip, 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)
@@ -63,19 +65,21 @@
                           , 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)"
+                          , sharpParam :: w ::: Maybe Double <?> "Probability to stop generating characters (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)"
+                          , typeBased :: w ::: Bool <?> "Run over types (i.e. don't consider frequency) rather than tokens (default: false)"
+                          , goldString :: w ::: Maybe String <?> "Treat occurrences of the given string as gold boundaries for evaluation"
+                          , logLevel :: w ::: Maybe String <?> "Log severity threshold, 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"
+                            , borderCharacters :: w ::: Maybe String <?> "Characters to consider a static (unsampled) border"
+                            , goldString :: w ::: Maybe String <?> "Treat occurrences of the given string as gold boundaries for evaluation"
                             , 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)"
+                            , logLevel :: w ::: Maybe String <?> "Log severity threshold, one of [debug, info, warn, error] (default: info)"
                             , randomSeed :: w ::: Maybe Int <?> "Set a deterministic random seed (default: use system RNG)"
                             }
   deriving (Generic)                              
@@ -85,24 +89,22 @@
 
 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
+--
+--   Top-level actions
+--
 
+-- | Train a model on given data
+trainModel :: (Categorical p, Show p, Probability p, MonadIO m, MonadLog (WithSeverity String) m) => Vector (Location Char) -> Double -> (Params p) -> Int -> StdGen -> m (SamplingState Char)
+trainModel locations eta params iterations gen = do
+  logInfo (printf "Initial random seed: %v" (show gen))
   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))
+      gold = Set.empty
   let rLookup = initReverseLookup lookupS lookupE
-  --logInfo (show rLookup)
-  let state = SamplingState counts locations' lookupS lookupE rLookup Set.empty
+      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'
 
@@ -120,36 +122,23 @@
 
   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
+                    dataSet <- liftIO $ readDataset inputFile lineCount goldString
+                    let seq = (Vector.fromList . concat) dataSet
+                        numChars = (length . nub . map _value . 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)
+                        params = Params (fromDouble $ fromMaybe 0.1 alphaParam) (fromDouble $ fromMaybe 0.5 sharpParam) (fromDouble $ 1.0 - (fromMaybe 0.5 sharpParam)) False typeBased Set.empty charProb (fromMaybe 1 minCount)
                     state <- trainModel seq (fromMaybe 1.0 etaParam) params iterations gen
-                    liftIO $ writeState stateFile params (_locations state) 
+                    liftIO $ writeState stateFile params (_locations state)
+                    return ()
                   Segment{..} -> do
-                    dataSet <- liftIO $ readDataset inputFile lineCount
+                    toSeg <- liftIO $ (liftM ((map words) . lines . unpack) . readFileOrStdin) inputFile
                     (params :: Params ProbRep, locs :: Locations Char) <- liftIO $ readState stateFile
-                    let cs = nub $ concat (map concat dataSet)
+                    let cs = nub $ concat (map concat toSeg)
                     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
+                    segmented <- sequence $ map (applyModel model) toSeg
+                    liftIO $ writeFileOrStdout segmentationFile (pack (intercalate "\n" segmented ++ "\n"))
               )
                     (\msg -> case msgSeverity msg <= level of
-                               True -> putStrLn (discardSeverity msg)
+                               True -> hPutStrLn stderr (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
--- a/haskseg.cabal
+++ b/haskseg.cabal
@@ -1,14 +1,16 @@
--- This file has been generated from package.yaml by hpack version 0.28.2.
+cabal-version: 1.12
+
+-- This file has been generated from package.yaml by hpack version 0.31.2.
 --
 -- see: https://github.com/sol/hpack
 --
--- hash: 24fc8394d107cbe0c68b555dc7c3d0d3ed5edd305cce441cd1616785bf7ae2ed
+-- hash: 5c17f1d3fb53df989570fc2769e03a9e9156a94b8abf108d7d0997474b9c68fd
 
 name:           haskseg
-version:        0.1.0.1
+version:        0.1.0.2
 synopsis:       Simple unsupervised segmentation model
-description:    Implementation of the non-parametric segmentation model described in "Type-based MCMC" (Liang, Jordan, and Klein, 2010).
-category:       natural-language-processing, machine-learning
+description:    Implementation of the non-parametric segmentation model described in "Type-based MCMC" (Liang, Jordan, and Klein, 2010) and "A Bayesian framework for word segmentation Exploring the effects of context" (Goldwater, Griffiths, and Johnson, 2009).
+category:       NLP
 homepage:       https://github.com/githubuser/haskseg#readme
 author:         Tom Lippincott
 maintainer:     tom@cs.jhu.edu
@@ -16,7 +18,6 @@
 license:        BSD3
 license-file:   LICENSE
 build-type:     Simple
-cabal-version:  >= 1.10
 extra-source-files:
     README.md
 
diff --git a/src/Text/HaskSeg/Location.hs b/src/Text/HaskSeg/Location.hs
--- a/src/Text/HaskSeg/Location.hs
+++ b/src/Text/HaskSeg/Location.hs
@@ -29,25 +29,27 @@
 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')
+createData :: (Probability p, MonadLog (WithSeverity String) m) => (Params p) -> Vector (Char, Bool) -> m (Locations Char, Set Int)
+createData = undefined
+-- createData Params{..} cs' = do
+--   let cs = (map fst . 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" (length cs)) -- (length ws))
+--   return $! (ws'', Set.fromList bs')
 
 
 formatWord :: [Location Char] -> String
@@ -84,8 +86,8 @@
   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
+    pos' = (Vector.map (\i -> (i, (ls Vector.! i) { _morphFinal=True} )) . Vector.fromList . Set.toList) pos
+    neg' = (Vector.map (\i -> (i, (ls Vector.! i) { _morphFinal=False} )) . Vector.fromList . Set.toList) neg
     updates = pos' Vector.++ neg'
 
 
@@ -101,13 +103,13 @@
     
 
 -- | 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
+-- 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 True
 
 
 -- -- | Find the two words implied by a boundary at the given site
diff --git a/src/Text/HaskSeg/Metrics.hs b/src/Text/HaskSeg/Metrics.hs
--- a/src/Text/HaskSeg/Metrics.hs
+++ b/src/Text/HaskSeg/Metrics.hs
@@ -11,19 +11,25 @@
 import Data.Set (Set)
 import qualified Data.Set as Set
 
-precision :: Set Int -> Set Int -> Double
-precision guesses golds = tp / gs
+precision :: [Bool] -> [Bool] -> Double
+precision guesses golds = numer / denom
   where
-    tp = (fromIntegral . Set.size . Set.intersection guesses) golds
-    gs = (fromIntegral . Set.size) guesses
+    denom = (fromIntegral . length . filter (== True)) guesses
+    numer = (fromIntegral . length . filter (== (True, True))) (zip guesses golds)
+--    tp = (fromIntegral . Set.size . Set.intersection guesses) golds
+--    gs = (fromIntegral . Set.size) guesses
 
-recall :: Set Int -> Set Int -> Double
-recall guesses golds = tp / gs
+recall :: [Bool] -> [Bool] -> Double
+recall guesses golds = numer / denom
   where
-    tp = (fromIntegral . Set.size . Set.intersection guesses) golds
-    gs = (fromIntegral . Set.size) golds
+    denom = (fromIntegral . length . filter (== True)) golds
+    numer = (fromIntegral . length . filter (== (True, True))) (zip guesses golds)    
+--1.0 --tp / gs
+--  where
+--    tp = (fromIntegral . Set.size . Set.intersection guesses) golds
+--    gs = (fromIntegral . Set.size) golds
 
-f1 :: Set Int -> Set Int -> Double
+f1 :: [Bool] -> [Bool] -> Double
 f1 guesses golds = 2.0 * numer / denom
   where
     p = precision guesses golds
diff --git a/src/Text/HaskSeg/Model.hs b/src/Text/HaskSeg/Model.hs
--- a/src/Text/HaskSeg/Model.hs
+++ b/src/Text/HaskSeg/Model.hs
@@ -62,16 +62,21 @@
   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]  
+  logInfo (printf "\nIteration #%d, current vocab size %d" i (Map.size $ _counts state))
+  let locs = _locations state
+      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)
+  let locs = (filter (\x -> _static x /= True) . Vector.toList) (_locations state)
+      locs' = (filter (\x -> _static x /= True) . Vector.toList) (_locations state')
+      guesses = (map _morphFinal) locs
+      guesses' = (map _morphFinal) locs'
+      golds = (map _goldFinal) locs
+      golds' = (map _goldFinal) locs'
+      score = f1 guesses golds
+      score' = f1 guesses' golds'
   logInfo (printf "Log-likelihood old/new: %.3v/%.3v\tF-Score old/new: %.3f/%.3f" ll ll' score score')
   return $! ()
 
@@ -93,21 +98,18 @@
   (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
+  
+applyModel :: (MonadLog (WithSeverity String) m, Probability p, Show p) => Model p Char -> [String] -> m String
+applyModel model words = do
+  (sc, segs) <- mapAccumLM (segment model) (Map.empty :: SegCache p) (map Vector.fromList words) --uniqueWords
+  let segs' = concat (map snd segs)
+      segs'' = map Vector.toList segs'
+      segs''' = intercalate " " segs''
+  return segs'''
 
 
 type Table p = Map (Int, Int) p
 type SegCache p = Map (Vector Char) p
-
-
 type DPState prob = (SegCache prob, Table prob, Table Int)
 
 
diff --git a/src/Text/HaskSeg/Types.hs b/src/Text/HaskSeg/Types.hs
--- a/src/Text/HaskSeg/Types.hs
+++ b/src/Text/HaskSeg/Types.hs
@@ -1,4 +1,4 @@
-module Text.HaskSeg.Types (Locations, Morph, Counts, Site, Location(..), Lookup, showLookup, showCounts, SamplingState(..), Params(..), Model, Token, Sentence, Dataset, Filename, Vocabulary, Segmentation, ReverseLookup) where
+module Text.HaskSeg.Types (Locations, Morph, Counts, Site, Location(..), Lookup, showLookup, showCounts, SamplingState(..), Params(..), Model, Token, Sentence, Dataset, Filename, Vocabulary, Segmentation, ReverseLookup, Character) where
 
 import Data.List (unfoldr, nub, mapAccumL, intercalate, sort)
 import Data.Map (Map)
@@ -11,15 +11,13 @@
 import Data.Foldable (toList)
 import Text.HaskSeg.Probability (Probability)
 
-type Token = String
+type Character = (Char, Bool)
+type Token = [Character]
 type Sentence = [Token]
-type Dataset = [Sentence]
+type Dataset = [Token] --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
@@ -30,6 +28,7 @@
 data Location elem = Location { _value :: !elem
                               , _morphFinal :: !Bool
                               , _static :: !Bool
+                              , _goldFinal :: !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
@@ -53,7 +52,9 @@
                                         } 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 SamplingState{..} fmt | fmtChar (vFmt 'P' fmt) == 'P' = formatString (printf "SamplingState: %v" lcs :: String) (fmt { fmtChar = 's', fmtPrecision = Nothing })
+    where
+      lcs = show _locations
   formatArg _ fmt = errorBadFormat $ fmtChar fmt
   
 -- | Parameters that are set at training time
diff --git a/src/Text/HaskSeg/Utils.hs b/src/Text/HaskSeg/Utils.hs
--- a/src/Text/HaskSeg/Utils.hs
+++ b/src/Text/HaskSeg/Utils.hs
@@ -6,7 +6,7 @@
 {-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE ExplicitNamespaces #-}
 
-module Text.HaskSeg.Utils (readDataset, writeDataset, writeState, readState, datasetToVocabulary, applySegmentation) where
+module Text.HaskSeg.Utils (readDataset, writeDataset, writeState, readState, datasetToVocabulary, applySegmentation, writeFileOrStdout, readFileOrStdin) where
 
 import Prelude hiding (lookup, getContents, readFile, strip, lines, writeFile, words)
 import System.IO (withFile, IOMode(..), stdin, stderr, openFile, stdout, hClose, Handle(..))
@@ -23,17 +23,15 @@
 import qualified Data.Set as Set
 import Data.Map (Map)
 import qualified Data.Map as Map
-import Data.List (nub)
+import Data.List (nub, unfoldr, isPrefixOf)
+import Data.Maybe (fromMaybe)
 
 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 Vocabulary = [Token]
 type Segmentation = Map Token [Token]
 
 
@@ -53,25 +51,59 @@
 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
+lineToChars :: String -> Text -> [(Char, Bool)]
+lineToChars g t = t''
+  where
+    t' = unpack t
+    lineToChars' :: String -> Maybe ((Char, Bool), [Char])
+    lineToChars' [] = Nothing
+    lineToChars' (v:i) = Just ((v, g'), i')
+      where
+        g' = (g /= "" && g `isPrefixOf` i) || (i == [])
+        i' = if g' then drop (length g) i else i        
+    t'' = unfoldr lineToChars' t'
+
+lineToLocations :: String -> Text -> [Location Char]
+lineToLocations g t = t''
+  where
+    t' = unpack t
+    lineToChars' :: String -> Maybe (Location Char, [Char])
+    lineToChars' [] = Nothing
+    lineToChars' (v:[]) = Just (Location v True True True, [])
+    lineToChars' (v:i) = Just (Location v False False g', i')
+      where
+        g' = (g /= "" && g `isPrefixOf` i) || (i == [])
+        i' = if g' then drop (length g) i else i        
+    t'' = unfoldr lineToChars' t'
+
+
+readDataset :: Maybe Filename -> Maybe Int -> Maybe String -> IO [[Location Char]]
+readDataset mf n g = do
+  bs <- case mf of Just f -> readFile f
+                   Nothing -> getContents                   
+  --let ls = (map words . (case n of Nothing -> id; Just i -> take i) . lines) bs
+  let ls = lines bs  
+      g' = fromMaybe "" g
   --let ls = (map words . (case n of Nothing -> id; Just i -> take i) . lines . T.unpack . T.decodeUtf8) bs
-  return $ map (map unpack) ls
+  return $ map (lineToLocations g') ls
 
+  
 datasetToVocabulary :: Dataset -> Vocabulary
-datasetToVocabulary ss = Set.fromList $ nub ws
-  where
-    ws = concat ss
+datasetToVocabulary = undefined
+--datasetToVocabulary ss = 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
+writeDataset = undefined
+-- writeDataset mf cs = case mf of Just f -> BS.writeFile f bs
+--                                 Nothing -> BS.hPut stdout 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
+applySegmentation = undefined
+--applySegmentation seg ds = map (concat . (map (\w -> Map.findWithDefault [[c] | c <- w] w seg))) ds
 
 
 --readVocabulary :: Filename -> IO Dataset
