-- | Some data structures and functions related to stops and lines.

{-# LANGUAGE ScopedTypeVariables #-}

module Lab3Help
  ( -- * Data types
    Stop(..)
  , LineTable(..)
  , LineStop(..)
    -- * Parsing of stop/line files
  , readStops
  , readLines
    -- * Tries
  , Trie
  , member
  ) where

import Control.Applicative ((<$>))
import Control.Monad
import Data.List hiding (insert)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe

------------------------------------------------------------------------
-- Data types

-- | Description of a (bus\/tram\/…) stop.

data Stop = Stop
  { name     :: String
                -- ^ The stop's name.
  , position :: (Integer, Integer)
                -- ^ Coordinates (x, y).
                --
                -- Invariant: The coordinates range from 0 to 1000.
  } deriving Show

-- | Description of a line.

data LineTable = LineTable
  { lineNumber :: Integer
                  -- ^ The line number.
  , stops      :: [LineStop]
                  -- ^ Stops (in the order in which they are visited).
                  --
                  -- Invariant: The list contains at least one
                  -- element.
  } deriving Show

-- | Description of a stop, as part of a line.

data LineStop = LineStop
  { stopName :: String
                -- ^ The stop's name.
  , time     :: Integer
                -- ^ The travel time (in minutes) from the previous
                -- stop.
                --
                -- Invariant: A non-negative number, 0 if this is the
                -- first stop.
  } deriving Show

------------------------------------------------------------------------
-- Parsing of stop/line files

-- | Tries to parse a stop.
--
-- If an error is encountered, then an error message is returned
-- (@'Left' msg@).

readStop :: [String] -> Either String Stop
readStop ss = case ss of
  name : coordinates@[_, _] ->
    case mapM readNat coordinates of
      Just cs@[x, y] | all inRange cs ->
        Right (Stop { name = name, position = (x, y) })
      _ -> err
  _ -> err
  where
  inRange c = 0 <= c && c <= 1000
  err       = Left $ "Incorrectly formatted stop:\n" ++ unwords ss

-- | Tries to read a list of stops from the given file.
--
-- If the list is read successfully, then a trie with all the stop
-- names is returned along with the list.
--
-- If an error is encountered, then an error message is returned
-- (@'Left' msg@).

readStops :: FilePath -> IO (Either String ([Stop], Trie Char))
readStops f = process <$> readFile f
  where
  process s = do
    ss <- mapM readStop $ split 3 $ words s
    case duplicate (map name ss) of
      Nothing -> return (ss, foldr insert empty (map name ss))
      Just d  -> Left $ "Duplicate stop: " ++ d

-- | Tries to parse a \"line stop\" (not the first one for a line,
-- i.e., the 'time' must be given).
--
-- If an error is encountered, then an error message is returned
-- (@'Left' msg@).

readLineStop :: [String] -> Either String LineStop
readLineStop ss = case ss of
  [name, time] -> case readNat time of
    Just time -> Right (LineStop { stopName = name, time = time })
    _         -> err
  _ -> err
  where
  err = Left $ "Incorrectly formatted line stop:\n" ++ unwords ss

-- | Tries to read information about lines from the given file.
--
-- If an error is encountered, then an error message is returned
-- (@'Left' msg@).

readLines
  :: Trie Char
     -- ^ Known stop names. If an unknown stop is encountered, then an
     -- error message is returned.
  -> FilePath
  -> IO (Either String [LineTable])
readLines known f = parse . words <$> readFile f
  where
  parse []           = return []
  parse (l : n : ss) =
    case mapM readNat [l, n] of
      Just [line, noStops]
        | noStops <= 0 -> Left $ "Empty lines not allowed:\n" ++
                                 unwords [l, n]
        | otherwise    -> do
          let noTokens      = 2 * noStops - 1
              (stops, rest) = genericSplitAt noTokens ss
          when (genericLength stops < noTokens) $ do
            Left $ "Incorrectly formatted line:\n" ++
                   unwords (l : n : stops)
          case stops of
            stop : stops -> do
              stop  <- isKnown (LineStop { stopName = stop, time = 0 })
              stops <- mapM (isKnown <=< readLineStop) (split 2 stops)
              rest  <- parse rest
              return (LineTable { lineNumber = line
                                , stops      = stop : stops
                                } : rest)
      _ -> Left $ "Incorrectly formatted line header:\n" ++
                  unwords [l, n]

  isKnown stop
    | member (stopName stop) known = return stop
    | otherwise                    =
      Left $ "Unknown stop: " ++ stopName stop

------------------------------------------------------------------------
-- Some utility functions

-- | Tries to parse a non-negative integer.

readNat :: String -> Maybe Integer
readNat s = case filter (null . snd) $ reads s of
  [(x, _)] | x >= 0 -> Just x
  _                 -> Nothing

-- | @split n xs@ splits up @xs@ into groups of size @n@ (and possibly
-- a final, shorter group).
--
-- Time complexity: /O(length of the list)/.

split :: Integer -> [a] -> [[a]]
split n [] = []
split n xs = ys : split n zs
  where (ys, zs) = genericSplitAt n xs

-- | If the list contains any duplicated elements, then one of them is
-- returned.
--
-- Time complexity: /O(c n log n)/, where /n/ is the length of the
-- list, and /c/ is an upper bound on the time complexity of a
-- comparison.

duplicate :: Ord a => [a] -> Maybe a
duplicate =
  listToMaybe .
  map head .
  filter ((>= 2) . length) .
  group .
  sort

------------------------------------------------------------------------
-- An implementation of tries

-- | Tries.

data Trie a = Trie !Bool !(Map a (Trie a))

-- | An empty trie.
--
-- Time complexity: /O(1)/.

empty :: Trie a
empty = Trie False Map.empty

-- | Inserts the list into the trie.
--
-- Time complexity: /O(n c log v)/, where /n/ is the length of the
-- list, /c/ is an upper bound on the time complexity of a comparison,
-- and /v/ is the number of distinct list elements in the trie. For
-- 'Char' this can be simplified to /O(n)/.

insert :: Ord a => [a] -> Trie a -> Trie a
insert []       (Trie _ m) = Trie True m
insert (x : xs) (Trie b m) =
  Trie b (Map.alter (Just . insert xs . fromMaybe empty) x m)

-- | Is the list a member of the trie?
--
-- Time complexity: /O(n c log v)/, where /n/ is the length of the
-- list, /c/ is an upper bound on the time complexity of a comparison,
-- and /v/ is the number of distinct list elements in the trie. For
-- 'Char' this can be simplified to /O(n)/.

member :: Ord a => [a] -> Trie a -> Bool
member []       (Trie b _) = b
member (x : xs) (Trie _ m) = case Map.lookup x m of
  Nothing -> False
  Just t  -> member xs t

------------------------------------------------------------------------
-- A very small trie test suite

prop_member_empty :: [Integer] -> Bool
prop_member_empty xs = not (member xs empty)

prop_member_insert :: [[Integer]] -> [Integer] -> [[Integer]] -> Bool
prop_member_insert xss ys zss =
  member ys (foldr insert (insert ys (foldr insert empty zss)) xss)