parse.hs

{-# LANGUAGE OverloadedStrings #-}

module Main where 

import Text.Parsec
import Text.Parsec.String
import Text.Parsec.Char
import Control.Monad
import Data.List
import Data.Either
import System.Console.ArgParser hiding (Parser)
import Control.Applicative hiding (many)
import qualified Data.List.Split as S
import qualified Data.Map.Strict as M

data ChannelMask = ChannelMask [Bool] deriving (Show, Eq)
data Event = Event { _eventClockTick :: Int, _eventMask :: ChannelMask } deriving Show
data Record = Record { _recordEventCount :: Int, _recordEvents :: [Event] } deriving Show
data ListModeData = ListModeData { _listModeDataSegmentCount :: Int, _listModeRecords :: [Record] } deriving Show

simpleParse :: Parser a -> String -> String -> Either ParseError a
simpleParse = parse

parseChannelMask :: Parser ChannelMask
parseChannelMask = do
    bits <- many1 digit
    return $ ChannelMask (map ind bits)
        where ind '1' = True
              ind '0' = False
              ind _ = error "Bad bit in channel mask"

parseTick :: Parser Int
parseTick = do
    v <- many digit `sepBy` char ','
    return $ (read . concat) v

parseEvent :: Parser Event
parseEvent = do
    skipMany space
    tick <- parseTick
    skipMany space
    mask <- parseChannelMask
    return Event {
          _eventClockTick = tick
        , _eventMask = mask
        }

parseEventCount :: Parser Int
parseEventCount = do
    string "events:"
    skipMany space
    count <- many1 digit
    return . read $ count

parseRecord :: Parser Record
parseRecord = do
    counts <- parseEventCount
    endOfLine
    string "clock tick"
    skipMany space
    string "channel mask"
    skipMany space
    events <- parseEvent `sepEndBy` newline
    return Record {
          _recordEventCount = counts
        , _recordEvents = events
    }

parseSegmentCount :: Parser Int
parseSegmentCount = do
    string "segments:"
    skipMany space
    n <- many digit
    return . read $ n

readListModeData :: String -> ListModeData
readListModeData l = 
    let txt = map unlines $ S.splitWhen ("" ==) $ lines l
        Right segments = simpleParse parseSegmentCount "<segment count>" (head txt)
        parsedRecords = rights $ map (simpleParse parseRecord "<record>") (tail txt)
     in ListModeData {
              _listModeDataSegmentCount = segments
            , _listModeRecords = parsedRecords
        }

sortPredicate :: Int -> Int -> (Int, Int) -> (Int, Int) -> Bool
sortPredicate offset base (a, _) (b, _) =
    quot (a - offset) base == quot (b - offset) base

type BinTally = M.Map Int Int

printTally :: BinTally -> String
printTally x = unlines $ map (\(x, y) -> show x ++ " " ++ show y) (M.toList x)

countEvents :: [Event] -> [(Int, Int)]
countEvents e =
    let eventTimes = map _eventClockTick e
        masks = map _eventMask e
        mcounts = fmap countMasks masks
    in zip eventTimes mcounts
    where c True = 1
          c False = 0
          countMasks (ChannelMask m) = sum $ fmap c m

updateCounts :: (Int, Int) -> BinTally -> BinTally
updateCounts = uncurry (M.insertWith (+))

combineTallies :: BinTally -> BinTally -> BinTally
combineTallies = M.unionWith (+)

histogram :: Ord a => [a] -> M.Map a Int
histogram xs = M.fromList [(head l, length l) | l <- group (sort xs)]

countBetween' :: Int -> [Int] -> Int
countBetween' tot [x, y] = tot + y - x - 1 
countBetween' tot (x:y:xs) = countBetween' (tot + y - x - 1) (y:xs)

countBetween :: [Int] -> Int
countBetween l = countBetween' (head l) l

sequentialBin :: Int -> [Event] -> BinTally
sequentialBin binWidth events =
    let e = countEvents events
        t0 = fst . head $ e
        cpred = sortPredicate t0 binWidth
        {-groupedCounts = (map . map) snd $ groupBy cpred e-}
        groupedCounts = groupBy cpred e
        ticks = map (\c -> quot ((fst . head) c - t0) binWidth) groupedCounts
        zeroBins = countBetween ticks
        redCounts = (map . map) snd $ groupedCounts
    in M.insert 0 zeroBins $ histogram (map sum redCounts)

data CliArgs = CliArgs { _cliFname :: String, _cliBinWidth :: Int }

cliParser :: ParserSpec CliArgs
cliParser = CliArgs
  `parsedBy` reqPos "filename" `Descr` "Name of file to read from"
  `andBy` reqPos "bin_width" `Descr` "Width of binning window"

cliInterface :: IO (CmdLnInterface CliArgs)
cliInterface =
    (`setAppDescr` "Read nPOD list mode file")
    <$> (`setAppEpilog` "Prints multiplicity bins to stdout in csv format")
    <$> mkApp cliParser

runTally :: CliArgs -> IO ()
runTally cli = do
    file <- readFile . _cliFname $ cli
    let records = readListModeData file
        ev = map _recordEvents . _listModeRecords $ records
        recordTallies = map (sequentialBin (_cliBinWidth cli)) ev
    putStr . printTally $ foldr combineTallies M.empty recordTallies 

main :: IO ()
main = cliInterface >>= flip runApp runTally