{-# LANGUAGE ScopedTypeVariables, TypeFamilies, FlexibleContexts, GADTs #-}
module Main where

import Control.Monad
import Control.Arrow
import Data.List
import Data.Maybe
import qualified Data.Map as M
import qualified Data.ByteString.Char8 as S
import qualified Data.ByteString.Lazy.Char8 as B
import Data.ByteString.Lex.Lazy.Double

import Text.Regex.PCRE
import Text.Regex.PCRE.ByteString

import System
import System.Exit
import System.Console.GetOpt

import Data.Time hiding (parseTime)
import Data.Time.Parse

import Data.Accessor

import Graphics.Rendering.Chart
import Graphics.Rendering.Chart.Gtk
import Graphics.Rendering.Chart.Grid
import Graphics.Rendering.Chart.Plot
import Graphics.Rendering.Chart.Event

import Data.Colour
import Data.Colour.Names

data Edge = Rise | Fall | Pulse deriving (Eq,Show)

data InEvent = InEdge  Edge
             | InValue Double
             | InAtom  S.ByteString

data OutFormat = PNG | PDF | PS | SVG | Window

class HasDelta t where
  type Delta t :: *
  add :: Delta t -> t -> t

instance HasDelta Double where
  type Delta Double = Double
  add d t = t + d

instance HasDelta LocalTime where
  type Delta LocalTime = NominalDiffTime
  add d t = utcToLocalTime utc (addUTCTime d (localTimeToUTC utc t))

instance Read NominalDiffTime where
  readsPrec n s = [(fromIntegral i, s') | (i,s') <- readsPrec n s]

class (Ord t, HasDelta t, PlotValue t, Show t, Show (Delta t), Read (Delta t)) => TimeAxis t 

instance TimeAxis Double

instance TimeAxis LocalTime

data (TimeAxis t) => ChartKind t = KindEvent
               | KindHistogram { binSize :: Delta t }
               | KindQuantile  { binSize :: Delta t, quantiles :: [Double] }
               | KindBinFreq   { binSize :: Delta t, delims    :: [Double] }
               | KindBinCount  { binSize :: Delta t, delims    :: [Double] }
               | KindFreq      { binSize :: Delta t, style :: PlotBarsStyle }
               | KindCount     { binSize :: Delta t, style :: PlotBarsStyle }
               | KindValue
               | KindNone


data Conf = forall t . TimeAxis t =>
  Conf {
    inFile        :: FilePath,
    parseTime     :: B.ByteString -> Maybe (t, B.ByteString),
    chartKindF    :: S.ByteString -> ChartKind t,

    fromTime      :: Maybe t,
    toTime        :: Maybe t,

    outFile       :: FilePath,
    outFormat     :: OutFormat,
    outResolution :: (Int,Int)
  }

readConf :: [String] -> Conf
readConf args = case (words $ single "time format" "-tf" ("date %Y-%m-%d %H:%M:%OS")) of
    ["num"]  -> readConf' (fmap (int2double *** id) . B.readInt)
    "date":f -> readConf' (strptime (B.pack $ unwords f))
    _        -> error "Unrecognized time format (-tf)"
  where
    int2double = fromIntegral::Int-> Double
    single desc name def = case (getArg name 1 args) of
      [[r]] -> r
      []    -> def
      _     -> error $ "Single argument expected for: "++desc++" ("++name++")"

    readConf' parseTime = Conf {inFile=inFile, outFile=outFile, outFormat=outFormat, outResolution=outRes,
                      chartKindF=chartKindF, parseTime=parseTime, fromTime=fromTime, toTime=toTime}
      where
        inFile      = single "input file"  "-if" (error "No input file (-if) specified")
        outFile     = single "output file" "-o"  (error "No output file (-of) specified")
        outFormat   = maybe PNG id $ lookup (single "output format" "-of" (name2format outFile)) $
            [("png",PNG), ("pdf",PDF), ("ps",PS), ("svg",SVG), ("x",Window)]
          where
            name2format = reverse . takeWhile (/='.') . reverse
        outRes      = parseRes $ single "output resolution" "-or" "640x480"
          where
            parseRes s = case break (=='x') s of (h,_:v) -> (read h,read v)
        chartKindF  = kindByRegex [((=~regex), parseKind (words kind)) | [regex,kind] <- getArg "-k" 2 args]
          where
            kindByRegex rks s = case [k | (p,k) <- rks, p s] of
              k:_ -> k
              _   -> defaultKind

        fromTime    = fst `fmap` (parseTime . B.pack $ single "minimum time (inclusive)" "-fromTime" "")
        toTime      = fst `fmap` (parseTime . B.pack $ single "maximum time (exclusive)" "-toTime"   "")

        parseKind ["hist",    n  ] = KindHistogram {binSize=read n}
        parseKind ["freq",    n  ] = KindFreq      {binSize=read n,style=BarsClustered}
        parseKind ["freq",    n,s] = KindFreq      {binSize=read n,style=parseStyle s}
        parseKind ["count",   n  ] = KindCount     {binSize=read n,style=BarsClustered}
        parseKind ["count",   n,s] = KindCount     {binSize=read n,style=parseStyle s}
        parseKind ["event"       ] = KindEvent
        parseKind ["quantile",b,q] = KindQuantile  {binSize=read b, quantiles=read ("["++q++"]")}
        parseKind ["binf",    b,q] = KindBinFreq   {binSize=read b, delims   =read ("["++q++"]")}
        parseKind ["binc",    b,q] = KindBinCount  {binSize=read b, delims   =read ("["++q++"]")}
        parseKind ["value"       ] = KindValue
        parseKind ["none"        ] = KindNone
        parseKind ws               = error ("Unknown diagram kind" ++ unwords ws)

        defaultKind = parseKind $ words $ single "default kind" "-dk" "event"

        parseStyle "stacked"   = BarsStacked
        parseStyle "clustered" = BarsClustered


-- getArg "-a" 2 ["-b", "1", "-a", "2", "q", "r", "-c", "3", "-a", "x"] =
-- [["2", "q"], ["x"]]
getArg :: String -> Int -> [String] -> [[String]]
getArg name arity args = [take arity as | (t:as) <- tails args, t==name]

readSource :: (Show t) => (B.ByteString -> Maybe (t,B.ByteString)) -> FilePath -> IO [(t, S.ByteString, InEvent)]
readSource readTime f = (justs . map parseLine . blines) `fmap` (if f=="-" then B.getContents else B.readFile f)
  where
    justs xs = [x | Just x <- xs]
    blines   = B.split '\n'
    strict   = S.concat . B.toChunks
    parseLine s = do
      (t, s') <- readTime s
      (_, s'') <- B.uncons s'
      (c,rest) <- B.uncons s''
      case c of
        '>' -> return (t, strict rest, InEdge Rise )
        '<' -> return (t, strict rest, InEdge Fall )
        '!' -> return (t, strict rest, InEdge Pulse)
        '=' -> do
          let (track, val') = B.break (==' ') rest
          (_,val) <- B.uncons val'
          if B.null val
            then Nothing
            else do
              case B.head val of
                '`' -> do
                  return (t, strict track, InAtom (strict $ B.tail val))
                _   -> do
                  (v,_  ) <- readDouble val
                  return (t, strict track, InValue v)
        _   -> Nothing

makeChart :: forall t . TimeAxis t =>
             (S.ByteString -> ChartKind t) -> [(t, S.ByteString, InEvent)] -> Renderable ()
makeChart chartKindF []                  = emptyRenderable
makeChart chartKindF events@((t0,_,_):_) = renderLayout1sStacked plots
  where
    track2events = reverse `fmap` foldl insert M.empty [(s, (t, e)) | (t, s, e) <- events]
      where insert m (s, r) = M.alter (Just . maybe [r] (r:)) s m

    plots          = [ plotTrack k es | (k, es) <- M.toList track2events, 
                                        case (chartKindF k) of {KindNone -> False ; _ -> True} ]

    timeBounds     = (head times, last times)
    times          = sort $ [t | tes <- M.elems track2events, (t,_)<- tes]
    commonTimeAxis = autoAxis times

    plotTrack      name es = case (chartKindF name) of
      KindHistogram bs    -> withAnyOrdinate $ plotTrackHistogram name es bs
      KindFreq      bs k  -> withAnyOrdinate $ plotTrackFreq      name es bs k
      KindCount     bs k  -> withAnyOrdinate $ plotTrackCount     name es bs k
      KindEvent           -> withAnyOrdinate $ plotTrackEvent     name es
      KindQuantile  bs qs -> withAnyOrdinate $ plotTrackQuantile  name es qs bs
      KindBinFreq   bs vs -> withAnyOrdinate $ plotTrackBinFreqs  name es vs bs
      KindBinCount  bs vs -> withAnyOrdinate $ plotTrackBinCounts name es vs bs
      KindValue           -> withAnyOrdinate $ plotTrackValue     name es
      KindNone               -> error "KindNone should not be plotted"

    edges  es = [(t,e) | (t,InEdge  e) <- es]
    values es = [(t,v) | (t,InValue v) <- es]
    atoms  es = [(t,a) | (t,InAtom  a) <- es]

    lag xs = xs `zip` tail xs

    ourPlotBars :: (BarsPlotValue a) => PlotBars t a
    ourPlotBars = plot_bars_spacing ^= BarsFixGap 0 0 $
                  plot_bars_style   ^= BarsStacked    $
                  plot_bars_alignment ^= BarsLeft     $
                  defaultPlotBars

    plotTrackHistogram name es bs = layoutWithTitle (plotBars plot) name
      where plot = plot_bars_values  ^= barsData $ ourPlotBars
            barsData = [(t,[n]) | ((t,_),n) <- edges2bins bs t0 (edges es)]

    plotTrackFreq  = plotTrackAtoms atoms2freqs
    plotTrackCount = plotTrackAtoms atoms2counts

    plotTrackAtoms f name es bs k = layoutWithTitle (plotBars plot) name
      where plot = plot_bars_style       ^= k           $
                   plot_bars_values      ^= vals        $
                   plot_bars_item_styles ^= itemStyles  $
                   plot_bars_titles      ^= "":map show vs $
                   ourPlotBars
            itemStyles = none:[(solidFillStyle (opaque c), Nothing) | c <- colors]
            vals = byTimeBins ((0:).f vs) bs t0 as
            as   = atoms es
            vs   = M.keys $ M.fromList $ [(a,()) | (_,a) <- as]

    plotTrackEvent     name es       = layoutWithTitle (toPlot plot) name
      where plot = plot_event_data ^= edges2events (edges es) $ defaultPlotEvent

    plotTrackQuantile  name es qs bs = layoutWithTitle (plotBars plot) name
      where plot = plot_bars_values  ^= toBars (byTimeBins (getQuantiles qs) bs t0 (values es)) $
                   plot_bars_item_styles ^= quantileStyles $
                   plot_bars_titles  ^= quantileTitles $
                   ourPlotBars
            quantileStyles = none:(zip (map (solidFillStyle . opaque) colors) [Just $ solidLine 1 (opaque black) | i <- [0..n+1]])
            quantileTitles = [""]++[show p1++".."++show p2++"%" | (p1,p2) <- lag percents ]
              where percents = map (floor . (*100)) $ [0]++qs++[1]
            n = length qs
    
    colors = cycle [green,blue,yellow,red,orange,brown,grey,purple,violet,lightblue]

    binTitles vs = [low]++[show v1++".."++show v2 | (v1,v2) <- lag vs]++[high]
      where
        low = "<"++show (head vs)
        high = ">"++show (last vs)

    binColor n i = opaque (colors !! i)

    plotTrackBinFreqs  name es vs bs = plotTrackBars vals (binTitles vs) name (binColor n)
      where
        vals = byTimeBins ((0:).values2binFreqs  vs) bs t0 (values es)
        n    = length vs
    plotTrackBinCounts name es vs bs = plotTrackBars vals (binTitles vs) name (binColor n)
      where
        vals = byTimeBins ((0:).values2binCounts vs) bs t0 (values es)
        n    = length vs

    plotTrackBars :: (BarsPlotValue a) => [(t,[a])] -> [String] -> S.ByteString -> (Int -> AlphaColour Double) -> Layout1 t a
    plotTrackBars values titles name clr = layoutWithTitle (plotBars plot) name
      where plot = plot_bars_values      ^= values    $
                   plot_bars_item_styles ^= binStyles $
                   plot_bars_titles      ^= "":titles $
                   ourPlotBars
            binStyles = none:[(solidFillStyle (clr i), Just $ solidLine 1 (opaque black))
                             | (i,_) <- [0..]`zip`titles]

    none = (solidFillStyle transparent, Nothing)
    toBars tvs = [(t,diffs vs) | (t,vs) <- tvs]
    diffs xs = zipWith (-) xs (0:xs)

    plotTrackValue name es = layoutWithTitle (toPlot plot) name
      where plot = plot_lines_values ^= [values es] $ defaultPlotLines

    layoutWithTitle plot name =
        layout1_title ^= "" $
        layout1_plots ^= [Left plot] $
        layout1_bottom_axis .> laxis_generate ^= (\_ -> commonTimeAxis) $
        layout1_top_axis    .> laxis_generate ^= (\_ -> commonTimeAxis) $
        layout1_left_axis   .> laxis_title ^= S.unpack name $
        layout1_margin ^= 0 $
        defaultLayout1

edges2events :: (Ord t) => [(t,Edge)] -> [Event t ()]
edges2events tes = longs `merge` pulses
  where
    merge [] ps = ps
    merge ls [] = ls
    merge (l@(LongEvent t1 t2 _):ls) (p@(PulseEvent t _):ps)
      | t1<t = l:merge ls (p:ps)
      | True = p:merge (l:ls) ps
    pulses = [PulseEvent t () | (t,Pulse) <- tes]
    edges  = [(t,e) | (t,e) <- tes, e /= Pulse]
    longs  = longs' Nothing 0 (error "Unreachable") edges
      where
        longs' _         0 _ [] = []
        longs' (Just t0) _ t [] = [LongEvent t0 t ()]
        longs' Nothing   n _ ((t,Rise):tes) = longs' (Just t)  (n+1) t tes
        longs' Nothing   n _ ((t,Fall):tes) = longs' Nothing   0     t tes
        longs' (Just t0) n _ ((t,Rise):tes) = longs' (Just t0) (n+1) t tes
        longs' (Just t0) 1 _ ((t,Fall):tes) = LongEvent t0 t () : longs' Nothing 0 t tes
        longs' (Just t0) n _ ((t,Fall):tes) = longs' (Just t0) (n-1) t tes

edges2bins :: (Ord t,HasDelta t) => Delta t -> t -> [(t,Edge)] -> [((t,t), Int)]
edges2bins binSize t0 es = gather 0 0 0 es $ iterate (add binSize) t0
  where
    gather :: (Ord t) => Int -> Int -> Int -> [(t,Edge)] -> [t] -> [((t,t), Int)]
    gather 0 _ _ [] (t1:t2:ts) = []
    gather n _ _ [] (t1:t2:ts) = [((t1,t2),n)]
    gather nmax nopen npulse ((t,e):tes) (t1:t2:ts)
      | t<t1 = error "Times are not in ascending order"
      | t>=t2 = ((t1,t2),nmax):gather nopen nopen 0 ((t,e):tes) (t2:ts)
    gather nmax nopen npulse ((t,Rise ):tes) (t1:t2:ts)
      = gather (nmax `max` (nopen+npulse+1)) (nopen+1) npulse     tes (t1:t2:ts)
    gather nmax nopen npulse ((t,Fall ):tes) (t1:t2:ts)
      = gather nmax                          (nopen-1) npulse     tes (t1:t2:ts)
    gather nmax nopen npulse ((t,Pulse):tes) (t1:t2:ts)
      = gather (nmax `max` (nopen+npulse+1)) nopen    (npulse+1)  tes (t1:t2:ts)

values2timeBins :: (Ord t) => [t] -> [(t,a)] -> [[a]]
values2timeBins (t1:t2:ts) []          = []
values2timeBins (t1:t2:ts) tvs@((t,_):_)
  | t<t1 = error "Times are not in ascending order"
  | True = let (bin,rest) = span ((<t2).fst) tvs
           in (map snd bin : values2timeBins (t2:ts) rest)

byTimeBins :: (Ord t, HasDelta t, Ord a) => ([a] -> b) -> Delta t -> t -> [(t,a)] -> [(t, b)]
byTimeBins f binSize t0 tvs = times `zip` map f (values2timeBins times tvs)
  where times = iterate (add binSize) t0

getQuantiles :: (Ord a) => [Double] -> [a] -> [a]
getQuantiles qs = \xs -> quantiles' (sort xs)
  where
    qs' = sort qs
    quantiles' [] = []
    quantiles' xs = index (0:ns++[n-1]) 0 xs
      where
        n  = length xs
        ns = map (floor . (*(fromIntegral n-1))) qs'

        index _         _ []     = []
        index []        _ _      = []
        index [i]       j (x:xs)
          | i<j  = []
          | i==j  = [x]
          | True = index [i] (j+1) xs
        index (i:i':is) j (x:xs)
          | i<j  = index (i':is)   j     (x:xs)
          | i>j  = index (i:i':is) (j+1) xs
          | i==i' = x:index (i':is) j     (x:xs)
          | True = x:index (i':is) (j+1) xs

values2binFreqs :: (Ord a) => [a] -> [a] -> [Double]
values2binFreqs bins xs = map toFreq $ values2binCounts bins xs
  where
    n = length xs
    toFreq = if n==0 then const 0 else (\k -> fromIntegral k/fromIntegral n)
values2binCounts bins xs = values2binCounts' bins $ sort xs
  where
    values2binCounts' []     xs = [length xs]
    values2binCounts' (a:as) xs = length xs0 : values2binCounts' as xs'
      where (xs0,xs') = span (<a) xs

atoms2counts :: (Ord a) => [a] -> [a] -> [Int]
atoms2counts as xs = map (maybe 0 id . (`M.lookup` m)) as
  where
    m          = foldl' insert M.empty xs
    insert m a = M.alter (Just . maybe 1 inc) a m
    inc    n   = n `seq` (n+1)

atoms2freqs :: (Ord a) => [a] -> [a] -> [Double]
atoms2freqs as xs = map toFreq (atoms2counts as xs)
  where
    n = length xs
    toFreq = if n==0 then const 0 else (\k -> fromIntegral k/fromIntegral n)

zoom :: (TimeAxis t) => [(t, S.ByteString, InEvent)] -> Maybe t -> Maybe t -> [(t, S.ByteString, InEvent)]
zoom events fromTime toTime = filter p events
  where 
    p (t, _, _) = (maybe True (\ft -> t >= ft) fromTime) &&
                  (maybe True (\tt -> t <  tt) toTime)

showHelp = mapM_ putStrLn [ "",
  "tplot - a tool for drawing timing diagrams. See http://www.haskell.org/haskellwiki/Timeplot",
  "Usage: tplot [-o OFILE] [-of {png|pdf|ps|svg|x}] [-or 640x480] -if IFILE [-tf TF] ",
  "             [-k Pat1 Kind1 -k Pat2 Kind2 ...] [-dk KindN] [-fromTime TIME] [-toTime TIME]",
  "  -o  OFILE - output file (required if -of is not x)",
  "  -of       - output format (x means draw result in a window, default: extension of -o)",
  "  -or       - output resolution (default 640x480)",
  "  -if IFILE - input file; '-' means 'read from stdin'",
  "  -tf TF    - time format: 'num' means that times are integer numbers less than 2^31",
  "              (for instance, line numbers); 'date PATTERN' means that times are dates",
  "              in the format specified by PATTERN - see http://linux.die.net/man/3/strptime,",
  "              for example, [%Y-%m-%d %H:%M:%S] parses dates like [2009-10-20 16:52:43]. ",
  "              We also support %OS for fractional seconds (i.e. %OS will parse 12.4039).",
  "              Default: 'date %Y-%m-%d %H:%M:%OS'",
  "  -k P K    - set diagram kind for tracks matching pattern P to K ",
  "              (-k clauses are matched till first success)",
  "  -dk       - set default diagram kind",
  "  -fromTime - filter records whose time is >= this time (formatted according to -tf)",
  "  -toTime   - filter records whose time is <  this time (formatted according to -tf)",
  "",
  "Input format: lines of the following form:",
  "1234 >A - at time 1234, during event A has begun",
  "1234 <A - at time 1234, during event A has ended",
  "1234 !B - at time 1234, pulse event B has occured",
  "1234 =C VAL - at time 1234, parameter C had numeric value VAL (for example, HTTP response time)",
  "1234 =D `EVENT - at time 1234, event EVENT occured in process D (for example, HTTP response code)",
  "It is assumed that many events of the same kind may occur at once.",
  "Diagram kinds:",
  "  'event' is for event diagrams: during events are drawn like --[===]--- , pulse events like --|--",
  "  'hist N' is for histograms: a histogram is drawn with granularity of N time units, where",
  "     the bin corresponding to [t..t+N) has value 'what was the maximal number of active events",
  "     in that interval'.",
  "  'freq N [TYPE]' is for event frequency histograms: a histogram of type TYPE (stacked or ",
  "     clustered, default clustered) is drawn for each time bin of size N, about the distribution ",
  "     of various ` events",
  "  'count N [TYPE]' is for event frequency histograms: a histogram of type TYPE (stacked or ",
  "     clustered, default clustered) is drawn for each time bin of size N, about the counts of ",
  "     various ` events",
  "  'quantile N q1,q2,..' (example: quantile 100 0.25,0.5,0.75) - a bar chart of corresponding",
  "     quantiles in time bins of size N",
  "  'binf N v1,v2,..' (example: binf 100 1,2,5,10) - a bar chart of frequency of values falling",
  "     into bins min..v1, v1..v2, .., v2..max in time bins of size N",
  "  'binc N v1,v2,..' (example: binf 100 1,2,5,10) - a bar chart of counts of values falling",
  "     into bins min..v1, v1..v2, .., v2..max in time bins of size N",
  "  'value' - a simple line plot of numeric values",
  "N is measured in units or in seconds."
  ]


main = do
  args <- getArgs
  mainWithArgs args
mainWithArgs args = do
  when (null args || args == ["--help"]) $ showHelp >> exitSuccess
  let conf = readConf args
  let render = case (outFormat conf) of {
      PNG    -> renderableToPNGFile ;
      PDF    -> renderableToPDFFile ;
      PS     -> renderableToPSFile  ;
      SVG    -> renderableToSVGFile ;
      Window -> \c w h f -> renderableToWindow c w h
    }
  case conf of
    Conf{parseTime=parseTime, inFile=inFile, chartKindF=chartKindF, outFile=outFile, outResolution=outResolution, fromTime=fromTime, toTime=toTime} -> do
      source <- readSource parseTime inFile
      let source' = zoom source fromTime toTime
      let chart = makeChart chartKindF source'
      let (w,h) = outResolution
      render chart w h outFile