{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -Wno-missing-export-lists #-}

module Lib where

import Control.Monad (join)
import Data.Function (on)
import Data.List (elemIndex, find)
import Data.Maybe (fromJust)
import qualified Data.Set as Set
import Prettyprinter (Pretty (..), hsep, parens, (<+>))

type Var = String

type Name = String

data Term = Var !Var | App !Term !Term | Abs !Var !Term

data DBTerm = DBFree !Name | DBVar !Int | DBApp !DBTerm !DBTerm | DBAbs !DBTerm deriving (Eq, Show)

instance (Pretty Term) where
  pretty (Var x) = pretty x
  pretty (App t@(Abs _ _) s@(Abs _ _)) =
    parens (pretty t) <+> parens (pretty s)
  pretty (App t@(Abs _ _) s) =
    parens (pretty t) <+> pretty s
  pretty (App s t@(Abs _ _)) =
    pretty s <+> parens (pretty t)
  pretty (App s t@(App _ _)) =
    pretty s <+> parens (pretty t)
  pretty (App s t) = pretty s <+> pretty t
  pretty t@(Abs _ _) =
    let (vs, body) = getCurriedVars t
     in "λ" <> hsep (pretty <$> vs) <> "." <+> pretty body

instance (Show Term) where
  show = show . pretty

(<.>) :: Term -> Term -> Term
(<.>) = App

getCurriedVars :: Term -> ([Var], Term)
getCurriedVars (Abs x t) =
  let (vs, body) = getCurriedVars t
   in (x : vs, body)
getCurriedVars t = ([], t)

subterms :: Term -> [Term]
subterms t@(Var _) = [t]
subterms t@(App m n) = t : subterms m ++ subterms n
subterms t@(Abs _ m) = t : subterms m

freeVars :: Term -> Set.Set Var
freeVars (Var x) = Set.singleton x
freeVars (App m n) = on Set.union freeVars m n
freeVars (Abs x m) = Set.delete x $ freeVars m

bindingVars :: Term -> Set.Set Var
bindingVars (Var _) = Set.empty
bindingVars (App m n) = on Set.union bindingVars m n
bindingVars (Abs x m) = Set.insert x $ bindingVars m

allVars :: Term -> Set.Set Var
allVars = liftA2 Set.union freeVars bindingVars

freeIn :: Var -> Term -> Bool
freeIn = (. freeVars) . Set.member

bindingIn :: Var -> Term -> Bool
bindingIn = (. bindingVars) . Set.member

freeOrBindingIn :: Var -> Term -> Bool
freeOrBindingIn = join . on (||) . freeIn

isCombinator :: Term -> Bool
isCombinator = Set.null . freeVars

alphaConv :: Var -> Var -> Term -> Term
alphaConv _ _ t@(Var _) = t
alphaConv _ _ t@(App _ _) = t
alphaConv x y s@(Abs x' t) =
  if y `freeIn` t || x /= x'
    then s
    else maybe s (Abs y) (alphaConv' t)
  where
    alphaConv' (Var x'') = if x == x'' then Just (Var y) else Just (Var x'')
    alphaConv' (App m n) = App <$> alphaConv' m <*> alphaConv' n
    alphaConv' t'@(Abs y' m)
      | y == y' = Nothing
      | x == y' = Just t'
      | otherwise = Abs y' <$> alphaConv' m

vars :: [Var]
vars = ["x", "y", "z", "w", "u", "v"] ++ map ("x" <>) (show <$> [1 :: Integer ..])

genNewVar :: Set.Set Var -> Var
genNewVar vs = fromJust $ find (not . (`Set.member` vs)) vars

subst :: Var -> Term -> Term -> Term
subst x n m@(Var x') = if x == x' then n else m
subst x n (App p q) = on App (subst x n) p q
subst x n m@(Abs y p)
  | y `freeIn` n = subst x n $ alphaConv y z m
  | x == y = m
  | otherwise = Abs y $ subst x n p
  where
    z = genNewVar $ on Set.union freeVars n p

betaReduce :: Term -> Term
betaReduce (Var x) = Var x
betaReduce (Abs x m) = Abs x (betaReduce m)
betaReduce (App (Abs x m) n) = subst x n m
betaReduce (App m n) = on App betaReduce m n

betaNf :: Term -> Bool
betaNf (Var _) = True
betaNf (App (Abs _ _) _) = False
betaNf (App m n) = on (&&) betaNf m n
betaNf (Abs _ m) = betaNf m

findBetaNf :: Term -> Term
findBetaNf = fromJust . find betaNf . iterate betaReduce

findBetaNf' :: Term -> Term
findBetaNf' t =
  if betaNf t
    then t
    else findBetaNf' $ betaReduce t

alphaEquiv :: Term -> Term -> Bool
alphaEquiv = on (==) debruijnize

debruijnize :: Term -> DBTerm
debruijnize = helper []
  where
    helper vs (Var x) = maybe (DBFree x) DBVar $ elemIndex x vs
    helper vs (App m n) = on DBApp (helper vs) m n
    helper vs (Abs x m) = DBAbs $ helper (x : vs) m