NatHelper.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

{- |
Module      : Cybus.NatHelper
Description : Nat helper methods
Copyright   : (c) Grant Weyburne, 2022
License     : BSD-3
-}
module Cybus.NatHelper (
  -- * peano
  NatToPeanoT,
  PeanoToNatT,
  Peano (..),

  -- * arithmetic
  DiffT,
  DiffTC,
  FacT,
  type (<=!),
  type (<!),
  LTEQT,

  -- * matrix dimension synonyms
  D1,
  D2,
  D3,
  D4,
  D5,
  D6,
  D7,
  D8,
  D9,
  D10,

  -- * matrix helpers
  NS,
  Product1T,
  NN,
  NN',
  Reverse1T,
  ListTupleT,

  -- * list and nonempty conversions
  ValidateNestedListC (..),
  ValidateNestedNonEmptyC (..),
  validateNestedList,
  validateNestedNonEmpty,
  ValidateNestedListT,
  ValidateNestedNonEmptyT,
  nestedNonEmptyToList,
  nestedListToNonEmpty,
  NestedListC (..),
  NonEmptyNST,
  ListNST,
) where

import Data.Kind
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as N
import Data.Pos
import Data.Proxy
import qualified GHC.TypeLits as GL
import GHC.TypeNats (Nat)
import qualified GHC.TypeNats as GN
import Primus.Error
import Primus.Fold
import Primus.List
import Primus.NonEmpty
import Primus.One
import qualified Primus.TypeLevel as TP (Cons1T, FailUnless)

-- | get the factorial of a 'Nat'
type FacT :: Nat -> Nat
type family FacT x where
  FacT 0 = 1
  FacT 1 = 1
  FacT n = n GN.* FacT (n GN.- 1)

-- | constraint for ensuring that "i" <= "n"
type (<=!) :: Nat -> Nat -> Constraint
type i <=! n =
  ( TP.FailUnless
      (i GN.<=? n)
      ( 'GL.Text "i>n"
          'GL.:<>: 'GL.Text ": i="
          'GL.:<>: 'GL.ShowType i
          'GL.:<>: 'GL.Text " n="
          'GL.:<>: 'GL.ShowType n
      )
  , PosT i
  )

-- | constraint for ensuring that "i" <= "n" with a custom error message
type LTEQT :: GL.ErrorMessage -> Nat -> Nat -> Constraint
type LTEQT msg i n =
  ( TP.FailUnless
      (i GN.<=? n)
      ( 'GL.Text "i>n"
          'GL.:<>: 'GL.Text ": i="
          'GL.:<>: 'GL.ShowType i
          'GL.:<>: 'GL.Text " n="
          'GL.:<>: 'GL.ShowType n
          'GL.:<>: msg
      )
  , PosT i
  )

-- | constraint for ensuring that "i" <= "n"
type (<!) :: Nat -> Nat -> Constraint
type i <! n =
  ( TP.FailUnless
      (i GN.+ 1 GN.<=? n)
      ( 'GL.Text "i>=n"
          'GL.:<>: 'GL.Text ": i="
          'GL.:<>: 'GL.ShowType i
          'GL.:<>: 'GL.Text " n="
          'GL.:<>: 'GL.ShowType n
      )
  , GN.KnownNat i
  )

-- | constraint for DiffC with better error messages
type DiffTC :: Nat -> Nat -> Nat -> Constraint
type DiffTC i j n = (i <=! j, j <=! n)

-- | find the number of N between "i" and "j" while ensuring i<=j and j<=n
type DiffT :: Nat -> Nat -> Nat -> Nat
type DiffT i j n = j GN.+ 1 GN.- i

-- | product of a type level nonempty list as a 'Nat'
type Product1T :: NonEmpty Nat -> Nat
type family Product1T ns where
  Product1T (n ':| '[]) = n
  Product1T (n ':| n1 ': ns) = n GN.* Product1T (n1 ':| ns)

-- | convert a list of 'Nat' into a nonempty list of 'Nat'
type NS :: [Nat] -> NonEmpty Nat
type family NS ns where
  NS '[] = GL.TypeError ( 'GL.Text "NS: must have at least one Nat value for NonEmpty Nat")
  NS (n ': '[]) = n ':| '[]
  NS (n ': m ': ns) = TP.Cons1T n (NS (m ': ns))

-- | used for reversing the indices of a matrix using type level list of nonempty indices
type Reverse1T :: forall k. NonEmpty k -> NonEmpty k
type family Reverse1T ns where
  Reverse1T (n ':| ns) = Reverse1T' (n ': ns) '[]

-- | used by 'Reverse1T'
type Reverse1T' :: forall k. [k] -> [k] -> NonEmpty k
type family Reverse1T' ns ret where
  Reverse1T' '[] (r ': rs) = r ':| rs
  Reverse1T' (n ': ns) ret = Reverse1T' ns (n ': ret)

-- | extracts the dimensions of a nested list
type ValidateNestedListT :: Type -> Peano
type family ValidateNestedListT x where
  ValidateNestedListT [x] = 'S (ValidateNestedListT x)
  ValidateNestedListT _ = 'S 'Z

-- | extracts the dimensions of a nested nonempty list
type ValidateNestedNonEmptyT :: Type -> Peano
type family ValidateNestedNonEmptyT x where
  ValidateNestedNonEmptyT (NonEmpty x) = 'S (ValidateNestedNonEmptyT x)
  ValidateNestedNonEmptyT _ = 'S 'Z

-- | validate that the nested nonempty list is consistent in size along all dimensions
validateNestedNonEmpty :: forall x. ValidateNestedNonEmptyC x (ValidateNestedNonEmptyT x) => x -> Either String (NonEmpty Pos)
validateNestedNonEmpty x = validateNestedNonEmptyC @x @(ValidateNestedNonEmptyT x) [] x []

-- | validate that the nested list is consistent in size along all dimensions
validateNestedList :: forall x. ValidateNestedListC x (ValidateNestedListT x) => x -> Either String (NonEmpty Pos)
validateNestedList x = validateNestedListC @x @(ValidateNestedListT x) [] x []

-- | extracts the dimensions of a nested nonempty list: doesnt allow empty dimensions
type ValidateNestedNonEmptyC :: Type -> Peano -> Constraint
class ValidateNestedNonEmptyC x y where
  validateNestedNonEmptyC :: [Pos] -> x -> [x] -> Either String (NonEmpty Pos)

instance GL.TypeError ( 'GL.Text "ValidateNestedNonEmptyC: not defined at 'Z") => ValidateNestedNonEmptyC x 'Z where
  validateNestedNonEmptyC = compileError "validateNestedNonEmptyC: ValidateNestedNonEmptyC x 'Z"
instance ValidateNestedNonEmptyC x ( 'S 'Z) where
  validateNestedNonEmptyC ixes _ _ =
    case ixes of
      i : is -> Right (i :| is)
      [] -> programmError "ValidateNestedNonEmptyC: ('S 'Z): empty list of indices"
instance ValidateNestedNonEmptyC x ( 'S zs) => ValidateNestedNonEmptyC (NonEmpty x) ( 'S ( 'S zs)) where
  validateNestedNonEmptyC ixes x@(n :| ns) xs =
    let cs = map clOrdering $ compareLengths (x :| xs)
     in if all (Just EQ ==) cs
          then
            let zs = ns <> concatMap N.toList xs
             in validateNestedNonEmptyC @x @( 'S zs) (ixes `snocL` lengthP x) n zs
          else Left $ "validateNestedNonEmptyC: lengths=" ++ show (map length (x : xs)) ++ " ixes=" ++ show (map unP ixes)

-- | extracts the dimensions of a nested list: doesnt allow empty dimensions
type ValidateNestedListC :: Type -> Peano -> Constraint
class ValidateNestedListC x y where
  validateNestedListC :: [Pos] -> x -> [x] -> Either String (NonEmpty Pos)

instance GL.TypeError ( 'GL.Text "ValidateNestedListC: not defined at 0") => ValidateNestedListC x 'Z where
  validateNestedListC = compileError "validateNestedListC: ValidateNestedListC x 'Z"
instance ValidateNestedListC x ( 'S 'Z) where
  validateNestedListC ixes _ _ =
    case ixes of
      i : is -> Right (i :| is)
      [] -> programmError "ValidateNestedListC: ('S 'Z): empty list of indices"
instance ValidateNestedListC x ( 'S n) => ValidateNestedListC [x] ( 'S ( 'S n)) where
  validateNestedListC ixes [] _ = Left $ "validateNestedListC: ixes=" ++ show ixes ++ ":no data!"
  validateNestedListC ixes x@(n : ns) xs =
    let cs = map clOrdering $ compareLengths (x :| xs)
     in if all (Just EQ ==) cs
          then
            let zs = ns <> concat xs
             in validateNestedListC @x @( 'S n) (ixes `snocL` lengthP (n :| ns)) n zs
          else Left $ "validateNestedListC: lengths=" ++ show (map length (x : xs)) ++ " ixes=" ++ show ixes

-- | peano numbers for converting between 'Nat' and peano
data Peano = Z | S !Peano deriving stock (Ord, Show, Eq)

-- | convert Nat to Peano
type NatToPeanoT :: Nat -> Peano
type family NatToPeanoT n where
  NatToPeanoT 0 = 'Z
  NatToPeanoT n = 'S (NatToPeanoT (n GN.- 1))

-- | convert Peano to Nat
type PeanoToNatT :: Peano -> Nat
type family PeanoToNatT n where
  PeanoToNatT 'Z = 0
  PeanoToNatT ( 'S n) = 1 GN.+ PeanoToNatT n

-- | convert a matrix index into nested lists
type ListNST :: NonEmpty Nat -> Type -> Type
type family ListNST ns a where
  ListNST (_ ':| '[]) a = [a]
  ListNST (_ ':| n1 ': ns) a = [ListNST (n1 ':| ns) a]

-- | convert a matrix index into nested lists
type NonEmptyNST :: NonEmpty Nat -> Type -> Type
type family NonEmptyNST ns a where
  NonEmptyNST (_ ':| '[]) a = NonEmpty a
  NonEmptyNST (_ ':| n1 ': ns) a = NonEmpty (NonEmptyNST (n1 ':| ns) a)

-- | convert a nested nonempty list into a nested list
nestedNonEmptyToList :: forall ns a. NestedListC ns => NonEmptyNST ns a -> Either String (ListNST ns a)
nestedNonEmptyToList = nestedNonEmptyToListC @ns (Proxy @a)

-- | convert a nested list into a nested nonempty list
nestedListToNonEmpty :: forall ns a. NestedListC ns => ListNST ns a -> Either String (NonEmptyNST ns a)
nestedListToNonEmpty = nestedListToNonEmptyC @ns @_ @a Proxy

-- | methods for working with nested lists
type NestedListC :: NonEmpty Nat -> Constraint
class NestedListC ns where
  -- | convert a nested list to a nested nonempty list
  nestedListToNonEmptyC :: proxy a -> ListNST ns a -> Either String (NonEmptyNST ns a)

  -- | convert a nested nonempty list to a nested list
  nestedNonEmptyToListC :: proxy a -> NonEmptyNST ns a -> Either String (ListNST ns a) -- need a proxy to make it work and find the correct 'a'

  flattenNestedListC :: proxy a -> ListNST ns a -> Either String [a]

instance PosT n => NestedListC (n ':| '[]) where
  nestedListToNonEmptyC _ = \case
    [] -> Left "nestedListToNonEmptyC 'SZ no data"
    x : xs -> lmsg "nestedListToNonEmptyC 'SZ" $ lengthExact1 (fromNP @n) (x :| xs)
  nestedNonEmptyToListC _ lst = N.toList <$> lmsg "nestedNonEmptyToListC 'SZ" (lengthExact1 (fromNP @n) lst)
  flattenNestedListC _ = \case
    [] -> Left "flattenNestedListC 'SZ no data"
    x : xs -> lmsg "flattenNestedListC 'SZ" $ lengthExact (fromN @n) (x : xs)

instance (PosT n, NestedListC (n1 ':| ns)) => NestedListC (n ':| n1 ': ns) where
  nestedListToNonEmptyC p = \case
    [] -> Left "nestedListToNonEmptyC 'SS no data"
    x : xs -> do
      ys <- lmsg "nestedListToNonEmptyC 'SS" $ lengthExact1 (fromNP @n) (x :| xs)
      traverse (nestedListToNonEmptyC @(n1 ':| ns) p) ys
  nestedNonEmptyToListC p lst = do
    xs <- lmsg "nestedNonEmptyToListC 'SS" $ lengthExact1 (fromNP @n) lst
    N.toList <$> traverse (nestedNonEmptyToListC @(n1 ':| ns) p) xs
  flattenNestedListC p = \case
    [] -> Left "flattenNestedListC 'SS no data"
    x : xs -> do
      ys <- lmsg "flattenNestedListC 'SS" $ lengthExact (fromN @n) (x : xs)
      concat <$> traverse (flattenNestedListC @(n1 ':| ns) p) ys

-- mapM_ (putStrLn . genListTupleT) [2..20]  -- to generate from two onwards

-- | translates a type of size "n" to a tuple of size "n"
type ListTupleT :: Nat -> Type -> Type
type family ListTupleT n a = result | result -> n a where
  ListTupleT 1 a = One a
  ListTupleT 2 a = (a, a)
  ListTupleT 3 a = (a, a, a)
  ListTupleT 4 a = (a, a, a, a)
  ListTupleT 5 a = (a, a, a, a, a)
  ListTupleT 6 a = (a, a, a, a, a, a)
  ListTupleT 7 a = (a, a, a, a, a, a, a)
  ListTupleT 8 a = (a, a, a, a, a, a, a, a)
  ListTupleT 9 a = (a, a, a, a, a, a, a, a, a)
  ListTupleT 10 a = (a, a, a, a, a, a, a, a, a, a)
  ListTupleT 11 a = (a, a, a, a, a, a, a, a, a, a, a)
  ListTupleT 12 a = (a, a, a, a, a, a, a, a, a, a, a, a)
  ListTupleT 13 a = (a, a, a, a, a, a, a, a, a, a, a, a, a)
  ListTupleT 14 a = (a, a, a, a, a, a, a, a, a, a, a, a, a, a)
  ListTupleT 15 a = (a, a, a, a, a, a, a, a, a, a, a, a, a, a, a)
  ListTupleT 16 a = (a, a, a, a, a, a, a, a, a, a, a, a, a, a, a, a)
  ListTupleT 17 a = (a, a, a, a, a, a, a, a, a, a, a, a, a, a, a, a, a)
  ListTupleT 18 a = (a, a, a, a, a, a, a, a, a, a, a, a, a, a, a, a, a, a)
  ListTupleT 19 a = (a, a, a, a, a, a, a, a, a, a, a, a, a, a, a, a, a, a, a)
  ListTupleT 20 a = (a, a, a, a, a, a, a, a, a, a, a, a, a, a, a, a, a, a, a, a)

-- | generates a nonempty list of indices using each digit of the given 'Nat'
type NN :: Nat -> NonEmpty Nat
type NN n = NS (NN' '[] n)

-- | generates a list of indices using the individual digits of the given 'Nat'
type NN' :: [Nat] -> Nat -> [Nat]
type family NN' ns n where
  NN' ns 0 = ns
  NN' ns n = NN' (GN.Mod n 10 ': ns) (GN.Div n 10)

-- | matrix dimension of degree 1
type D1 :: Nat -> NonEmpty Nat
type D1 a = a ':| '[]

-- | matrix dimension of degree 2
type D2 :: Nat -> Nat -> NonEmpty Nat
type D2 a b = a ':| '[b]

-- | matrix dimension of degree 3
type D3 :: Nat -> Nat -> Nat -> NonEmpty Nat
type D3 a b c = a ':| '[b, c]

-- | matrix dimension of degree 4
type D4 :: Nat -> Nat -> Nat -> Nat -> NonEmpty Nat
type D4 a b c d = a ':| '[b, c, d]

-- | matrix dimension of degree 5
type D5 :: Nat -> Nat -> Nat -> Nat -> Nat -> NonEmpty Nat
type D5 a b c d e = a ':| '[b, c, d, e]

-- | matrix dimension of degree 6
type D6 :: Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> NonEmpty Nat
type D6 a b c d e f = a ':| '[b, c, d, e, f]

-- | matrix dimension of degree 7
type D7 :: Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> NonEmpty Nat
type D7 a b c d e f g = a ':| '[b, c, d, e, f, g]

-- | matrix dimension of degree 8
type D8 :: Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> NonEmpty Nat
type D8 a b c d e f g h = a ':| '[b, c, d, e, f, g, h]

-- | matrix dimension of degree 9
type D9 :: Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> NonEmpty Nat
type D9 a b c d e f g h i = a ':| '[b, c, d, e, f, g, h, i]

-- | matrix dimension of degree 10
type D10 :: Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> NonEmpty Nat
type D10 a b c d e f g h i j = a ':| '[b, c, d, e, f, g, h, i, j]