content: {-# LANGUAGE ViewPatterns, OverloadedStrings #-}
  
  module LaTeXBase
   ( MathType(..), LaTeXUnit(..), LaTeX, TeXArg, ArgKind(..), concatRaws, hasCommand, isJustRaw
   , matchCommand, lookForCommand, matchEnv, mapTeX, mapCommandName, renderLaTeX, mapTeXRaw, isTeXEnv, texSpan, unconsRaw
   , trim, trimr, triml, texStripInfix, isCodeblock, isMath, texStripPrefix, texStripAnyPrefix, AllUnits(..) ) where
  
  import Data.String (fromString)
  import Prelude hiding ((.), (++), writeFile, dropWhile)
  import Data.Text (Text, pack)
  import qualified Data.Text as Text
  import Data.Char (isSpace)
  import Util ((.), (++), textStripInfix)
  import Control.Arrow (first, second)
  
  data MathType = Parentheses | Square | Dollar
  	deriving (Eq, Show, Ord)
  
  data ArgKind = FixArg | OptArg
  	deriving (Eq, Show, Ord)
  
  type TeXArg = (ArgKind, LaTeX)
  
  data LaTeXUnit
  	= TeXRaw Text
  	| TeXComm String String [TeXArg] -- first string is command name, second is trailing whitespace
  	| TeXEnv String [TeXArg] LaTeX
  	| TeXMath MathType LaTeX
  	| TeXLineBreak
  	| TeXBraces LaTeX
  	deriving (Eq, Show, Ord)
  
  isTeXEnv :: String -> LaTeXUnit -> Bool
  isTeXEnv x (TeXEnv y _ _) = x == y
  isTeXEnv _ _ = False
  
  type LaTeX = [LaTeXUnit]
  
  lookForCommand :: String -> LaTeX -> [[TeXArg]]
  lookForCommand n = (snd .) . matchCommand (n ==)
  
  class AllUnits a where
  	allUnits :: a -> [LaTeXUnit]
  
  instance AllUnits LaTeXUnit where
  	allUnits u = u : case u of
  		TeXMath _ l -> allUnits l
  		TeXBraces l -> allUnits l
  		TeXComm _ _ a -> (snd . a) >>= allUnits
  		TeXEnv _ a l -> (l : snd . a) >>= allUnits
  		_ -> []
  
  instance AllUnits a => AllUnits [a] where
  	allUnits = concatMap allUnits
  
  matchCommand :: AllUnits a => (String -> Bool) -> a -> [(String, [TeXArg])]
  matchCommand f x = [(str, as) | TeXComm str _ as <- allUnits x, f str]
  
  hasCommand :: (String -> Bool) -> LaTeX -> Bool
  hasCommand f = not . null . matchCommand f
  
  matchEnv :: AllUnits a => (String -> Bool) -> a -> [(String, [TeXArg], LaTeX)]
  matchEnv f x = [(str, as, l) | TeXEnv str as l <- allUnits x, f str]
  
  mapTeX :: (LaTeXUnit -> Maybe LaTeX) -> LaTeX -> LaTeX
  mapTeX f = concatMap g
  	where
  		g :: LaTeXUnit -> LaTeX
  		g (f -> Just x) = x
  		g (TeXComm c ws a) = [TeXComm c ws (h . a)]
  		g (TeXBraces x) = [TeXBraces (mapTeX f x)]
  		g (TeXMath t b) = [TeXMath t (mapTeX f b)]
  		g (TeXEnv n a b) = [TeXEnv n (h . a) (mapTeX f b)]
  		g x = [x]
  		h = second (mapTeX f)
  
  mapCommandName :: (String -> String) -> LaTeX -> LaTeX
  mapCommandName f = concatMap g
  	where
  		g :: LaTeXUnit -> LaTeX
  		g (TeXComm c ws a) = [TeXComm (f c) ws (h . a)]
  		g (TeXBraces x) = [TeXBraces (mapCommandName f x)]
  		g (TeXMath t b) = [TeXMath t (mapCommandName f b)]
  		g (TeXEnv n a b) = [TeXEnv n (h . a) (mapCommandName f b)]
  		g x = [x]
  		h = second (mapCommandName f)
  
  renderLaTeX :: LaTeX -> Text
  renderLaTeX = mconcat . (renderUnit .)
  
  renderUnit :: LaTeXUnit -> Text
  renderUnit (TeXRaw t) = t
  renderUnit (TeXComm "right" _ [(FixArg, [TeXRaw "."])]) = "\\right."
  renderUnit (TeXComm name ws [])
  	| name `elem` ["left", "sum", "int", "sin", "cos", "right", "bigl", "bigr", "big", "small", "smaller"] = pack $ "\\" <> name <> ws
  	| otherwise = "\\" <> fromString name <> "{}"
  renderUnit (TeXComm name ws args) = "\\" <> pack (fromString name) <> pack (fromString ws) <> renderArgs args
  renderUnit (TeXEnv name args c) =
  	"\\begin{" <> fromString name <> "}"
  	<> renderArgs args
  	<> renderLaTeX c
  	<> "\\end{" <> fromString name <> "}"
  renderUnit (TeXMath Dollar l) = "$" <> renderLaTeX l <> "$"
  renderUnit (TeXMath Square l) = "\\[" <> renderLaTeX l <> "\\]"
  renderUnit (TeXMath Parentheses l) = "\\(" <> renderLaTeX l <> "\\)"
  renderUnit TeXLineBreak = "\\\\"
  renderUnit (TeXBraces l) = "{" <> renderLaTeX l <> "}"
  
  renderArgs :: [TeXArg] -> Text
  renderArgs = mconcat . (renderArg .)
  
  renderArg :: TeXArg -> Text
  renderArg (FixArg, l) = "{" <> renderLaTeX l <> "}"
  renderArg (OptArg, l) = "[" <> renderLaTeX l <> "]"
  
  mapTeXRaw :: (Text -> LaTeXUnit) -> (LaTeX -> LaTeX)
  mapTeXRaw f = map go
  	where
  		go :: LaTeXUnit -> LaTeXUnit
  		go (TeXRaw t) = f t
  		go (TeXComm s ws args) = TeXComm s ws (second (go .) . args)
  		go (TeXEnv s args body) = TeXEnv s (second (go .) . args) (go . body)
  		go (TeXBraces l) = TeXBraces $ go . l
  		go t@(TeXMath _ _) = t
  		go t@TeXLineBreak = t
  
  concatRaws :: LaTeX -> LaTeX
  concatRaws (TeXRaw a : TeXRaw b : more) = concatRaws (TeXRaw (a ++ b) : more)
  concatRaws (TeXComm s ws args : more) = TeXComm s ws (second concatRaws . args) : concatRaws more
  concatRaws (TeXEnv s args bd : more) = TeXEnv s (second concatRaws . args) (concatRaws bd) : concatRaws more
  concatRaws (TeXBraces x : more) = TeXBraces (concatRaws x) : concatRaws more
  concatRaws (x : more) = x : concatRaws more
  concatRaws [] = []
  
  unconsRaw :: LaTeX -> (Text, LaTeX)
  unconsRaw (TeXRaw x : y) = first (x ++) (unconsRaw y)
  unconsRaw x = ("", x)
  
  texStripPrefix :: Text -> LaTeX -> Maybe LaTeX
  texStripPrefix t (TeXRaw s : y) = case Text.stripPrefix t s of
  	Just "" -> Just y
  	Just s' -> Just (TeXRaw s' : y)
  	Nothing -> Nothing
  texStripPrefix _ _ = Nothing
  
  texStripAnyPrefix :: [Text] -> LaTeX -> Maybe (Text, LaTeX)
  texStripAnyPrefix [] _ = Nothing
  texStripAnyPrefix (x:y) z
      | Just a <- texStripPrefix x z = Just (x, a)
      | otherwise = texStripAnyPrefix y z
  
  texStripInfix :: Text -> LaTeX -> Maybe (LaTeX, LaTeX)
  texStripInfix t = go
  	where
  		go [] = Nothing
  		go (x : rest)
  			| TeXRaw s <- x
  			, Just (y, z) <- textStripInfix t s
  				= Just (h y, h z ++ rest)
  			| otherwise = first (x :) . go rest
  		h "" = []
  		h x = [TeXRaw x]
  
  texSpan :: (Char -> Bool) -> LaTeX -> (Text, LaTeX)
  texSpan p (TeXRaw x : y) = case Text.span p x of
      (stuff, "") ->  first (stuff ++) (texSpan p y)
      (stuff, rest) -> (stuff, TeXRaw rest : y)
  texSpan _ x = ("", x)
  
  invisible :: LaTeXUnit -> Bool
  invisible (TeXComm "index" _ _) = True
  invisible _ = False
  
  dropWhileEnd :: (Char -> Bool) -> LaTeX -> LaTeX
  dropWhileEnd _ [] = []
  dropWhileEnd p x
  	| invisible (last x) = dropWhileEnd p (init x) ++ [last x]
  	| TeXRaw y <- last x = init x ++ case Text.dropWhileEnd p y of
  		"" -> []
  		a -> [TeXRaw a]
  	| otherwise = x
  
  trimr, trim :: LaTeX -> LaTeX
  trimr = dropWhileEnd isSpace
  trim = triml . trimr
  
  triml :: LaTeX -> LaTeX
  triml (TeXRaw x : y) = case Text.dropWhile isSpace x of
  	"" -> triml y
  	x' -> TeXRaw x' : y
  triml x = x
  
  isMath :: LaTeXUnit -> Bool
  isMath (TeXMath _ _) = True
  isMath (TeXComm "ensuremath" _ _) = True
  isMath (TeXEnv "eqnarray*" _ _) = True
  isMath (TeXEnv "equation*" _ _) = True
  isMath _ = False
  
  isCodeblock :: LaTeXUnit -> Bool
  isCodeblock (TeXEnv "codeblock" _ _) = True
  isCodeblock (TeXEnv "indexedcodeblock" _ _) = True
  isCodeblock (TeXEnv "codeblocktu" _ _) = True
  isCodeblock (TeXEnv "codeblockdigitsep" _ _) = True
  isCodeblock _ = False
  
  isJustRaw :: LaTeX -> Maybe Text
  isJustRaw [TeXRaw x] = Just x
  isJustRaw _ = Nothing
  