{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-partial-fields #-}

module Symantic.Compiler.Term where

import Data.Function qualified as Fun
import Data.Monoid (Monoid)
import GHC.Types (Type)
import Symantic.Semantics.Eval (Eval (..))
import Symantic.Semantics.Forall
import Symantic.Syntaxes.Classes (Instantiable (..), Syntax, Syntaxes, Unabstractable (..))
import Symantic.Syntaxes.Classes qualified as Sym
import Type.Reflection (Typeable)
import Unsafe.Coerce (unsafeCoerce)

import Symantic.Typer.Type (Ty, monoTy)

-- * Class 'AbstractableTy'
class AbstractableTy ty sem where
  -- | Lambda term abstraction, in HOAS (Higher-Order Abstract Syntax) style.
  lamTy :: ty a -> (sem a -> sem b) -> sem (a -> b)

fun ::
  forall prov sem a b.
  Monoid prov =>
  AbstractableTy (Ty prov '[]) sem =>
  Typeable a =>
  (sem a -> sem b) ->
  sem (a -> b)
fun = lamTy (monoTy @_ @prov)

instance
  ( forall sem. Syntaxes syns sem => AbstractableTy (Ty prov '[]) sem
  -- , forall sem a. syn sem => AbstractableLam sem a
  -- , forall sem. syn sem => AbstractableLam sem a
  -- , forall sem. syn sem => Typeable sem -- user instance not accepted
  -- , forall s1 s2. (syn s1, syn s2) => s1 ~ s2 -- crazy...
  ) =>
  AbstractableTy (Ty prov '[]) (Forall syns)
  where
  lamTy aTy f = Forall (lamTy aTy (\a -> let Forall b = f (forallSem a) in b))
    where
      -- Safe here because (a :: sem a) and (b :: sem b) share the same 'sem'.
      forallSem :: sem a -> Forall syns a
      forallSem a = Forall (unsafeCoerce a)

-- * Type 'OpenTerm'
data OpenTerm (syns :: [Syntax]) (vs :: [Type]) (a :: Type) where
  -- | 'E' contains embedded closed (i.e. already abstracted) terms.
  E :: Forall syns a -> OpenTerm syns vs a
  -- | 'V' represents a reference to the innermost/top environment variable, i.e. Z
  V :: OpenTerm syns (a ': vs) a
  -- | 'N' represents internalizing the innermost bound variable as a
  -- function argument. In other words, we can represent an open
  -- term referring to a certain variable as a function which
  -- takes that variable as an argument.
  N :: OpenTerm syns vs (a -> b) -> OpenTerm syns (a ': vs) b
  -- | 'W' is a special variant of N for efficiency,
  -- in the case where the term does not refer
  -- to the topmost variable at all.
  W :: OpenTerm syns vs b -> OpenTerm syns (a ': vs) b
instance
  ( forall sem. Syntaxes syns sem => AbstractableTy (Ty prov '[]) sem
  , Syntaxes syns (Forall syns)
  ) =>
  AbstractableTy (Ty prov '[]) (OpenTerm syns '[])
  where
  lamTy aTy f = E (lamTy aTy (unE Fun.. f Fun.. E))
instance
  ( forall sem. Syntaxes syns sem => Unabstractable sem
  , Syntaxes syns (Forall syns)
  ) =>
  Unabstractable (OpenTerm syns vs)
  where
  ap = E Sym.ap
  const = E Sym.const
  id = E Sym.id
  (.) = E (Sym..)
  flip = E Sym.flip
  ($) = E (Sym.$)

instance AbstractableTy (Ty prov '[]) Eval where
  lamTy _aTy f = Eval (unEval Fun.. f Fun.. Eval)

runOpenTerm :: Syntaxes syns Eval => OpenTerm syns '[] a -> a
runOpenTerm t | E (Forall sem) <- t = unEval sem

unE :: OpenTerm syns '[] a -> Forall syns a
unE t = case t of E t' -> t'

instance
  ( forall sem. Syntaxes syns sem => Unabstractable sem
  , forall sem. Syntaxes syns sem => Instantiable sem
  , Syntaxes syns (Forall syns)
  ) =>
  Instantiable (OpenTerm syns vs)
  where
  (.@) = appOpenTerm

appOpenTerm ::
  forall syns as a b.
  ( forall sem. Syntaxes syns sem => Unabstractable sem
  , forall sem. Syntaxes syns sem => Instantiable sem
  , Syntaxes syns (Forall syns)
  ) =>
  OpenTerm syns as (a -> b) ->
  OpenTerm syns as a ->
  OpenTerm syns as b
E d `appOpenTerm` N e = N (E ((Sym..) .@ d) `appOpenTerm` e)
E d `appOpenTerm` V = N (E d)
E d `appOpenTerm` W e = W (E d `appOpenTerm` e)
E d1 `appOpenTerm` E d2 = E (d1 .@ d2)
N e `appOpenTerm` E d = N (E (Sym.flip .@ Sym.flip .@ d) `appOpenTerm` e)
N e `appOpenTerm` V = N (E Sym.ap `appOpenTerm` e `appOpenTerm` E Sym.id)
N e1 `appOpenTerm` N e2 = N (E Sym.ap `appOpenTerm` e1 `appOpenTerm` e2)
N e1 `appOpenTerm` W e2 = N (E Sym.flip `appOpenTerm` e1 `appOpenTerm` e2)
V `appOpenTerm` E d = N (E (Sym.flip .@ Sym.id .@ d))
V `appOpenTerm` N e = N (E (Sym.ap .@ Sym.id) `appOpenTerm` e)
V `appOpenTerm` W e = N (E (Sym.flip .@ Sym.id) `appOpenTerm` e)
W e `appOpenTerm` E d = W (e `appOpenTerm` E d)
W e `appOpenTerm` V = N e
W e1 `appOpenTerm` N e2 = N (E (Sym..) `appOpenTerm` e1 `appOpenTerm` e2)
W e1 `appOpenTerm` W e2 = W (e1 `appOpenTerm` e2)