{-# LANGUAGE NoMonomorphismRestriction #-}
module PTS.Pretty 
  ( singleLine 
  , multiLine
  , showCtx ) where

import Control.Arrow(first)

import Parametric.Pretty

import Data.List(intersperse)

import PTS.AST

-- priorities
pAppR = 3
pApp  = 2
pArrL = 2
pArr  = 1
pLam  = 0
pPi   = 0
pIf0  = 2

-- finding chains of similiar terms
appChain = appChain' [] where
  appChain' xs (App f x) = appChain' (x:xs) f
  appChain' xs (Pos p t) = appChain' xs t
  appChain' xs t = (t, xs)

lamChain (Lam n q t) = first ((n, q) :) $ lamChain t
lamChain (Pos p t) = lamChain t
lamChain t = ([], t)

piChain (Pi n q t) | n `elem` freevars t = first ((n, q) :) $ piChain t
piChain (Pos p t) = piChain t
piChain t = ([], t)

arrChain (Pi n q t) | n `notElem` freevars t = first (q :) $ arrChain t
arrChain (Pos p t) = arrChain t
arrChain t = ([], t)

-- pretty printing
instance (Pretty c, Show c, Eq c) => Pretty (Term c) where
  pretty p (Nat n) = 
    int n
  
  pretty p (NatOp n _ a b) = parens `when` (pApp < p) $ 
    text n <+> pretty pAppR a <+> pretty pAppR b
    
  pretty p (IfZero c t e) = parens `when` (pIf0 < p) $ 
    sep [ text "if" <+> pretty pIf0 c
        , nest p $ text "then" <+> pretty pIf0 t
        , nest p $ text "else" <+> pretty pIf0 e ]
        
  pretty p (Var n) = 
    text n
    
  pretty p (Const c) = 
    pretty p c
    
  pretty p t@(App _ _) = parens `when` (pApp < p) $ 
    let (f, xs) = appChain t in
      (case f of Var n -> hsep
                 Const c -> hsep
                 t -> sep)
        [ pretty pApp f
        , nest 2 . sep . map (pretty pAppR) $ xs ]
    
  pretty p t@(Lam _ _ _) = parens `when` (pLam < p) $ 
    let (lams, body) = lamChain t in
      sep [ sep . map (\(n, q) -> text "lambda" <+> text n <+> text ":" <+> pretty pLam q <+> text ".") $ lams
          , nest 2 $ pretty pLam body ]
  
  pretty p t@(Pi n _ t') | n `elem` freevars t' = parens `when` (pPi < p) $
    let (pis, body) = piChain t in
     sep [ sep . map (\(n, q) -> text "Pi" <+> text n <+> text ":" <+> pretty pPi q <+> pretty p ".") $ pis
         , nest 2 $ pretty pPi body ]
  
  pretty p t@(Pi _ _ _) = parens `when` (pArr < p) $
    let (arrs, body) = arrChain t in
     sep $ [ sep . map (\q -> pretty pArrL q <+> pretty p "->") $ arrs] ++ [ pretty pArr body ]
  
  pretty p (Pos _ t) = pretty p t

instance (Pretty c, Eq c, Show c) => Pretty (Stmt c) where
  pretty p (Bind n t) = text n <+> text "=" <+> pretty 0 t
  pretty p (Term t) = pretty 0 t

instance (Pretty c, Eq c, Show c) => Show (Term c) where
  show t = singleLine t
  
showCtx :: (Show c, Eq c, Pretty c) => [(String, Term c)] -> String
showCtx = concat . intersperse ", " . map (\(n, t) -> n ++ " : " ++ show t) 
