{-# LANGUAGE 
  FlexibleContexts, 
  FlexibleInstances, 
  FunctionalDependencies, 
  MultiParamTypeClasses, 
  OverlappingInstances, 
  RankNTypes, 
  TypeOperators, 
  UndecidableInstances 
 #-}

{---------------------------------------------------------------------

Type level machinery

---------------------------------------------------------------------}

--
-- Type level Bools
--
data HTrue
data HFalse

class HOr x y res | x y -> res
instance HOr HFalse HFalse HFalse
instance HOr HTrue HFalse HTrue
instance HOr HFalse HTrue HTrue
instance HOr HTrue HTrue HTrue

--
-- Type level Nats
--
data Z = Z
data S a = S a

--
-- Type level Nat equality
--
class EQ x y res | x y -> res
instance EQ Z Z HTrue
instance EQ Z (S y) HFalse
instance EQ (S x) Z HFalse
instance EQ x y res => EQ (S x) (S y) res

--
-- Type level maybe for linear variables: F is a variable, U is used
--
data U 
data F a 

--
-- Machinery needed to check/enforce conditions for additive conjunction/disjunction to correctly deal with Top.
--
class VarOk tf v
instance VarOk HTrue (F v)
instance VarOk HTrue U
instance VarOk HFalse U

class x :<: y 
instance () :<: ()
instance (xs :<: ys) => (U, xs) :<: (U, ys)
instance (xs :<: ys) => (U, xs) :<: (F y, ys)
instance (xs :<: ys) => (F x, xs) :<: (F x, ys)

class Intersect xs ys zs | xs ys -> zs
instance Intersect () () ()
instance Intersect xs ys zs => Intersect (U, xs) (U, ys) (U, zs)
instance Intersect xs ys zs => Intersect (U, xs) (F y, ys) (U, zs)
instance Intersect xs ys zs => Intersect (F x, xs) (U, ys) (U, zs)
instance Intersect xs ys zs => Intersect (F x, xs) (F x, ys) (F x, zs)

class WithOk tf0 tf1 h0 h1 tf2 | tf0 tf1 h0 h1 -> tf2 
instance WithOk HFalse HFalse h h HFalse 
instance (h0 :<: h1) => WithOk HFalse HTrue h0 h1 HFalse 
instance (h1 :<: h0) => WithOk HTrue HFalse h0 h1 HFalse 
instance WithOk HTrue HTrue h0 h1 HTrue 

--
-- Type level machinery for consuming a variable in a list (nested pair) of variables.
--
class Consume vid i o | vid i -> o 
class Consume1 b vid v vs res | b vid v vs -> res

instance Consume vid vs res => Consume vid (U, vs) (U, res) 
instance (EQ vid v b, Consume1 b vid v vs res) => Consume vid (F v, vs) res 

instance Consume1 HTrue v vid vs (U, vs) 
instance Consume vid vs res => Consume1 HFalse vid v vs (F v, res) 


{---------------------------------------------------------------------

Final tagless HOAS linear lambda calculus.

---------------------------------------------------------------------}

--
-- Linear types
--
newtype a -<> b = Lolli {unLolli :: a -> b}
data a <*> b = Tensor a b
data One = One
type a <&> b = (a, b)
type Top = ()
data a <+> b = Inl a | Inr b
data Zero 
newtype Bang a = Bang {unBang :: a}


type VarTp repr vid a = forall v i o . Consume vid i o => repr v HFalse i o a

class Lin repr where
    int :: Int -> repr vid HFalse h h Int
    add :: (HOr tf0 tf1 tf2) => repr vid tf0 hi h Int -> repr vid tf1 h ho Int -> repr vid tf hi ho Int

    llam :: (VarOk tf var) => (VarTp repr vid a -> repr (S vid) tf (F vid, hi) (var, ho) b) -> repr vid tf hi ho (a -<> b)
    (<^>) :: (HOr tf0 tf1 tf2) => repr vid tf0 hi h (a -<> b) -> repr vid tf1 h ho a -> repr vid tf2 hi ho b

    (!) :: repr vid tf h h a -> repr vid HFalse h h (Bang a)
    letBang :: (HOr tf0 tf1 tf2) => repr vid tf0 hi h (Bang a) -> ((forall v tf h . repr v tf h h a) -> repr vid tf1 h ho b) -> repr vid tf2 hi ho b

    lam :: ((forall v tf h . repr v tf h h a) -> repr vid tf hi ho b) -> repr vid tf hi ho (a -> b)
    (<$>) :: repr vid tf0 hi ho (a -> b) -> repr vid tf1 ho ho a -> repr vid tf0 hi ho b

    one :: repr vid hFalse h h One
    letOne :: (HOr tf0 tf1 tf2) => repr vid tf0 hi h One -> repr vid tf1 h ho a -> repr vid tf2 hi ho a

    (<*>) :: (HOr tf0 tf1 tf2) => repr vid tf0 hi h a -> repr vid tf1 h ho b -> repr vid tf2 hi ho (a <*> b)
    letStar :: (VarOk tf1 var0, VarOk tf1 var1, HOr tf0 tf1 tf2) => 
               repr vid tf0 hi h (a <*> b) -> 
               (VarTp repr vid a -> VarTp repr (S vid) b -> repr (S (S vid)) tf1 (F vid, (F (S vid), h)) (var0, (var1, ho)) c) -> 
               repr vid tf2 hi ho c
                    
    top :: repr vid HTrue h h Top

    (<&>) :: (WithOk tf0 tf1 h0 h1 tf2, Intersect h0 h1 ho) => 
             repr vid tf0 hi h0 a -> repr vid tf1 hi h1 b -> repr vid tf2 hi ho (a <&> b)
    pi1 :: repr vid tf hi ho (a <&> b) -> repr vid tf hi ho a 
    pi2 :: repr vid tf hi ho (a <&> b) -> repr vid tf hi ho b 

    inl :: repr vid tf hi ho a -> repr vid tf hi ho (a <+> b)
    inr :: repr vid tf hi ho b -> repr vid tf hi ho (a <+> b)

    letPlus :: (VarOk tf1 var1, VarOk tf2 var2, WithOk tf1 tf2 ho1 ho2 tf12, Intersect ho1 ho2 ho, HOr tf0 tf12 tf3) => 
               repr vid tf0 hi h (a <+> b) ->
               (VarTp repr vid a -> repr (S vid) tf1 (F vid, h) (var1, ho1) c) ->
               (VarTp repr vid b -> repr (S vid) tf2 (F vid, h) (var2, ho2) c) ->
               repr vid tf3 hi ho c

    abort :: repr vid tf hi ho Zero -> repr vid HTrue hi ho a


newtype R vid tf hi ho a = R {unR :: a}
instance Lin R where
    int = R 
    add x y = R $ unR x + unR y

    llam f = R $ Lolli $ \x -> unR (f (R x))
    f <^> x = R $ unLolli (unR f) (unR x)

    (!) = R . Bang . unR
    letBang x f = R $ unR $ f' (unR x)
      where f' (Bang x) = f (R x)

    lam f = R $ \x -> unR (f (R x))
    f <$> x = R $ unR f (unR x)

    x <*> y = R $ Tensor (unR x) (unR y)
    letStar xy f = R $ unR $ f' (unR xy)
      where f' (Tensor x y) = f (R x) (R y)

    one = R One
    letOne x y = R $ unR $ (\One -> y) $ unR x

    x <&> y = R $ (unR x, unR y)
    pi1 = R . fst . unR
    pi2 = R . snd . unR

    top = R ()

    inl = R . Inl . unR
    inr = R . Inr . unR
    letPlus xy fInl fInr = case unR xy of
                             Inl x -> R $ unR $ fInl (R x)
                             Inr y -> R $ unR $ fInr (R y)

    abort x = R $ error "abort"

eval :: R Z tf () () a -> a
eval = unR

{- Some weirdness with higher-order polymorphism.
--    llam f = R $ Lolli $ unR . f . R  -- doesn't work
    llam f = R $ Lolli $ \x -> unR $ f $ R x
--    llam f = R $ Lolli $ \x -> unR (f (R x))  -- works
--    llam (f :: (forall v i o . Consume vid i o => R v i o a) -> R (S vid) (F vid, hi) (U, ho) b) =   -- works
--        R $ Lolli $ (\(x::a) -> unR (f (R x :: forall v i o . Consume vid i o => R v i o a) :: R (S vid) (F vid, hi) (U, ho) b) )  :: R vid hi ho (a -<> b)
-}


-- Str unrestrictedCtx linearCtx
newtype Str vid tf hi ho a = Str {unStr :: NatToInt vid => Int -> vid -> String}

class NatToInt nat where natToInt :: nat -> Int
instance NatToInt Z where natToInt _ = 0
instance NatToInt x => NatToInt (S x) where natToInt (S x) = 1 + natToInt x
showNat :: NatToInt x => x -> String
showNat = show . natToInt

instance Lin Str where
    int x = Str $ \_ _ -> show x
    add x y = Str $ \uv lv -> "(" ++ unStr x uv lv ++ " + " ++ unStr y uv lv ++ ")"

    llam f = Str $ \uv lv -> let v = "x"++showNat lv in 
                           "\\" ++ v ++ " -<> " ++ unStr (f $ Str (\_ _ -> v)) uv (S lv)
    f <^> x = Str $ \uv lv -> "(" ++ unStr f uv lv ++ " ^ " ++ unStr x uv lv ++ ")"

    (!) x = Str $ \uv lv -> "!" ++ unStr x uv lv
    letBang x f = Str $ \uv lv -> 
        let y = "y"++show uv in
        "(" ++ "let !" ++ y ++ " = " ++ unStr x uv lv ++ " in " ++ unStr (f (Str $ \_ _ -> y)) (uv + 1) lv ++ ")"

    lam f = Str $ \uv lv -> let v = "y"++show uv in 
                           "\\" ++ v ++ " -> " ++ unStr (f $ Str (\_ _ -> v)) (uv + 1) lv
    f <$> x = Str $ \uv lv -> "(" ++ unStr f uv lv ++ " " ++ unStr x uv lv ++ ")"

    x <*> y = Str $ \uv lv -> "<" ++ unStr x uv lv ++ "," ++ unStr y uv lv ++ ">"
    letStar xy f = Str $ \uv lv -> 
      let x = "x"++showNat lv
          y = "x"++showNat (S lv)
      in "(" ++ "let " ++ x++"<*>"++y ++" = " ++ unStr xy uv lv ++ " in " ++ unStr (f (Str $ \_ _ -> x) (Str $ \_ _ -> y)) uv (S (S lv)) ++ ")"

    one = Str $ \_ _ -> "<>"
    letOne x y = Str $ \uv lv -> "(" ++ "let <> = " ++ unStr x uv lv ++ " in " ++ unStr y uv lv ++ ")"

    x <&> y = Str $ \uv lv -> "(" ++ unStr x uv lv ++ "," ++ unStr y uv lv ++ ")"
    pi1 x = Str $ \uv lv -> "(" ++ "pi1 " ++ unStr x uv lv ++ ")"
    pi2 x = Str $ \uv lv -> "(" ++ "pi2 " ++ unStr x uv lv ++ ")"

    top = Str $ \_ _ -> "()"

    inl x = Str $ \uv lv -> "(" ++ "inl " ++ unStr x uv lv ++ ")"
    inr x = Str $ \uv lv -> "(" ++ "inr " ++ unStr x uv lv ++ ")"
    letPlus xy fInl fInr = Str $ \uv lv -> 
      let x = "x"++showNat lv in
        "(" ++ "case " ++ unStr xy uv lv ++ " of " ++ 
                "inl " ++ x ++ " => " ++ unStr (fInl (Str $ \_ _ -> x)) uv (S lv) ++ " | " ++
                "inr " ++ x ++ " => " ++ unStr (fInr (Str $ \_ _ -> x)) uv (S lv) ++ 
        ")"

    abort x = Str $ \_ _ -> "abort"  -- \uv lv -> "(" ++ "abort " ++ unStr x uv lv ++ ")"

showLin :: Str Z tf () () a -> String
showLin x = (unStr x) 0 Z