{-# LANGUAGE DisambiguateRecordFields, NamedFieldPuns, RecordWildCards, PostfixOperators, LiberalTypeSynonyms, TypeOperators #-}

module DerivationTrees (
-- * Basics
module Data.Monoid,
module Data.LabeledTree,

-- * TeX extras
brace, paren, brack, tex, text, (!), TeX(..), (<>),

-- * Derivation' building
-- axiom, rule, etc, aborted, 
emptyDrv, haltDrv, haltDrv', abortDrv, delayPre, 
dummy, rule, Derivation, Premise, Rule(..),

-- * Figure building
Figure(..),

-- * Engine
compile

) where

import DerivationTrees.Basics
import Data.List
import Data.Traversable hiding (mapM)
import Control.Monad.Writer 
import Control.Monad.State 
import Control.Applicative 
import System.IO.Unsafe
import Data.LabeledTree
import Control.Applicative.State
import Data.Monoid

data Rule tag = Rule {tag :: tag, style :: LineStyle, delimiter :: TeX, label :: TeX, conclusion :: TeX}
  deriving Show


instance Monoid t => Applicative (Writer t) where
    pure = return
    (<*>) = ap

type Premise = Link ::> Derivation' ()
type Derivation' tag = Tree Link (Rule tag) 
type Derivation = Derivation' ()

data Figure tag = Figure {ident :: Int, contents :: Derivation' tag}

alloc = do
  x <- get
  put (x + 1)
  return x

------------------------------------------------------------
-- Phase 1: Detach

type Detach x = x -> Writer [Figure ()] x

detachP :: Detach Premise
detachP (Detached{..} ::> d) = do
  d'@(Node r ps) <- detachD d
  tell [Figure {contents = Node r {delimiter = label} ps,..}]
  return $ (defaultLink ::> haltDrv label d)
detachP (l ::> d) = (l ::>) <$> detachD d

detachD :: Detach Derivation
detachD (Node n ps) = Node n <$> for ps detachP

detachF :: Figure () -> Writer [Figure ()] ()
detachF Figure{..} = do
  contents <- detachD contents
  tell [Figure{..}]

detachTop :: [Figure ()] -> [Figure ()]
detachTop fs = snd $ runWriter $ for fs detachF

--------------------------------------------------
-- Phase 2: Delay

insertAt n x xs = take n xs ++ x : drop n xs

rm idx [] = []
rm 0 (x:xs) = xs
rm n (x:xs) = x : rm (n-1) xs

depth (Detached{} ::> _) = 2
depth (Link{steps} ::> Node _ ps) = 1 + steps + maximum (0 : map depth ps)

isDelayed :: Premise -> Bool
isDelayed (Delayed{} ::> _) = True
isDelayed _ = False

delayPre a s (Link {..} ::> j) = Link {steps = s, align = a, ..} ::> j

delayD :: Derivation -> Derivation
delayD (Node r ps0) = Node r (map delayP ps)
    where ps = fmap (fmap delayD) ps0
          ps' = filter (not . isDelayed) ps
          delayP (Delayed{..} ::> d) = Link{..} ::> d
             where steps = 1 + maximum (0 : map depth ps')
                   label = mempty
                   style = Dotted
          delayP p = p

delayF :: Figure () -> Figure ()
delayF (Figure{..}) = Figure{contents = delayD contents,..}

delayTop = map delayF



---------------------------------------------------------
-- Phase 3: Tag

type Tag x = x () -> State Int (x Int)

tagify :: Tag Rule
tagify (Rule {..}) = do
  tag <- alloc
  return $ Rule {..}

tagifyFig :: Tag Figure
tagifyFig (Figure {..}) = Figure ident <$> traverse tagify contents

tagifyTop :: [Figure ()] -> State Int [Figure Int]
tagifyTop = mapM tagifyFig

----------------------------------------------------------
-- Phase 4: Texify

fromTeX :: TeX -> String
fromTeX (TeX s) = "\""  ++ s ++ "\""

mkTuple :: [String] -> String
mkTuple l = " (" ++ intercalate "," l ++ ") "


link (Link {..} ::> Node (Rule{tag}) _) 
    | steps == 0 = show tag
    | otherwise = "MVD " ++ show tag ++ " " ++ 
                  mkTuple [show steps,show "",fromTeX label,show align,show (fromEnum style)]

type Stringize x = x Int -> Writer [String] ()

stringize :: Stringize Derivation'
stringize (Node Rule {tag = t, ..} premises) = do
  traverse (traverse stringize) premises
  tell ["jgm " ++ show t ++ " " ++ fromTeX conclusion ++ ";"]
  tell ["Nfr " ++ show t ++ mkTuple (map link premises)  ++ " " ++
                            mkTuple [fromTeX label,fromTeX delimiter,fromTeX (TeX ""),show (fromEnum style)] ++ ";"]

stringizeFig :: Figure Int -> Writer [String] () 
stringizeFig (Figure {..}) = do
  tell ["beginfig(" ++ show ident ++")"]
  stringize contents
  tell ["draw drv_tree;", 
        "endfig;"]
  
preamble = ["outputtemplate := \"%j-%c.mps\";",
            "input drv;",
              "verbatimtex %&latex"] ++ lines texPreamble ++ 
             [
              "\\begin{document}",
              "etex;",
              "prologues:=3;"
             ]

texPreamble = unsafePerformIO $ readFile "preamble.tex"

stringizeTop figs = do
  tell preamble
  mapM_ stringizeFig figs
  tell ["end"]

------------------------------------------
-- Pipeline

compile :: [Figure ()] -> [String]
compile j = snd $ runWriter $ stringizeTop $ fst $ flip runState 0 $ tagifyTop $ delayTop $ detachTop $ j

-----------------------

rule label conclusion = Rule {tag = (), delimiter = mempty, style = Simple, ..}

dummy :: Rule ()
dummy = (rule mempty mempty) {DerivationTrees.style = None}
emptyDrv = Node dummy []

abortDrv (Node Rule {..} _) = Node Rule {style = Waved, ..} []

-- | Used when the rest of the derivation is known.
haltDrv' :: TeX -> Derivation -> Derivation
haltDrv' tex (Node r _) = Node r {DerivationTrees.style = None} 
     [defaultLink {DerivationTrees.Basics.style = TeXDotted, steps = 1, DerivationTrees.Basics.label = tex} ::> emptyDrv]

-- | More compact variant
haltDrv :: TeX -> Derivation -> Derivation
haltDrv t (Node r _) = Node r [defaultLink ::> Node dummy {conclusion = tex "vdots" [] <> tex "hspace" [TeX "2pt"] <> t} []]


-------------------------
brace :: TeX -> TeX
brace (TeX x) = TeX $ '{' : x ++ ['}']
paren (TeX x) = TeX $ '(' : x ++ [')']
brack (TeX x) = TeX $ "[" ++ x ++ "]"

tex :: String -> [TeX] -> TeX
tex macro args = TeX ('\\' : macro) <> mconcat (map brace args)
text x = tex "text" [TeX x]

f ! x = brace f <> TeX "_" <> brace x