module Reduction where

import           AST

import           Control.Parallel

updateTerm :: (Ord i, Enum i) => i -> i -> Term m v i -> Term m v i
updateTerm i k subject =
  case subject of
    Type metadata -> Type metadata
    Application metadata argument function ->
      Application metadata (updateTerm i k argument) (updateTerm i k function)
    Abstraction variable metadata parameter body ->
      Abstraction variable metadata (updateTerm i k parameter) (updateTerm i (succ k) body)
    Variable variable index metadata
      | index > k -> Variable variable (toEnum $ fromEnum index + fromEnum i - 1) metadata
      | otherwise -> subject

substituteTerm :: (Ord i, Enum i) => i -> Term m v i -> Term m v i -> Term m v i
substituteTerm index replacement' subject =
  case subject of
    Type metadata -> Type metadata
    Application metadata argument function ->
      Application metadata (substituteTerm index replacement' argument)
        (substituteTerm index replacement' function)
    Abstraction variable metadata parameter body ->
      Abstraction variable metadata (substituteTerm index replacement' parameter)
        (substituteTerm (succ index) replacement' body)
    Variable variable index' metadata
      | index' > index -> Variable variable (pred index') metadata
      | index' == index -> updateTerm index (toEnum 0) replacement'
      | index' < index -> subject
    Variable{} -> undefined -- this should never be matched!

substituteKind :: (Ord i, Enum i) => i -> Term m v i -> Kind m v i -> Kind m v i
substituteKind index replacement' subject =
  case subject of
    Kind metadata -> Kind metadata
    Function variable metadata parameter body ->
      Function variable metadata (substituteTerm index replacement' parameter)
        (substituteKind (succ index) replacement' body)

whnf :: (Ord i, Enum i) => Term m v i -> Term m v i
whnf reduct =
  case reduct of
    Application metadata argument function ->
      case whnf function of
        Abstraction _ _ _ body ->
          whnf $ substituteTerm (toEnum 1) argument body
        f@_ -> Application metadata argument f
    reduct'@_ -> reduct'

nfTerm :: (Enum i, Ord i) => Term m v i -> Term m v i
nfTerm reduct =
  case reduct of
    Abstraction variable metadata parameter body ->
      let parameter' = nfTerm parameter
          body' = parameter' `par` nfTerm body
      in Abstraction variable metadata parameter' body'
    Application metadata argument function ->
      case whnf function of
        Abstraction _ _ _ body ->
          let arg = nfTerm argument
              bod = arg `seq` substituteTerm (toEnum 1) arg (nfTerm body)
          in nfTerm bod
        f@_ -> Application metadata (nfTerm argument) (nfTerm f)
    _ -> reduct

nfKind :: (Enum i, Ord i) => Kind m v i -> Kind m v i
nfKind reduct =
  case reduct of
    Kind metadata -> Kind metadata
    Function variable metadata parameter body ->
      let param = nfTerm parameter
          body' = param `par` nfKind body
      in Function variable metadata param body'

instance (Enum i, Ord i, Eq v) =>
         Eq (Term m v i) where
  Type _ == Type _ = True
  Variable variable index _ == Variable variable' index' _ = variable == variable' && index == index'
  Application _ argument function == Application _ argument' function' =
    let ftest = function == function'
        atest = ftest `par` argument == argument'
    in ftest && atest
  Abstraction variable metadata parameter body == Abstraction _ _ parameter' body' =
    let ptest = parameter == parameter'
        var = Variable variable (toEnum 1) metadata
        btest = ptest `par` body == substituteTerm (toEnum 1) var body'
    in ptest && btest
  _ == _ = False

instance (Enum i, Ord i, Eq v, Eq i) =>
         Eq (Kind m v i) where
  Kind _ == Kind _ = True
  Function variable metadata parameter body == Function _ _ parameter' body' =
    let ptest = parameter == parameter'
        var = Variable variable (toEnum 1) metadata
        btest = ptest `par` body == substituteKind (toEnum 1) var body'
    in ptest && btest
  _ == _ = False

infix 4 ===

class (Enum i, Ord i, Eq i, Eq v) =>
      BetaEq t v i where
  (===) :: (Enum i, Ord i, Eq i, Eq v) => t v i -> t v i -> Bool

instance (Enum i, Ord i, Eq i, Eq v) =>
         BetaEq (Term m) v i where
  a === b = nfTerm a == nfTerm b

instance (Enum i, Ord i, Eq i, Eq v) =>
         BetaEq (Kind m) v i where
  a === b = nfKind a == nfKind b

instance (Enum i, Ord i, Eq i, Eq v) =>
         BetaEq (PseudoTerm m) v i where
  Term a === Term b = nfTerm a == nfTerm b
  Kind' a === Kind' b = nfKind a == nfKind b
  _ === _ = False