{-# OPTIONS_GHC -Wno-redundant-constraints #-}
module Plutarch.Docs.PMatch (Tree(..), swap, TreeRepr) where 
import Plutarch.Prelude
import Plutarch.Internal.PlutusType (PlutusType (pcon', pmatch'))
import Plutarch.Builtin (pforgetData, pasConstr, pconstrBuiltin)
import Plutarch.Unsafe (punsafeCoerce)

I - Identity functor (`newtype I a = I a`)
K - Constant functor (`newtype K a b = K a`)

Z - zero (as in Peano numbers)
S - successor (as in Peano numbers)

<https://hackage.haskell.org/package/generics-sop-0.5.1.2/docs/Generics-SOP.html#t:NS>
NS - n-ary Sum
Picking nth element of sum comes from composing Z and S

<https://hackage.haskell.org/package/generics-sop-0.5.1.2/docs/Generics-SOP.html#t:NP>
NP - n-ary Product
Value level witness for a list of types parametrized by some functor f.

SOP - sum of products
`SOP I` corresponds to haskell's structure of given SOP encoded datatype.
`SOP (Term s)` corresponds to Plutarch structure of given SOP encoded datatype.

`Code a`
The code of a datatype.

This is a list of lists of its components. The outer list contains one element per constructor. The inner list contains one element per constructor argument (field).

data Tree = Leaf Int | Node Tree Tree

type TreeRepr =
  '[ '[ Int ]
   , '[ Tree, Tree ]
   ]

-- Generic representation of given Haskell datatype
type Rep a = SOP I (Code a)

data MyType (a :: PType) (b :: PType) (s :: S)
  = One (Term s a)
  | Two (Term s b)
  deriving stock Generic
  deriving anyclass PlutusType
instance DerivePlutusType (MyType a b) where type DPTStrat _ = PlutusTypeScott

-- If you instead want to use data encoding, you should derive 'PlutusType' and provide data strategy:

data MyTypeD (a :: PType) (b :: PType) (s :: S)
  = OneD (Term s (PDataRecord '[ "_0" ':= a ]))
  | TwoD (Term s (PDataRecord '[ "_0" ':= b ]))
  deriving stock Generic
  deriving anyclass PlutusType
instance DerivePlutusType (MyTypeD a b) where type DPTStrat _ = PlutusTypeData

-- Alternatively, you may derive 'PlutusType' by hand as well. A simple example, encoding a
-- Sum type as an Enum via PInteger:

data AB (s :: S) = A | B

instance PlutusType AB where
  type PInner AB = PInteger

  pcon' A = 0
  pcon' B = 1

  pmatch' x f =
    pif (x #== 0) (f A) (f B)


-- instead of using `pcon'` and `pmatch'` directly,
-- use 'pcon' and 'pmatch', to hide the `PInner` type:

swap :: Term s AB -> Term s AB
swap x = pmatch x $ \case
  A -> pcon B
  B -> pcon A

-- | Scott
data PSMaybe a s = PSJust (Term s a) | PSNothing

-- | Newtype wrapper around function that represents Scott encoding,
-- | Plutarch uses generic one for deriving.
newtype ScottEncodedMaybe a b s = ScottEncodedMaybe (Term s ((a :--> b) :--> PDelayed b :--> b))

instance PlutusType (ScottEncodedMaybe a r) where
  type PInner (ScottEncodedMaybe a r) = (a :--> r) :--> PDelayed r :--> r
  pcon' (ScottEncodedMaybe x) = x
  pmatch' x f = f (ScottEncodedMaybe x)

instance PlutusType (PSMaybe a) where
  -- The resulting type of pattern matching on Maybe is quantified via `PForall`
  type PInner (PSMaybe a) = PForall (ScottEncodedMaybe a)
  pcon' (PSJust x) = pcon $ PForall $ pcon $ ScottEncodedMaybe $ plam $ \f _ -> f # x
  pcon' PSNothing = pcon $ PForall $ pcon $ ScottEncodedMaybe $ plam $ \_ g -> pforce g
  pmatch' x' f =
    pmatch x' $ \(PForall sem) ->
      pmatch sem $ \(ScottEncodedMaybe x) ->
        x # plam (f . PSJust) # pdelay (f PSNothing)

-- | Maybe encoded using Constr
data PMaybeData a (s :: S)
  = PDJust (Term s a)
  | PDNothing

-- | Note - thing hold in PMaybeData must be able to be represented as Data too, not needed in case of Scott version
instance PIsData a => PlutusType (PMaybeData a) where
  type PInner (PMaybeData a) = PData
  pcon' (PDJust x) = pforgetData $ pconstrBuiltin # 0 #$ psingleton # pforgetData (pdata x)
  pcon' PDNothing = pforgetData $ pconstrBuiltin # 1 # pnil
  pmatch' x f = (`runTermCont` f) $ do
    constrPair <- TermCont $ plet (pasConstr # x)
    indexNum <- TermCont $ plet (pfstBuiltin # constrPair)
    TermCont $ \g -> pif (indexNum #== 0)
        (g $ PDJust $ punsafeCoerce $ phead # (psndBuiltin # constrPair))
        (pif (indexNum #== 1)
          (g PDNothing)
          perror
        )