module Data.FastVect.Index (Index, reflectInt, index, (!!), widen, parse, modulo) where

import Prelude

import Control.Monad.ST as ST
import Data.Array.ST as ArrayST
import Data.Array.ST.Partial as ArraySTPartial
import Data.Enum (class BoundedEnum, class Enum, Cardinality(..), toEnum)
import Data.FastVect.FastVect (Vect)
import Data.FastVect.FastVect as FastVect
import Data.Maybe (Maybe(..))
import Data.Reflectable (class Reflectable, reflectType)
import Partial.Unsafe (unsafeCrashWith, unsafePartial)
import Prelude as Prelude
import Prim.Int (class Add, class Compare)
import Prim.Ordering (GT, LT)
import Prim.TypeError as TypeError
import Safe.Coerce (coerce)
import Type.Prelude (Proxy(..))

-- | An int between 0 (inclusive) and bound (exclusive)
newtype Index :: Int -> Type
newtype Index bound = Index Int
type role Index nominal

derive newtype instance Eq (Index bound)
derive newtype instance Ord (Index bound)
instance Reflectable bound Int => Show (Index bound) where
  show (Index idx) =
    "Index " <> show idx <> " within bound: " <> show (reflectType (Proxy @bound))

infixr 2 type TypeError.Beside as <>
infixr 1 type TypeError.Above as |>
type ZeroErrorMsg :: Symbol -> Symbol -> TypeError.Doc
type ZeroErrorMsg className closing =
         TypeError.Text "Cannot create " <> TypeError.QuoteLabel className
      <> TypeError.Text " instance for " <> TypeError.Quote (Index 0)
      |> TypeError.Text "A collection of size 0 has no elements to index into; "
      <> TypeError.Text "equivalently, the type " <> TypeError.Quote (Index 0)
      <> TypeError.Text "is uninhabited (i.e. cannot be instantiated like " <> TypeError.Quote Prelude.Void
      <> TypeError.Text ")."
      <> TypeError.Text closing

instance semiringIndexZeroError ::
    TypeError.Fail (ZeroErrorMsg "Semiring" "As such, the `one` and `zero` type class members cannot exist")
    => Semiring (Index 0) where
  zero = unsafeCrashWith "absurd"
  one = unsafeCrashWith "absurd"
  add _ _ = unsafeCrashWith "absurd"
  mul _ _ = unsafeCrashWith "absurd"
else instance semiringIndexDegenerate :: Semiring (Index 1) where
  zero = Index 0
  one = Index 0
  add _ _ = zero
  mul _ _ = zero
else instance semiringIndexGt1 :: (Compare bound 1 GT, Reflectable bound Int) =>
  Semiring (Index bound) where
  zero = Index 0
  one = Index 1
  add (Index a) (Index b) = (Index $ (a + b) `mod` (reflectType (Proxy @bound)))
  mul (Index a) (Index b) = (Index $ (a * b) `mod` (reflectType (Proxy @bound)))

instance (Compare bound 0 GT, Reflectable bound Int) => Bounded (Index bound) where
  top = Index $ (reflectType $ Proxy @bound) - 1
  bottom = Index 0

instance (Compare bound 0 GT, Reflectable bound Int) => Enum (Index bound) where
  succ (Index idx) =
    if idx < (reflectType $ Proxy @bound) - 1 then
      Just (Index $ idx + 1)
    else
      Nothing
  pred (Index idx) =
    if idx > 0 then
      Just (Index $ idx - 1)
    else
      Nothing

instance (Compare bound 0 GT, Reflectable bound Int) => BoundedEnum (Index bound) where
  cardinality = Cardinality 0
  toEnum int =
    if 0 <= int && int < (reflectType $ Proxy @bound) then
      Just (Index int)
    else
      Nothing
  fromEnum (Index int) = int

parse :: forall (@bound :: Int).
         Compare bound 0 GT => Reflectable bound Int =>

         Int -> Maybe (Index bound)
parse = toEnum @(Index bound)

modulo :: forall (@bound :: Int).
          Compare bound 1 GT => Reflectable bound Int =>

          Int -> Index bound
modulo idx = Index (idx `mod` (reflectType $ Proxy @bound))

-- | Reflect a type-level integer into a valid Index
-- | Note: with how `at` is implemented, the created `Index`
-- | is valid for any `Vect n a` where `@int < n`
reflectInt :: forall (@int :: Int) (bound :: Int).
              Reflectable int Int =>
              Compare int (-1) GT =>
              Add int 1 bound =>

              Index bound
reflectInt = Index (reflectType (Proxy @int))

widen :: forall (@widenedBound :: Int) (currentBound :: Int) (currentBoundMinusOne :: Int).
        Add currentBound (-1) currentBoundMinusOne =>
        Compare widenedBound currentBoundMinusOne GT =>

        Index currentBound -> Index widenedBound
widen = coerce

index :: forall (idxBound :: Int) (vectSize :: Int) (idxBoundMinusOne :: Int) a.
     Compare idxBound 0 GT => Compare vectSize 0 GT =>
     Add idxBound (-1) idxBoundMinusOne =>
     Compare idxBoundMinusOne vectSize LT =>

     Index idxBound -> Vect vectSize a -> a
index (Index idx) vec =
  ST.run
    (ArrayST.unsafeThaw (FastVect.toArray vec)
      >>= (unsafePartial ArraySTPartial.peek idx))

infixl 8 index as !!