{-# LANGUAGE UndecidableInstances, OverloadedStrings, NoMonomorphismRestriction, DeriveGeneric, ConstraintKinds #-}
module Data.CaPriCon(
  -- * Expression nodes
  IsCapriconString(..),BindType(..),Term(..),ApHead(..),Application(..),ContextTerm(..),Env,UniverseConstraints(..),DependentLogic(..),
  -- ** Managing De Bruijin indices
  adjust_depth,adjust_telescope_depth,inc_depth,map_univs,free_vars,is_free_in,
  -- ** Expression directories
  StringPattern,NodeDir(..),AHDir(..),ApDir,
  findPattern,freshContext,
  -- * Showing nodes
  ListBuilder(..),NodeDoc(..),doc2raw,doc2latex,doc2svg,showNode,showNode'
  ) where

import Definitive
import Language.Format
import GHC.Generics (Generic)

sequence2 :: (Applicative f,Traversable t,Traversable u) => t (u (f a)) -> f (t (u a))
sequence2 = sequence.map sequence
sequence3 :: (Applicative f,Traversable t,Traversable u,Traversable v) => t (u (v (f a))) -> f (t (u (v a)))
sequence3 = sequence.map sequence2

type FreeMap f a = Cofree f (Maybe a)
instance (Semigroup a, Semigroup (Coforest f a)) => Semigroup (Cofree f a) where
  Step a b + Step a' b' = Step (a+a') (b+b')
instance (Monoid a, Monoid (Coforest f a)) => Monoid (Cofree f a) where
  zero = Step zero zero
instance (DataMap (f (FreeMap f a)) k (FreeMap f a)) => DataMap (FreeMap f a) [k] a where
  at [] = lens (\(Step a _) -> a) (\(Step _ m) a -> Step a m)
  at (h:t) = coforest.at h.l'Just zero.at t
coforest :: Lens' (Cofree f a) (Coforest f a)
coforest = lens (\(Step _ m) -> m) (\(Step a _) m -> Step a m)

class (Ord str,Show str,Monoid str,Sequence str,IsString str) => IsCapriconString str where
  toString :: str -> String
instance IsCapriconString String where
  toString = id

type ListStream = [Word8]
newtype ListBuilder = ListBuilder (ListStream -> ListStream)
instance Semigroup ListBuilder where ListBuilder a + ListBuilder b = ListBuilder (a . b)
instance Monoid ListBuilder where zero = ListBuilder id
instance SerialStreamType ListStream where
  type StreamBuilder ListStream = ListBuilder
instance SerialStream ListStream where
  encodeByte _ b = ListBuilder (b:)
  toSerialStream (ListBuilder k) = k []

-- | Inductive types
type UniverseSize = Int
type SymbolRef = Int
data BindType = Lambda | Prod
              deriving (Show,Eq,Ord,Generic)
data Term str a = Bind BindType str (TypeTerm str a) (Term str a)
                | Cons (Application str a)
                | Universe UniverseSize
          deriving (Show,Generic)
type TypeTerm str a = Term str a
data ApHead str a = Sym SymbolRef
                  | Mu [(str,TypeTerm str a,TypeTerm str a)] [Term str a] (Application str a)
                  | Axiom (Term str a) a
            deriving (Show,Generic)
data Application str a = Ap (ApHead str a) [Term str a]
                 deriving (Show,Generic) 

type ListSerializable a = (Serializable ListStream a)
type ListFormat a = (Format ListStream a)
instance ListSerializable BindType
instance ListFormat BindType
instance (ListSerializable a,ListSerializable str) => ListSerializable (Term str a)
instance (ListFormat a,ListFormat str) => ListFormat (Term str a)
instance (ListSerializable a,ListSerializable str) => ListSerializable (ApHead str a)
instance (ListFormat a,ListFormat str) => ListFormat (ApHead str a)
instance (ListSerializable a,ListSerializable str) => ListSerializable (Application str a)
instance (ListFormat a,ListFormat str) => ListFormat (Application str a)

type Env str e = [(str,Binding e)]
class Monad m => DependentLogic str m e | e -> str where
  type Axiom e :: *
  type Binding e :: *
  
  mkUniverse :: UniverseSize -> m e
  mkVariable :: str -> m e
  mkBind :: BindType -> e -> m e
  mkApply :: e -> e -> m e
  mkMu :: e -> m e
  checkType :: e -> m e
  conversionDelta :: e -> e -> m (UniverseSize,UniverseSize)

  substHyp :: str -> e -> m (e -> e,Env str e)
  pullTerm :: Maybe str -> e -> m e
  insertHypBefore :: Maybe str -> str -> e -> m (e -> e,Env str e)

hypIndex :: (IsCapriconString str,MonadReader [(str,a)] m) => str -> MaybeT m Int
hypIndex h = ask >>= \l -> case [i | (i,x) <- zip [0..] l, fst x==h] of
  i:_ -> return i
  _ -> zero

instance (Show a,IsCapriconString str,Monad m,MonadReader (Env str (Term str a)) m) => DependentLogic str (MaybeT m) (Term str a) where
  type Axiom (Term str a) = a
  type Binding (Term str a) = TypeTerm str a

  mkUniverse = pure . Universe
  mkVariable v = hypIndex v <&> \i -> Cons (Ap (Sym i) [])
  mkBind b e = ask >>= \case
    (x,tx):_ -> pure $ Bind b x tx e
    _ -> zero
  mkApply f x = return (subst f (Cons (Ap (Sym 0) [inc_depth 1 x])))
  mkMu e = do
    te <- checkType e
    mte <- mu_type te^.maybeT
    let args (Bind Prod _ tx e') = tx:args e'
        args _ = []
    return (subst e (Cons (Ap (Mu [] (args mte) (Ap (Sym 0) [])) [])))
  checkType e = type_of e^.maybeT
  conversionDelta a b = return (convertDelta a b)^.maybeT

  substHyp h x = do
    i <- hypIndex h
    lift $ do
      ctx <- ask
      return (substn x i,let (ch,ct) = splitAt i ctx in zipWith (\j -> second $ substn (inc_depth (negate (1+j)) x) (i-j-1)) [0..] ch+drop 1 ct)
  pullTerm _ = return
  insertHypBefore Nothing h th = lift $ do
    ctx <- ask
    return (inc_depth 1,(h,th):ctx)
  insertHypBefore (Just h) h' th' = do
    hi <- hypIndex h
    lift $ do
      ctx <- ask
      let adj i j = if i+j>=hi then j+1 else j
      return (
        adjust_depth (adj (-1)),
        foldr (\x k i -> case compare hi i of
                           LT -> x:k (i+1)
                           EQ -> second (adjust_depth (adj i)) x:(h',inc_depth (negate (hi+1)) th'):k (i+1)
                           GT -> second (adjust_depth (adj i)) x:k (i+1))
          (\_ -> []) ctx 0)


newtype UniverseConstraints = UniverseConstraints [Maybe UniverseSize]
                            deriving (Generic,Show)
instance Format [Word8] UniverseConstraints
instance Serializable [Word8] UniverseConstraints
instance Semigroup UniverseConstraints where
  UniverseConstraints a + UniverseConstraints b = UniverseConstraints (zipWith (\x y -> zipWith max x y + x + y) a b) 
instance Monoid UniverseConstraints where
  zero = UniverseConstraints (repeat Nothing)
singleConstraint :: Int -> UniverseSize -> UniverseConstraints
singleConstraint h u = UniverseConstraints $ take h (repeat Nothing)+(Just u:repeat Nothing)
shiftConstraints :: Int -> UniverseConstraints -> UniverseConstraints
shiftConstraints n (UniverseConstraints x) | n < 0 = UniverseConstraints (drop (-n) x)
                                           | otherwise = UniverseConstraints (take n (repeat Nothing) + x)


data ContextTerm str a = ContextTerm SymbolRef UniverseConstraints (Term str a)
                       deriving (Show,Generic)
instance (ListSerializable a,ListSerializable str) => ListSerializable (ContextTerm str a)
instance (ListFormat a,ListFormat str) => ListFormat (ContextTerm str a)
restrictEnv :: SymbolRef -> [a] -> [a]
restrictEnv n e = drop (length e-n) e

instance (Show a,IsCapriconString str,MonadReader (Env str (Term str a)) m,Monad m) => DependentLogic str (MaybeT m) (ContextTerm str a) where
  type Axiom (ContextTerm str a) = a
  type Binding (ContextTerm str a) = Term str a

  mkUniverse u = ask >>= \ctx -> ContextTerm (length ctx) zero<$>mkUniverse u
  mkVariable i = local (dropWhile ((/=i) . fst)) (ask >>= \ctx -> ContextTerm (length ctx) (singleConstraint 0 0)<$>mkVariable i)
  mkBind t ce@(ContextTerm de uc e) | de>0 = ContextTerm (de-1) (shiftConstraints (-1) uc) <$> local (restrictEnv de) (mkBind t e)
                                    | otherwise = return ce
  mkApply (ContextTerm df ucf f) (ContextTerm dx ucx x) = do
    let dm = max df dx
    ContextTerm dm (ucf+ucx) <$> mkApply (inc_depth (dm-df) f) (inc_depth (dm-dx) x)
  mkMu (ContextTerm d uc e) = ContextTerm d uc <$> local (restrictEnv d) (mkMu e)
  checkType (ContextTerm d uc e) = ContextTerm d uc <$> local (restrictEnv d) (checkType e)
  conversionDelta (ContextTerm da _ a) (ContextTerm db _ b) =
    let dm = max da db in
      local (restrictEnv dm)
      $ conversionDelta (inc_depth (dm-da) a) (inc_depth (dm-db) b)
  
  pullTerm Nothing (ContextTerm d uc e) = ask <&> \l -> ContextTerm (length l) (shiftConstraints (length l - d) uc) (inc_depth (length l-d) e)
  pullTerm (Just v) (ContextTerm d uc e) = do
    nctx <- length <$> ask
    i <- hypIndex v
    let d' = nctx-(i+1)
    guard (d'>=d || all (\j -> d'+j >= d) (free_vars e))
    return (ContextTerm d' (shiftConstraints (d'-d) uc) $ inc_depth (d'-d) e)

  substHyp h vh = do
    ContextTerm dh uc vh' <- pullTerm (Just h) vh
    dm <- length <$> ask
    first (\f cv@(ContextTerm d uc' v) ->
              if d <= dh then cv
              else ContextTerm (d-1) (uc+uc') (inc_depth (d-dm) $ f $ inc_depth (dm-d) v)) <$>
      substHyp h vh'
  insertHypBefore h h' cth' = do
    ContextTerm dh uc th' <- pullTerm h cth'
    dm <- length <$> ask
    first (\f cx@(ContextTerm d uc' x) ->
             if d <= dh then cx
             else ContextTerm (d+1) (uc+uc') (inc_depth (d-dm) $ f $ inc_depth (dm-d) x))
            <$> insertHypBefore h h' th'

data NodeDir str ax a = NodeDir
  (Map BindType (NodeDir str ax (NodeDir str ax a)))
  (ApDir str ax a)
  (Map Int a)
  deriving (Eq,Ord,Show,Generic)
instance Functor (NodeDir str ax) where
  map f (NodeDir a b c) = NodeDir (map3 f a) (map3 f b) (map f c)
instance Foldable (NodeDir str ax) where
  fold (NodeDir a b c) = (fold.map fold.map2 fold) a + (fold.map fold.map2 fold) b + fold c
instance Ord ax => Traversable (NodeDir str ax) where
  sequence (NodeDir a b c) = NodeDir<$>sequence3 a<*>sequence3 b<*>sequence c

instance (Serializable ListStream str,Serializable ListStream a, Serializable ListStream ax) => Serializable ListStream (Cofree (NodeDir str ax) a) where encode = encodeCofree
instance (ListSerializable str, ListSerializable a, ListSerializable ax) => ListSerializable (NodeDir str ax a)
instance (Ord ax,Format ListStream str,Format ListStream a, Format ListStream ax) => Format ListStream (Cofree (NodeDir str ax) a) where datum = datumCofree
instance (Ord ax,ListFormat str, ListFormat a, ListFormat ax) => ListFormat (NodeDir str ax a)

i'NodeDir :: Iso (NodeDir str ax a) (NodeDir str' ax' a')
             ((,,) (Map BindType (NodeDir str ax (NodeDir str ax a)))
               (ApDir str ax a)
               (Map Int a))
             ((,,) (Map BindType (NodeDir str' ax' (NodeDir str' ax' a')))
               (ApDir str' ax' a')
               (Map Int a'))
i'NodeDir = iso (\(x,y,z) -> NodeDir x y z) (\(NodeDir x y z) -> (x,y,z))

type ApDir str ax a = AHDir str ax (FreeMap (NodeDir str ax) a)
data AHDir str ax a = AHDir
  (Map Int a)
  (Map Int (ApDir str ax a))
  (Map ax a)
  deriving (Eq,Ord,Show,Generic)
instance Functor (AHDir str ax) where
  map f (AHDir a b c) = AHDir (map f a) ((map2.map2) f b) (map f c)
instance Foldable (AHDir str ax) where
  fold (AHDir a b c) = fold a + (fold.map fold.map2 fold.map3 fold) b + fold c
instance Ord ax => Traversable (AHDir str ax) where
  sequence (AHDir a b c) = AHDir<$>sequence a<*>(sequence3.map3 sequence) b<*>sequence c
instance (ListSerializable str, ListSerializable ax, ListSerializable a) => ListSerializable (AHDir str ax a)
instance (Ord ax,ListFormat str, ListFormat ax, ListFormat a) => ListFormat (AHDir str ax a)
i'AHDir :: Iso (AHDir str ax a) (AHDir str' ax' a')
           ((,,) (Map Int a) (Map Int (ApDir str ax a)) (Map ax a))
           ((,,) (Map Int a') (Map Int (ApDir str' ax' a')) (Map ax' a'))
i'AHDir = iso (uncurry3 AHDir) (\(AHDir x y z) -> (x,y,z))

i'Cofree :: Iso (Cofree f a) (Cofree f' a') (a,Coforest f a) (a',Coforest f' a')
i'Cofree = iso (uncurry Step) (\(Step x y) -> (x,y))

instance Ord ax => Semigroup (NodeDir str ax a) where NodeDir a b c + NodeDir a' b' c' = NodeDir (a+a') (b+b') (c+c')
instance Ord ax => Monoid (NodeDir str ax a) where zero = NodeDir zero zero zero
instance Ord ax => DataMap (NodeDir str ax a) (Term str ax) a where
  at (Bind t _ tx e) = from i'NodeDir.l'1.at t.l'Just zero.at tx.l'Just zero.at e
  at (Cons a) = from i'NodeDir.l'2.atAp a
  at (Universe u) = from i'NodeDir.l'3.at u

instance Ord ax => Semigroup (AHDir str ax a) where AHDir a b c + AHDir a' b' c' = AHDir (a+a') (b+b') (c+c')
instance Ord ax => Monoid (AHDir str ax a) where zero = AHDir zero zero zero
instance Ord ax => DataMap (AHDir str ax a) (ApHead str ax) a where
  at (Sym i) = from i'AHDir.l'1.at i
  at (Mu xs _ a) = from i'AHDir.l'2.at (length xs).l'Just zero.atAp a
  at (Axiom _ a) = from i'AHDir.l'3.at a

type StringPattern str = [str :+: Int]

atAp :: Ord ax => Application str ax -> Lens' (ApDir str ax a) (Maybe a)
atAp (Ap h xs) = at h.l'Just zero.at xs

mayChoose (Just x) = return x
mayChoose Nothing = zero

(<++>) :: WriterT w [] a -> WriterT w [] a -> WriterT w [] a
a <++> b = a & from writerT %~ (+ b^..writerT)

findPattern :: Ord ax => NodeDir str ax a -> Term str ax -> [([([(str,Term str ax)],Int,Term str ax)],a)]
findPattern = \x y -> go [] x y^..writerT
  where go :: Ord ax => [(str,Term str ax)] -> NodeDir str ax a -> Term str ax -> WriterT [([(str,Term str ax)],Int,Term str ax)] [] a
        go_a :: Ord ax => [(str,Term str ax)] -> ApDir str ax a -> Application str ax -> WriterT [([(str,Term str ax)],Int,Term str ax)] [] a
        go_ah :: Ord ax => [(str,Term str ax)] -> AHDir str ax a -> ApHead str ax -> WriterT [([(str,Term str ax)],Int,Term str ax)] [] a
        withEnv env d x m = foldr (\(i,as) ma -> ma <++> (foldl'.foldl') (\l a -> (tell [(env,i-length env,x)] >> return a) <++> l) zero as)
                            m (d^??from i'NodeDir.l'2.from i'AHDir.l'1.ascList.each.sat ((>=length env) . fst))
        go env d wh@(Bind t x tx e) = withEnv env d wh $ do
          d' <- mayChoose (d^.from i'NodeDir.l'1.at t)
          d'' <- go env d' tx
          go ((x,tx):env) d'' e
        go env d wh@(Cons a) = withEnv env d wh $ go_a env (d^.from i'NodeDir.l'2) a
        go env d wh@(Universe u) = withEnv env d wh $ mayChoose (d^.from i'NodeDir.l'3.at u)

        go_a env d (Ap ah xs) = do
          d' <- go_ah env d ah
          foldr
            (\x ma d'' -> ma =<< go env (d''^.from i'Cofree.l'2) x)
            (\d''' -> mayChoose (d'''^.from i'Cofree.l'1))
            xs d' 

        go_ah _ d (Axiom _ a) = mayChoose (d^.from i'AHDir.l'3.at a)
        go_ah env d (Sym i) | i<length env = mayChoose (d^.from i'AHDir.l'1.at i)
                            | otherwise = []^.writerT
        go_ah env d (Mu tenv _ a) = do
          d' <- mayChoose (d^.from i'AHDir.l'2.at (length tenv))
          go_a env d' a

-- `adjust_depth f e` produces an expression `e'` whose variables (de
-- Bruijin indices) are adjusted from `e` by the function `f`.
--
-- `f` takes two arguments `i` and `d`, where `i` is the previous
-- variable depth, and `d` is the current depth of the considered node
-- (the number of binders between the top-level and the node in
-- question).
--
 -- For example, `adjust_depth (\i d -> i-d+1) (Bind Lambda "x" (Universe 0) (Cons (Ap (Sym 1) [])))
--               == Bind Lambda "x" (Universe 0) (Cons (Ap (Sym 2) []))`

adjust_depth f = go 0
  where go d (Bind t x tx e) = Bind t x (go d tx) (go (d+1) e)
        go _ (Universe u) = Universe u
        go d (Cons a) = Cons (go_a d a)
        go_a d (Ap (Sym i) subs) | i<d = Ap (Sym i) (map (go d) subs)
                                 | otherwise = Ap (Sym (d+f (i-d))) (map (go d) subs)
        go_a d (Ap (Mu env t a') subs) = Ap (Mu
                                            (reverse $ zipWith (\i (x,tx,tx') -> (x,go (d+i) tx,go (d+i) tx')) [0..] (reverse env))
                                            (zipWith (\i -> go (d+length env+i)) [0..] t)
                                            (go_a (d+length env) a')) (map (go d) subs)
        go_a d (Ap x@(Axiom _ _) subs) = Ap x (map (go d) subs)

map_univs f = go
  where go (Bind t x tx e) = Bind t x (go tx) (go e)
        go (Universe u) = Universe (f u)
        go (Cons a) = Cons (go_a a)
        go_a (Ap h subs) = Ap (go_ah h) (map go subs)
        go_ah (Mu e t a) = Mu (map (\(x,y,z) -> (x,go y,go z)) e) (map go t) (go_a a)
        go_ah (Axiom t a) = Axiom (go t) a
        go_ah x = x

inc_depth 0 = \x -> x
inc_depth dx = adjust_depth (+dx)
adjust_telescope_depth field f = zipWith (field . adjust_depth . \i j -> if j<i then j else i+f (j-i)) [0..]
free_vars :: Term str a -> Set Int
free_vars (Bind _ _ tx e) = free_vars tx + delete (-1) (map (subtract 1) (free_vars e))
free_vars (Cons a) = freeA a
  where freeA (Ap (Sym i) xs) = singleton' i + foldMap free_vars xs
        freeA (Ap (Mu env _ a') xs) = foldMap free_vars xs +
          map (subtract envS) (freeA a' - fromKList [0..envS-1])
          where envS = length env
        freeA (Ap (Axiom _ _) xs) = foldMap free_vars xs
free_vars _ = zero

is_free_in :: Int -> Term str a -> Bool
is_free_in = map2 not go
  where go v (Bind _ _ t e) = go v t && go (v+1) e
        go v (Cons a) = go_a v a
        go _ (Universe _) = True
        go_a v (Ap (Sym v') subs) = v/=v' && all (go v) subs
        go_a v (Ap (Mu env _ a) subs) = go_a (v+length env) a && all (go v) subs
        go_a v (Ap (Axiom _ _) subs) = all (go v) subs
        
subst :: (Show str,Show a) => Term str a -> Term str a -> Term str a
subst = flip substn 0
substn :: (Show str,Show a) => Term str a -> Int -> Term str a -> Term str a
substn val n | n>=0 = getId . go n
             | otherwise = error "'subst' should not be called with a negative index"
  where go d (Bind t x tx e) = do
          tx' <- go d tx
          Bind t x tx' <$> go (d+1) e
        go _ (Universe u) = pure (Universe u)
        go d (Cons a) = go_a d a

        go_a d (Ap (Sym i) xs) = traverse (go d) xs >>= \xs' ->
          case compare i d of
            EQ -> rec_subst xs' (inc_depth d val)
            LT -> return $ Cons $ Ap (Sym i) xs'
            GT -> return $ Cons $ Ap (Sym (i-1)) xs'
        go_a d (Ap (Mu e t a) xs) = do
          x <- go_a (d+length e) a
          xs' <- traverse (go d) xs
          e' <- sequence $ reverse $ zipWith (\d' (y,ty,ty') -> liftA2 (y,,) (go (d+d') ty) (go (d+d') ty'))
                [0..] (reverse e)
          t' <- sequence $ zipWith (\d' -> go (d+d')) [length e..] t
          case x of
            Cons a' -> return (Cons (Ap (Mu e' t' a') xs'))
            _ -> rec_subst xs' =<< go_mu d e' t' x
        go_a d (Ap x@(Axiom _ _) xs) = Cons . Ap x <$> traverse (go d) xs
        
        go_mu d env (tx':t) (Bind Lambda x tx e) = go_mu d ((x,tx,tx'):env) t e
        go_mu _ env _ (Cons (Ap (Sym i) xs))
          | i < length env = do
              let envS = length env
                  muEnv = reverse $ map (by l'3) env
              a' <- Cons . Ap (Sym i) <$>
                sequence (fold [if nonempty (free_vars x - fromKList [0..envS-1])
                                then [ return $ inc_depth envS $ foldl' (\e (x',tx,_) -> Bind Lambda x' tx e) x env
                                     , return $ subst x (Cons (Ap (Mu [] muEnv (Ap (Sym 0) [])) [Cons (Ap (Sym j) []) | j <- reverse [1..envS]]))]
                                else [return x]
                               | x <- xs])
              return $ foldl' (\e (x,_,tx) -> Bind Lambda x tx e) a' env
        go_mu _ e t (Cons a) = return $ Cons (Ap (Mu e t a) [])
        go_mu _ _ _ x' = error $ "Cannot produce an induction principle for a term : "+show x'

        rec_subst (y:t) (Bind Lambda _ _ e) = rec_subst t (subst y e)
        rec_subst xs (Cons (Ap h hxs)) = return (Cons (Ap h (hxs+xs)))
        rec_subst [] x = return x
        rec_subst _ x = error $ "Invalid substitution of non-lambda expression : "+show x

fresh env v = head $ select (not . (`elem` env)) (v:[v+fromString (show i) | i <- [0..]])
freshContext = go []
  where go env ((n,v):t) = let n' = fresh env n in (n',(n,v)):go (n':env) t
        go _ [] = []

data NodeDoc str = DocSeq [NodeDoc str]
                 | DocParen (NodeDoc str)
                 | DocMu (NodeDoc str)
                 | DocSubscript (NodeDoc str) (NodeDoc str)
                 | DocSuperscript (NodeDoc str) (NodeDoc str)
                 | DocAssoc str (NodeDoc str)
                 | DocVarName str
                 | DocText str
                 | DocArrow
                 | DocSpace
                 deriving Show
par lvl d msg | d>lvl = DocParen msg
              | otherwise = msg

instance Functor NodeDoc where
  map f (DocSeq l) = DocSeq (map2 f l)
  map f (DocParen x) = DocParen (map f x)
  map f (DocMu x) = DocMu (map f x)
  map f (DocSubscript x y) = DocSubscript (map f x) (map f y)
  map f (DocSuperscript x y) = DocSuperscript (map f x) (map f y)
  map f (DocAssoc v x) = DocAssoc (f v) (map f x)
  map f (DocText x) = DocText (f x)
  map f (DocVarName x) = DocVarName (f x)
  map _ DocArrow = DocArrow
  map _ DocSpace = DocSpace
instance IsString str => IsString (NodeDoc str) where fromString = DocText . fromString

doc2raw :: IsCapriconString str => NodeDoc str -> str
doc2raw (DocSeq l) = fold (map doc2raw l)
doc2raw (DocParen p) = "("+doc2raw p+")"
doc2raw (DocMu m) = "μ("+doc2raw m+")"
doc2raw (DocSubscript v x) = doc2raw v+doc2raw x
doc2raw (DocSuperscript v x) = doc2raw v+"^"+doc2raw x
doc2raw (DocAssoc x v) = "("+x+" : "+doc2raw v+")"
doc2raw DocArrow = " -> "
doc2raw (DocText x) = x
doc2raw (DocVarName x) = x
doc2raw DocSpace = " "

doc2latex :: IsCapriconString str => NodeDoc str -> str
doc2latex (DocSeq l) = fold (map doc2latex l)
doc2latex (DocParen p) = "("+doc2latex p+")"
doc2latex (DocMu m) = "\\mu("+doc2latex m+")"
doc2latex (DocSubscript v x) = doc2latex v+"_{"+doc2latex x+"}"
doc2latex (DocSuperscript v x) = doc2latex v+"^{"+doc2latex x+"}"
doc2latex (DocAssoc x v) = "("+latexName x+":"+doc2latex v+")"
doc2latex DocArrow = " \\rightarrow "
doc2latex (DocText x) = x
doc2latex (DocVarName x) = latexName x
doc2latex DocSpace = "\\,"

doc2svg :: IsCapriconString str => NodeDoc str -> str
doc2svg = \x -> snd $ (go x^.from state) (0::Double)
  where
    sym s = get >>= \x -> if x == 0 then return s
                          else ("<tspan dy=\""+fromString (show x)+"em\">"+s+"</tspan>") <$ put 0
    go (DocSeq l) = fold<$>traverse go l
    go (DocParen p) = liftA3 (\x y z -> x+y+z) (sym "(") (go p) (sym ")")
    go (DocMu m) = liftA3 (\x y z -> x+y+z) (sym "μ(") (go m) (sym ")")
    go (DocSubscript v x) = sub (go v) (go x)
    go (DocSuperscript v x) = super (go v) (go x)
    go (DocAssoc x v) = fold<$>sequence [sym "(",svgName x,sym ":",go v,sym ")"]
    go DocArrow = sym " → "
    go (DocText x) = sym x
    go (DocVarName x) = svgName x
    go DocSpace = sym " "

    super mv mx = liftA2 (\x y -> x+"<tspan dy=\"-0.5em\"><tspan class=\"small\">"+y+"</tspan></tspan>") mv (mx <* put (0.5))
    sub mv mx = liftA2 (\x y -> x+"<tspan dy=\"0.3em\"><tspan class=\"small\">"+y+"</tspan></tspan>") mv (mx <* put (-0.3))

    svgName s = map (\x -> "<tspan class=\"variable\">"+x+"</tspan>") $ nm $ toString s
      where nm ('.':t) = super (nm t) (sym "P")
            nm x =
              let (n,y) = span (\c -> c>='0' && c<='9') (reverse x) in
                case n of
                  "" -> sym (fromString (reverse y))
                  _ -> sub (sym (fromString (reverse y))) (sym (fromString (reverse n)))

latexName :: IsCapriconString str => str -> str
latexName s = fromString $ go $ toString s
  where go ('.':t) = go t+"^P"
        go x = let (n,y) = span (\c -> c>='0' && c<='9') (reverse x) in
          "\\mathrm{"+reverse y+"}"+case n of
                                      "" -> ""
                                      _ -> "_{"+n+"}"

showNode = showNode' zero
showNode' :: (IsCapriconString str,Show ax,Ord ax) => NodeDir str ax ([str],StringPattern str) -> [(str,Term str ax)] -> Term str ax -> NodeDoc str
showNode' dir = go 0
  where go d env x | Just ret <- toPat d env x = ret
        go _ _ (Universe u) = DocSubscript "Set" (fromString (show u))
        go d env whole@(Bind t aname atype body) | t == Lambda || 0`is_free_in`body = par 0 d $ DocSeq (DocText (bind_head t):drop 1 (bind_tail env whole))
                                                 | otherwise = par 0 d $ DocSeq [go 1 env atype,DocArrow,go 0 ((aname,atype):env) body]
          where bind_head Lambda = "λ"
                bind_head Prod = "∀"
                bind_sep Prod = "," ; bind_sep Lambda = "."
                bind_tail env' x | Just ret <- toPat 0 (env'+env) x = [bind_sep t,DocSpace,ret]
                bind_tail env' (Bind t' x tx e) | t==t' && (t==Lambda || 0`is_free_in`e) =
                                                    [DocSpace,DocAssoc x' (go 0 env' tx)] + bind_tail ((x',tx):env') e
                  where x' = fresh (map fst env') x
                bind_tail env' x = [bind_sep t,DocSpace,go 0 env' x]
        go d env (Cons a) = showA d env a
          where showA _ envA (Ap h xs) =
                  let ni = case h of
                             Sym i -> DocVarName $ case drop i envA of
                               (h',_):_ -> h'
                               _ -> "#"+fromString (show i)
                             Mu envD _ a' -> DocMu (showA 0 (map (\(x,tx,_) -> (x,tx)) envD + envA) a')
                             Axiom _ ax -> DocText (fromString $ show ax)
                      lvl = if empty xs then 1000 else 1
                  in par lvl d $ DocSeq $ intercalate [DocSpace] $ map pure (ni:map (go 2 envA) xs)

        toPat d env x
          | (pats,(_,k)):_ <- findPattern dir x =
              let holes = c'map $ fromAList [(i,(env',hole)) | (env',i,hole) <- pats] in
                Just $ par (if all (has t'1) k then 1000 else 1 :: Int) d $ DocSeq $ intercalate [DocSpace] $ map pure $
                [case word of
                   Left w -> DocText w
                   Right i | Just (env',hole) <- lookup i holes ->
                               go 2 env $
                               let (hole',env'') =
                                     fix (\kj -> \case
                                             (Cons (Ap h t@(_:_)),_:env0)
                                               | Cons (Ap (Sym 0) []) <- last t
                                               , not (is_free_in 0 (Cons (Ap h (init t))))
                                                 -> kj (inc_depth (-1) (Cons (Ap h (init t))),env0)
                                             (Cons (Ap (Sym j') []),_:env0) | j'>0 -> kj (Cons (Ap (Sym (j'-1)) []),env0)
                                             e -> e) (hole,env')
                               in foldl' (\e (n,t) -> Bind Lambda n t e) hole' env''
                           | otherwise -> DocText "?"
                | word <- k]
          | otherwise = Nothing

type_of :: (Show a,IsCapriconString str,MonadReader (Env str (Term str a)) m) => Term str a -> m (Maybe (Term str a))
type_of = yb maybeT . go
  where go (Bind Lambda x tx e) = Bind Prod x tx <$> local ((x,tx):) (go e)
        go (Bind Prod x tx e) = do
          a <- go tx
          b <- local ((x,tx):) $ go e
          case (a,b) of
            (Universe ua,Universe ub) -> return (Universe (max ua ub))
            _ -> zero
        go (Universe u) = return (Universe (u+1))
        go (Cons a) = go' a
          where go' (Ap (Sym i) subs) = do
                  e <- ask
                  case drop i e of
                    (_,ti):_ -> rec_subst subs (inc_depth (i+1) ti)
                    _ -> zero
                go' (Ap (Mu env _ a') subs) = do
                  ta <- local (map (\(x,tx,_) -> (x,tx)) env +) (go' a')
                  preret <- maybeT $^ mu_type $ foldl' (\e (x,tx,_) -> Bind Prod x tx e) ta env
                  rec_subst subs (subst (foldl' (\e (x,tx,_) -> Bind Lambda x tx e) (Cons a') env) preret)
                go' (Ap (Axiom t _) subs) = rec_subst subs t
                    
                rec_subst (y:t) (Bind Prod _ tx e) = do
                  ty <- go y
                  (dx,_) <- return (convertDelta ty tx)^.maybeT
                  guard (dx<=0)
                  rec_subst t (subst y e)
                rec_subst [] x = return x
                rec_subst _ _ = zero

mu_type :: MonadReader (Env str (Term str a)) m => Term str a -> m (Maybe (Term str a))
mu_type (inc_depth 1 -> root_type) = yb maybeT $ go 0 root_type
  where
    root_args = go' root_type
      where go' (Bind Prod x tx e) = (x,tx):go' e
            go' _ = []
    nargs = length root_args
    bind t = flip $ foldr (\(x,tx) e -> Bind t x tx e) 
    constr_ind d d' i = d' <= i && i < d+d'

    go d (Bind Prod x tx e) = do
      tx' <- go_col d x tx
      e' <- local ((x,tx):) (go (d+1) e)
      return (Bind Prod x tx' e')
    go _ (Cons (Ap (Sym i) args)) = return $ Cons (Ap (Sym i) $ args + [Cons (Ap (Sym nargs) [])])
    go _ _ = zero

    go_col d xn = go_col' 0 (c'set zero)
      where go_col' d' recs (Bind Prod x tx@(Cons (Ap (Sym i) subs)) e)
              | constr_ind d d' i = do
                  let tx' = bind Prod (adjust_telescope_depth second (+(d+d')) root_args)
                            (adjust_depth (\i' -> if constr_ind d d' i' then (i'-d')+(nargs-d) else i'+nargs) tx)
                      tIx = Cons $ Ap (Sym (i+1)) $ map (inc_depth 1) subs + [Cons (Ap (Sym 0) [])]
                  e' <- local (((x,tx):) . (undefined:)) (go_col' (d'+2) (touch (1 :: Int) (map (+2) recs))
                                                          (adjust_depth (\j -> if j==0 then j else j+1) e))
                  return $ Bind Prod x tx' (Bind Prod x tIx e')
            go_col' d' recs (Bind Prod x tx e) = Bind Prod x tx <$> local ((x,tx):) (go_col' (d'+1) (map (+1) recs) e)
            go_col' d' recs (Cons (Ap (Sym i) xs))
              | constr_ind d d' i = do
                  let args = reverse $ select (not . (`isKeyIn`recs)) [0..d'-1]
                      lastE = bind Lambda (adjust_telescope_depth second (+(d+d')) root_args)
                              (Cons (Ap (Sym (nargs-d-1))
                                     [Cons (Ap (Sym (j'+nargs)) args')
                                     | j <- args
                                     , let (j',args') | (j+1)`isKeyIn`recs = (j+1,[Cons (Ap (Sym k) []) | k <- reverse [0..nargs-1]])
                                                      | otherwise = (j,[])
                                     ]))
                  return $ Cons (Ap (Sym i) $ xs+[lastE])
                  
            go_col' d' recs (Universe u) = do
              let tIH = bind Prod (adjust_telescope_depth second (+(d+d')) root_args) ihRoot
                  ihRoot = Cons (Ap (Sym (nargs-d-1)) [Cons $ Ap (Sym (j+nargs)) $
                                                      if j `isKeyIn` recs
                                                      then [Cons (Ap (Sym k) []) | k <- reverse [0..nargs-1]]
                                                      else []
                                                      | j <- reverse $ select (not . (`isKeyIn`recs) . (+1)) [0..d'-1]])
              return $ Bind Prod xn tIH (Universe (u+1))
            go_col' _ _ _ = zero

convertDelta :: Term str a -> Term str a -> Maybe (Int,Int)
convertDelta = \x y -> map ((getMax<#>getMax) . fst) ((tell (Max 0,Max 0) >> go False x y)^..writerT)
  where go inv (Bind b _ tx e) (Bind b' _ tx' e') = guard (b==b') >> go (not inv) tx tx' >> go inv e e'
        go inv (Cons ax) (Cons ay) = go_a inv ax ay
        go inv (Universe u) (Universe v) | u>v = tellInv inv (Max (u-v),zero)
                                         | otherwise = return ()
        go _ _ _ = lift Nothing
        
        go_a inv (Ap hi ai) (Ap hj aj) = go_ah inv hi hj >> sequence_ (zipWith (go inv) ai aj)
  
        go_ah _ (Sym i) (Sym j) | i==j = return ()
        go_ah inv (Mu _ _ x) (Mu _ _ y) = go_a inv x y
        go_ah _ _ _ = lift Nothing
        
        tellInv True (x,y) = tell (y,x)
        tellInv False (x,y) = tell (x,y)