-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | Defines lists with length fixed on the type level
module Morley.Util.SizedList
  ( -- * Base types
    SizedList
  , SizedList'(.., (::<), Nil')
  , SomeSizedList(..)

  -- * Utility type synonyms (re-exports)
  , SingIPeano
  , IsoNatPeano

  -- * Basic
  , append
  , reverse
  , head
  , tail

  -- * Construction
  , withNonEmpty
  , withList
  , fromList
  , fromListMaybe
  , fromListMaybe'
  , unsafeFromList
  , generate
  , generate'
  , replicate
  , replicate'
  , replicateT
  , singleton

  -- * Conversion
  , toNonEmpty

  -- * Zips
  , zipWith
  , unzipWith
  , zip
  , unzip

  -- * Index access
  , index'
  , index
  , indexMaybe
  , length'

  -- * Sublists
  , take
  , drop
  , splitAt
  ) where

import Prelude hiding
  (drop, fromList, head, replicate, reverse, splitAt, tail, take, unzip, zip, zipWith)
import Prelude qualified (fromList)

import Data.Constraint ((\\))
import Data.List qualified as List
import Data.Singletons (SingI(..))
import Fmt (Buildable(..))
import GHC.TypeNats (Nat)
import Text.Show (Show(..), showParen, showString)

import Morley.Util.Peano

-- $setup
--
-- >>> import Prelude hiding (drop, fromList, head, replicate, reverse, splitAt, tail, take, unzip, zip, zipWith)
-- >>> import Morley.Util.Peano
-- >>> import Data.Constraint ((\\))

-- | The primary fixed-size list type. Parametrized by a type-level 'Nat' as length and type
-- as element type.
--
-- Internally powered by 'Peano' numbers
type SizedList (n :: Nat) a = SizedList' (ToPeano n) a

-- | Actual fixed-size list GADT, parametrized by 'Peano' natural. You generally don't want to
-- use this directly, since 'Peano' is not very ergonomic. Prefer using 'SizedList' unless
-- writing utility functions that need to do type-level arithmetic.
--
-- Note that while this has the usual instances, 'Applicative' and 'Monad' are not the same
-- as for regular lists: they implement "zipper" semantics, i.e. @f \<*\> x@ is the same as
-- @zipWith ($) f x@
data SizedList' (n :: Peano) a where
  Nil :: SizedList' 'Z a -- ^ Empty list
  (:<) :: ~a -> ~(SizedList' n a) -> SizedList' ('S n) a -- ^ Cons
infixr 5 :<

-- | Sized list cons pattern. Unlike ':<' this pattern can be used to auto-deduce the type of
-- the result, e.g.
--
-- >>> a ::< b ::< c ::< Nil' = pure 'a'
-- >>> (a, b, c)
-- ('a','a','a')
infixr 5 ::<
pattern (::<) :: a -> SizedList' n a -> SizedList' ('S n) a
pattern a ::< l <- a :< l
{-# COMPLETE (::<) #-}

-- | Sized list Nil pattern. Unlike 'Nil' this pattern can be used to auto-deduce the type of
-- the result, see '::<'.
pattern Nil' :: SizedList' 'Z a
pattern Nil' <- Nil
{-# COMPLETE Nil' #-}

deriving stock instance Eq a => Eq (SizedList' n a)
deriving stock instance Foldable (SizedList' n)
deriving stock instance Traversable (SizedList' n)
deriving stock instance Functor (SizedList' n)
instance Show a => Show (SizedList' n a) where
  showsPrec d (x :< xs) = showParen (d > prec) $
    showsPrec (prec + 1) x . showString " :< " . showsPrec prec xs
    where prec = 5
  showsPrec _ Nil = showString "Nil"
deriving stock instance Ord a => Ord (SizedList' n a)

instance Buildable a => Buildable (SizedList' n a) where
  build = build . toList

instance Container (SizedList' n a)

instance SingI n => Applicative (SizedList' n) where
  pure = replicate' (sing @n)
  (<*>) = zipWith ($) \\ minIdempotency (sing @n)

instance SingI n => Monad (SizedList' n) where
  f >>= (k :: a -> SizedList' n b) = go SZ f
    where
      go :: forall k m. (AddPeano k m ~ n)
         => SingNat k -> SizedList' m a -> SizedList' m b
      go _ Nil = Nil
      go i (el :< (rest :: SizedList' m' a)) =
        index' i (k el) :< go (SS i) rest
        \\ associativity @m' i \\ additivity @k @m' @n i

-- | Existential capturing a fixed-size list whose length is only known at runtime.
--
-- In most cases, it's probably better to use regular lists, but this can be occasionally
-- useful.
--
-- We do not provide the 'Applicative' and 'Monad' instances, since "zipper" applicative is
-- ill-defined for lists of different length, and having inconsistent instances between
-- this and 'SizedList' is more confusing than useful.
--
-- Unlike regular sized list, @SomeSizedList@ is a 'Semigroup' and a 'Monoid':
--
-- >>> fromList "ab" <> fromList "de" <> mempty :: SomeSizedList Char
-- SomeSizedList (SS (SS (SS (SS SZ)))) ('a' :< 'b' :< 'd' :< 'e' :< Nil)
data SomeSizedList a where
  SomeSizedList :: SingNat n -> SizedList' n a -> SomeSizedList a

deriving stock instance Foldable SomeSizedList
deriving stock instance Traversable SomeSizedList
deriving stock instance Functor SomeSizedList
deriving stock instance Show a => Show (SomeSizedList a)

instance Semigroup (SomeSizedList a) where
  (SomeSizedList sl l) <> (SomeSizedList sr r) = SomeSizedList (peanoSingAdd sl sr) $ l `append` r

instance Monoid (SomeSizedList a) where
  mempty = SomeSizedList SZ Nil

instance Buildable a => Buildable (SomeSizedList a) where
  build = build . toList

instance Container (SomeSizedList a)

instance (n ~ 'S 'Z) => One (SizedList' n a) where
  type OneItem (SizedList' n a) = a
  one = singleton

instance FromList (SomeSizedList a) where
  type ListElement (SomeSizedList a) = a

  fromList = fromList

-- A type-specific `fromList` for use in qualified imports
fromList :: forall a. [a] -> SomeSizedList a
fromList [] = SomeSizedList SZ Nil
fromList (x:xs) = case fromList xs of
  SomeSizedList n xs' -> SomeSizedList (SS n) (x :< xs')

-- | Try to make a fixed-size list from a regular list, given a 'Nat'. Returns 'Nothing' if
-- the regular list has incorrect length.
--
-- >>> fromListMaybe @3 [1, 2, 3]
-- Just (1 :< 2 :< 3 :< Nil)
--
-- >>> fromListMaybe @4 [1, 2, 3]
-- Nothing
fromListMaybe :: forall n n' a. (SingIPeano n, IsoNatPeano n n') => [a] -> Maybe (SizedList n a)
fromListMaybe = fromListMaybe' $ sing @n'

-- | Same as 'fromListMaybe', but accepts an explicit 'Peano' singleton.
fromListMaybe' :: SingNat m -> [a] -> Maybe (SizedList' m a)
fromListMaybe' SZ [] = Just Nil
fromListMaybe' (SS n) (x:xs) = (x :<) <$> fromListMaybe' n xs
fromListMaybe' _ _ = Nothing

-- | Run some computation with a 'NonEmpty' list converted to 'SizedList'. Similar
-- to pattern-matching on 'SomeSizedList', but asserts on the type level that list is
-- in fact not empty.
withNonEmpty
  :: forall a r. NonEmpty a
  -> (forall n. SingNat ('S n) -> SizedList' ('S n) a -> r)
  -> r
withNonEmpty (x :| xs) f = case fromList xs of
  SomeSizedList n xs' -> f (SS n) $ x :< xs'

-- | Run some computation with a list converted to 'SizedList'. The same as
-- pattern-matching on 'SomeSizedList'.
withList :: forall a r. [a] -> (forall n. SingNat n -> SizedList' n a -> r) -> r
withList xs f = case fromList xs of
  SomeSizedList n xs' -> f n xs'

-- | Construct a list of given length from a regular list. Raise error
-- if length is inconsistent.
--
-- >>> unsafeFromList @3 [1, 2, 3]
-- 1 :< 2 :< 3 :< Nil
--
-- >>> unsafeFromList @1 [1, 2, 3]
-- *** Exception: Invalid list size in Morley.Util.SizedList.unsafeFromList
-- ...
unsafeFromList
  :: forall n n' a. (SingIPeano n, IsoNatPeano n n', HasCallStack)
  => [a] -> SizedList n a
unsafeFromList = unsafeFromList' $ sing @n'

-- | Same as 'unsafeFromList', but accepts explicit 'Peano' singleton
unsafeFromList' :: forall n a. HasCallStack => SingNat n -> [a] -> SizedList' n a
unsafeFromList' sg xs =
  fromMaybe (error "Invalid list size in Morley.Util.SizedList.unsafeFromList") $
  fromListMaybe' sg xs

-- | Convert a 'SizedList' to 'NonEmpty'
toNonEmpty :: SizedList' ('S n) a -> NonEmpty a
toNonEmpty (x :< xs) = x :| toList xs

-- | Replicate a value @n@ times. This is a version
-- returning an existential. You probably want
-- 'replicateT' or `replicate'`.
replicate :: Natural -> a -> SomeSizedList a
replicate 0 _ = SomeSizedList SZ Nil
replicate n x = case replicate (n - 1) x of
  SomeSizedList m xs -> SomeSizedList (SS m) (x :< xs)

-- | Replicate a value @n@ times, where @n@ is passed as a type-level natural.
-- This is essentially a synonym for 'pure'.
--
-- >>> replicateT @5 'a'
-- 'a' :< 'a' :< 'a' :< 'a' :< 'a' :< Nil
replicateT :: forall n n' a. (SingIPeano n, IsoNatPeano n n') => a -> SizedList n a
replicateT = pure

-- | As 'replicateT', but accepts an explicit 'Peano' singleton.
replicate' :: SingNat n -> a -> SizedList' n a
replicate' SZ _ = Nil
replicate' (SS n) x = x :< replicate' n x

-- | Zip two lists of potentially different lengths. If one list is
-- longer, extraneous elements will be ignored.
--
-- >>> zip (unsafeFromList @3 "abc") (unsafeFromList @4 "defg")
-- ('a','d') :< ('b','e') :< ('c','f') :< Nil
zip :: SizedList' n a -> SizedList' m b -> SizedList' (MinPeano n m) (a, b)
zip = zipWith (,)

-- | Zip using a binary operation.
zipWith :: (a -> b -> c) -> SizedList' n a -> SizedList' m b -> SizedList' (MinPeano n m) c
zipWith _ Nil _ = Nil
zipWith _ _ Nil = Nil
zipWith f (x :< xs) (y :< ys) = f x y :< zipWith f xs ys

-- | Given a type-level 'Nat' and a generator function, make a 'SizedList'
--
-- Indexing starts at @0@
--
-- >>> generate @3 (+1)
-- 1 :< 2 :< 3 :< Nil
generate :: forall n n' a. (SingIPeano n, IsoNatPeano n n') => (Natural -> a) -> SizedList' n' a
generate f = generate' (sing @n') (f . snd)

-- | Same as 'generate', but accepts an explicit 'Peano' singleton, and passes an explicit
-- singleton as index to the generator function
generate' :: forall n a. SingNat n
          -> (forall m. (n > m ~ 'True) => (SingNat m, Natural) -> a)
          -> SizedList' n a
generate' n f = go 0 SZ n
  where
  go :: forall k m. (AddPeano k m ~ n)
     => Natural -> SingNat k -> SingNat m -> SizedList' m a
  go _ _ SZ = Nil
  go i k (SS (m :: SingNat m')) =
    f (k, i) :< go (i + 1) (SS k) m \\ associativity @m' k \\ additivity @k @m' @n k


-- | Reverse a 'SizedList'
--
-- >>> reverse $ generate @3 (+1)
-- 3 :< 2 :< 1 :< Nil
reverse :: forall n a. SingI n => SizedList' n a -> SizedList' n a
reverse = unsafeFromList' (sing @n) . List.reverse . toList

-- | Append two sized lists.
--
-- >>> generate @3 (+1) `append` generate @3 (+11)
-- 1 :< 2 :< 3 :< 11 :< 12 :< 13 :< Nil
append :: SizedList' n a -> SizedList' m a -> SizedList' (AddPeano n m) a
append Nil ys = ys
append (x :< xs) ys = x :< append xs ys
infixr 5 `append`

-- | Unzip a sized list using a function
unzipWith :: (a -> (b, c)) -> SizedList' n a -> (SizedList' n b, SizedList' n c)
unzipWith _ Nil = (Nil, Nil)
unzipWith f (a :< as) =
  let (b, c) = f a
      (bs, cs) = unzipWith f as
  in (b :< bs, c :< cs)

-- | Unzip a sized list of pairs
unzip :: SizedList' n (b, c) -> (SizedList' n b, SizedList' n c)
unzip = unzipWith id

-- | Index into a list using a type-level constant
--
-- >>> index @2 $ unsafeFromList @5 [1..5]
-- 3
index
  :: forall i m a. (m > ToPeano i ~ 'True, SingIPeano i)
  => SizedList' m a -> a
index = index' $ peanoSing @i

-- | Same as 'index', but accepts an explicit 'Peano' singleton
--
-- Note, that if you wish to use this with `generate'`, indexing into some
-- other list, for one, you may want to rethink your approach, since it's
-- pretty inefficient.
-- For two, you may run into arcane compiler error messages, e.g.
--
-- >>> let someList = 'a' :< 'b' :< 'c' :< Nil
-- >>> generate' @_ @Char (peanoSing @2) $ \(sg, _) -> index' sg someList
-- ...
-- ... Could not deduce (('S ('S ('S 'Z)) > m) ~ 'True)
-- ...
--
-- To make this work, use 'transitivity' proof:
--
-- >>> let slLen = length' someList
-- >>> let len = peanoSing @2
-- >>> generate' @_ @Char len $ \(sg, _) -> index' sg someList \\ transitivity slLen len sg
-- 'a' :< 'b' :< Nil
index' :: (m > n ~ 'True) => SingNat n -> SizedList' m a -> a
index' SZ (x :< _)= x
index' (SS n) (_ :< xs) = index' n xs

-- | Use a term-level 'Natural' to index into a list. Since the natural can
-- be larger than the length of the list, returns a 'Maybe'.
indexMaybe :: Natural -> SizedList' m a -> Maybe a
indexMaybe _ Nil = Nothing
indexMaybe 0 (x :< _) = Just x
indexMaybe n (_ :< xs) = indexMaybe (n - 1) xs

-- | Get the first element of a non-empty list
head :: SizedList' ('S n) a -> a
head (x :< _) = x

-- | Get elements after the first
tail :: SizedList' ('S n) a -> SizedList' n a
tail (_ :< xs) = xs

-- | Construct a list of one element
singleton :: a -> SizedList 1 a
singleton x = x :< Nil

-- | Take a fixed number of elements, given by a type-level 'Nat'. Must be
-- smaller than the size of the list
--
-- >>> take @3 $ unsafeFromList @5 [1..5]
-- 1 :< 2 :< 3 :< Nil
take
  :: forall n n' m a. (m >= ToPeano n ~ 'True, SingIPeano n, IsoNatPeano n n')
  => SizedList' m a -> SizedList n a
take = fst . splitAt @n

-- | Drop a fixed number of elements, given by a type-level 'Nat'. Must be
-- smaller than the size of the list
--
-- >>> drop @3 $ unsafeFromList @5 [1..5]
-- 4 :< 5 :< Nil
drop
  :: forall n n' m a. (m >= ToPeano n ~ 'True, SingIPeano n, IsoNatPeano n n')
  => SizedList' m a -> SizedList' (SubPeano m (ToPeano n)) a
drop = snd . splitAt @n

-- | @splitAt \@n xs = (take \@n xs, drop \@n xs)@, but traverses the list only
-- once
--
-- >>> splitAt @3 $ unsafeFromList @5 [1..5]
-- (1 :< 2 :< 3 :< Nil,4 :< 5 :< Nil)
splitAt
  :: forall n n' m a. (m >= ToPeano n ~ 'True, SingIPeano n, IsoNatPeano n n')
  => SizedList' m a -> (SizedList n a, SizedList' (SubPeano m n') a)
splitAt = go $ sing @n'
  where
  go :: (l >= k ~ 'True)
     => SingNat k -> SizedList' l a
     -> (SizedList' k a, SizedList' (SubPeano l k) a)
  go SZ xs = (Nil, xs)
  go (SS n) (x :< xs) = first (x :<) $ go n xs

-- | Return a Peano singleton representing the list length
length' :: forall n a. (SingI n) => SizedList' n a -> SingNat n
length' _ = sing @n