{-# LANGUAGE GADTs #-}

module Expr where

import Data.Function (on)
import Data.List (genericDrop)

data Expr where
    Var :: Integer -> Expr
    Star :: Expr
    Square :: Expr
    App :: Expr -> Expr -> Expr
    Abs :: Expr -> Expr -> Expr
    Pi :: Expr -> Expr -> Expr
    deriving (Show, Eq)

occursFree :: Integer -> Expr -> Bool
occursFree n (Var k) = n == k
occursFree _ Star = False
occursFree _ Square = False
occursFree n (App a b) = occursFree n a || occursFree n b
occursFree n (Abs a b) = occursFree n a || occursFree (n + 1) b
occursFree n (Pi a b) = occursFree n a || occursFree (n + 1) b

{- --------------------- PRETTY PRINTING ----------------------------- -}

-- TODO : store parsed identifiers for better printing
genName :: Integer -> String
genName k = case genericDrop k ["x", "y", "z", "w", "u", "v"] of
    [] -> 'x' : show (k - 6)
    (v : _) -> v

pretty :: Expr -> String
pretty = helper 0
  where
    helper :: Integer -> Expr -> String
    helper k (Var n) = genName $ k - n - 1
    helper _ Star = "*"
    helper _ Square = "□"
    helper k (App e1 e2) = "(" ++ helper k e1 ++ " " ++ helper k e2 ++ ")"
    helper k (Abs ty b) =
        "(λ" ++ genName k ++ " : " ++ helper k ty ++ " . " ++ helper (k + 1) b ++ ")"
    helper k (Pi ty b) =
        if occursFree 0 b
            then
                "(∏" ++ genName k ++ " : " ++ helper k ty ++ " . " ++ helper (k + 1) b ++ ")"
            else "(" ++ helper k ty ++ " -> " ++ helper (k + 1) b ++ ")"

{- --------------- ACTUAL MATH STUFF ---------------- -}

isSort :: Expr -> Bool
isSort Star = True
isSort Square = True
isSort _ = False

mapIndices :: (Integer -> Expr) -> Expr -> Expr
mapIndices f (Var n) = f n
mapIndices _ Star = Star
mapIndices _ Square = Square
mapIndices f (App m n) = App (mapIndices f m) (mapIndices f n)
mapIndices f (Abs m n) = Abs (mapIndices f m) (mapIndices f n)
mapIndices f (Pi m n) = Pi (mapIndices f m) (mapIndices f n)

incIndices :: Expr -> Expr
incIndices = mapIndices (Var . (+ 1))

decIndices :: Expr -> Expr
decIndices = mapIndices (Var . subtract 1)

-- substitute 0 for s *AND* decrement indices; only use after reducing a redex.
subst :: Expr -> Expr -> Expr
subst s (Var 0) = s
subst _ (Var n) = Var $ n - 1
subst _ Star = Star
subst _ Square = Square
subst s (App m n) = App (subst s m) (subst s n)
subst s (Abs m n) = Abs (subst s m) (subst s n)
subst s (Pi m n) = Pi (subst s m) (subst s n)

substnd :: Expr -> Expr -> Expr
substnd s (Var n) = if n == 0 then s else Var n
substnd _ Star = Star
substnd _ Square = Square
substnd s (App m n) = App (substnd s m) (substnd s n)
substnd s (Abs m n) = Abs (substnd s m) (substnd s n)
substnd s (Pi m n) = Pi (substnd s m) (substnd s n)

betaReduce :: Expr -> Expr
betaReduce (Var k) = Var k
betaReduce Star = Star
betaReduce Square = Square
betaReduce (App (Abs _ v) n) = subst n v
betaReduce (App m n) = App (betaReduce m) (betaReduce n)
betaReduce (Abs t v) = Abs (betaReduce t) (betaReduce v)
betaReduce (Pi t v) = Pi (betaReduce t) (betaReduce v)

betaNF :: Expr -> Expr
betaNF e = if e == e' then e else betaNF e'
  where
    e' = betaReduce e

betaEquiv :: Expr -> Expr -> Bool
betaEquiv = on (==) betaNF