content: {-# OPTIONS_GHC -Wall #-}
  
  {-# LANGUAGE ConstraintKinds       #-}
  {-# LANGUAGE DataKinds             #-}
  {-# LANGUAGE DeriveGeneric         #-}
  {-# LANGUAGE DeriveTraversable     #-}
  {-# LANGUAGE FlexibleContexts      #-}
  {-# LANGUAGE FlexibleInstances     #-}
  {-# LANGUAGE GADTs                 #-}
  {-# LANGUAGE InstanceSigs          #-}
  {-# LANGUAGE MagicHash             #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE OverloadedLists       #-}
  {-# LANGUAGE PartialTypeSignatures #-}
  {-# LANGUAGE PolyKinds             #-}
  {-# LANGUAGE RankNTypes            #-}
  {-# LANGUAGE ScopedTypeVariables   #-}
  {-# LANGUAGE StandaloneDeriving    #-}
  {-# LANGUAGE TypeFamilies          #-}
  {-# LANGUAGE TypeOperators         #-}
  {-# LANGUAGE UndecidableInstances  #-}
  
  {-# LANGUAGE DeriveAnyClass        #-}
  
  module Naperian where
  
  import qualified Prelude
  import           Prelude hiding      ( lookup, length, replicate, zipWith )
  
  import qualified Data.IntMap         as IntMap
  import           Data.List           ( intercalate )
  import           Data.Kind           ( Type, Constraint )
  import           Control.Applicative ( liftA2 )
  import qualified GHC.Exts            as L (IsList(..))
  import           GHC.Prim
  import           GHC.TypeLits
  
  import qualified Data.Vector         as Vector
  
  import           Data.Foldable       ( toList )
  
  import           Data.Maybe ( fromJust )
  
  --------------------------------------------------------------------------------
  -- Miscellaneous
  
  -- | The finite set of type-bounded Naturals. A value of type @'Fin' n@ has
  -- exactly @n@ inhabitants, the natural numbers from @[0..n-1]@.
  data Finite :: Nat -> Type where
    Fin :: Int -> Finite n
    deriving (Eq, Show)
  
  -- | Create a type-bounded finite number @'Fin' n@ from a runtime integer,
  -- bounded to a statically known limit. If the input value @x > n@, then
  -- @'Nothing'@ is returned. Otherwise, returns @'Just' (x :: 'Fin' n)@.
  finite :: forall n. KnownNat n => Int -> Maybe (Finite n)
  finite x = case (x > y) of
    True  -> Nothing
    False -> Just (Fin x)
    where y = fromIntegral (natVal' (proxy# :: Proxy# n))
  
  -- | \"'Applicative' zipping\".
  azipWith :: Applicative f => (a -> b -> c) -> f a -> f b -> f c
  azipWith h xs ys = (pure h <*> xs) <*> ys
  
  -- | Format a vector to make it look nice.
  showVector :: [String] -> String
  showVector xs = "<" ++ intercalate "," xs ++ ">"
  
  --------------------------------------------------------------------------------
  -- Pairs
  
  -- | The cartesian product of @'a'@, equivalent to @(a, a)@.
  data Pair a = Pair a a
    deriving (Show, Eq, Ord, Functor, Foldable, Traversable)
  
  instance Applicative Pair where
    pure a = Pair a a
    Pair k g <*> Pair a b = Pair (k a) (g b)
  
  --------------------------------------------------------------------------------
  -- Vectors
  
  newtype Vector (n :: Nat) a = Vector (Vector.Vector a)
    deriving (Eq, Ord, Functor, Foldable, Traversable)
  
  instance Show a => Show (Vector n a) where
    show = showVector . map show . toList
  
  instance KnownNat n => Applicative (Vector n) where
    pure  = replicate
    (<*>) = zipWith ($)
  
  instance (KnownNat n, Traversable (Vector n)) => L.IsList (Vector n a) where
    type Item (Vector n a) = a
    toList = Data.Foldable.toList
  
    fromList xs = case fromList xs of
      Nothing -> error "Demanded vector of a list that wasn't the proper length"
      Just ys -> ys
  
  tail :: Vector (n + 1) a -> Vector n a
  tail (Vector v) = Vector (Vector.tail v)
  
  fromList :: forall n a. KnownNat n => [a] -> Maybe (Vector n a)
  fromList xs = case (Prelude.length xs == sz) of
    False -> Nothing
    True  -> Just (Vector $ Vector.fromList xs)
    where sz = fromIntegral (natVal' (proxy# :: Proxy# n)) :: Int
  
  zipWith :: (a -> b -> c) -> Vector n a -> Vector n b -> Vector n c
  zipWith f (Vector a) (Vector b) = Vector (Vector.zipWith f a b)
  
  length :: forall n a. KnownNat n => Vector n a -> Int
  length _ = fromIntegral $ natVal' (proxy# :: Proxy# n)
  
  replicate :: forall n a. KnownNat n => a -> Vector n a
  replicate v = Vector (Vector.replicate sz v) where
    sz = fromIntegral (natVal' (proxy# :: Proxy# n)) :: Int
  
  index :: Vector n a -> Finite n -> a
  index (Vector v) (Fin n) = (Vector.!) v n
  
  viota :: forall n. KnownNat n => Vector n (Finite n)
  viota = Vector (fmap Fin (Vector.enumFromN 0 sz)) where
    sz = fromIntegral (natVal' (proxy# :: Proxy# n)) :: Int
  
  --------------------------------------------------------------------------------
  -- Naperian functors
  
  -- | Naperian functors.
  
  -- A useful way of thinking about a Naperian functor is that if we have a value
  -- of type @v :: f a@ for some @'Naperian' f@, then we can think of @f a@ as a
  -- bag of objects, with the ability to pick out the @a@ values inside the bag,
  -- for each and every @a@ inside @f@. For example, in order to look up a value
  -- @a@ inside a list @[a]@, we could use a function @[a] -> Int -> a@, which is
  -- exactly @'(Prelude.!!)'@
  --
  -- The lookup function acts like a logarithm of the @'Functor' f@. Intuitively,
  -- a Haskell function @f :: a -> b@ acts like the exponential @b^a@ if we intuit
  -- types as an algebraic quantity. The logarithm of some value @x = b^a@ is
  -- defined as @log_b(x) = a@, so given @x@ and a base @b@, it finds the exponent
  -- @a@. In Haskell terms, this would be like finding the input value @a@ to a
  -- function @f :: a -> b@, given a @b@, so it is a reverse mapping from the
  -- outputs of @f@ back to its inputs.
  --
  -- A @'Naperian'@ functor @f@ is precisely a functor @f@ such that for any value
  -- of type @f a@, we have a way of finding every single @a@ inside.
  class Functor f => Naperian f where
    {-# MINIMAL lookup, (tabulate | positions) #-}
  
    -- | The \"logarithm\" of @f@. This type represents the 'input' you use to
    -- look up values inside @f a@. For example, if you have a list @[a]@, and
    -- you want to look up a value, then you use an @'Int'@ to index into
    -- the list. In this case, @'Log' [a] = Int@. If you have a type-bounded
    -- Vector @'Vector' (n :: 'Nat') a@, then @'Log' ('Vector' n)@ is the
    -- range of integers @[0..n-1]@ (represented here as @'Finite' n@.)
    type Log f
  
    -- | Look up an element @a@ inside @f a@. If you read this function type in
    -- english, it says \"if you give me an @f a@, then I will give you a
    -- function, so you can look up the elements of @f a@ and get back an @a@\"
    lookup :: f a -> (Log f -> a)
  
    -- | Tabulate a @'Naperian'@. This creates @f a@ values by mapping the logarithm
    -- of @f@ onto every \"position\" inside @f a@
    tabulate :: (Log f -> a) -> f a
    tabulate h = fmap h positions
  
    -- | Find every position in the \"space\" of the @'Naperian' f@.
    positions :: f (Log f)
    positions = tabulate id
  
  -- | The transposition of two @'Naperian'@ functors @f@ and @g@.
  transpose :: (Naperian f, Naperian g) => f (g a) -> g (f a)
  transpose = tabulate . fmap tabulate . flip . fmap lookup . lookup
  
  instance Naperian Pair where
    type Log Pair = Bool
    lookup (Pair x y) b = if b then y else x
  
    positions = Pair False True
  
  instance KnownNat n => Naperian (Vector n) where
    type Log (Vector n) = Finite n
  
    lookup    = index
    positions = viota
  
  instance Naperian [] where
    type Log [] = Int
    lookup = (!!)
    positions = [0..]
  
  data Three = One | Two | Three
    deriving (Read, Show, Eq, Ord)
  
  data Triple a = Triple a a a
    deriving (Show, Functor, Foldable, Traversable)
  
  instance Applicative Triple where
    pure x = Triple x x x
    Triple f g h <*> Triple a b c = Triple (f a) (g b) (h c)
  
  instance Naperian Triple where
    type Log Triple = Three
    lookup (Triple x _ _) One = x
    lookup (Triple _ y _) Two = y
    lookup (Triple _ _ z) Three = z
    positions = Triple One Two Three
  
  instance Dimension Triple
  
  --------------------------------------------------------------------------------
  -- Dimensions
  
  class (Applicative f, Naperian f, Traversable f) => Dimension f where
    size :: f a -> Int
    size = Prelude.length . toList
  
  instance               Dimension Pair       where size = const 2
  instance KnownNat n => Dimension (Vector n) where size = length
  
  inner :: (Num a, Dimension f) => f a -> f a -> a
  inner xs ys = sum (liftA2 (*) xs ys)
  
  tensor11 :: (Num a, Dimension f, Dimension g) => f a -> g a -> f (g a)
  tensor11 xs ys = liftA2 (liftA2 (*)) (fmap pure xs) (pure ys)
  
  trace11 :: (Naperian f, Num a, Foldable f) => f (f a) -> a
  trace11 = sum . diags
  
  diags :: Naperian f => f (f b) -> f b
  diags xx = fmap (\i -> uncurry (lookup . lookup xx) (i, i)) positions
  
  tensor22 :: (Num a, Dimension e, Dimension f, Dimension g, Dimension h) =>
         f (e a) -> g (h a) -> f (e (g (h a)))
  tensor22 xs ys = fmap transpose $ liftA2 (liftA2 (tensor11)) (fmap pure xs) (pure ys)
  
  matrix :: ( Num a
            , Dimension f
            , Dimension g
            , Dimension h
            ) => f (g a)
              -> g (h a)
              -> f (h a)
  matrix xss yss = liftA2 (liftA2 inner) (fmap pure xss) (pure (transpose yss))
  
  --------------------------------------------------------------------------------
  -- Hyper-dimensional stuff
  
  -- | Arbitrary-rank Hypercuboids, parameterized over their dimension.
  data Hyper :: [Type -> Type] -> Type -> Type where
    Scalar :: a -> Hyper '[] a
    Prism  :: (Dimension f, Shapely fs) => Hyper fs (f a) -> Hyper (f : fs) a
  
  point :: Hyper '[] a -> a
  point (Scalar a) = a
  
  crystal :: Hyper (f : fs) a -> Hyper fs (f a)
  crystal (Prism x) = x
  
  instance Show a => Show (Hyper fs a) where
    show = showHyper . fmap show where
      showHyper :: Hyper gs String -> String
      showHyper (Scalar s) = s
      showHyper (Prism x)  = showHyper (fmap (showVector . toList) x)
  
  {--
  class HyperLift f fs where
    hyper :: (Shapely fs, Dimension f) => f a -> Hyper (f : fs) a
  
  instance HyperLift f '[] where
    hyper = Prism . Scalar
  
  instance (Shapely fs, HyperLift f fs) => HyperLift f (f : fs) where
    hyper = Prism . (\x -> (hyper $ _ x))
  --}
  
  class Shapely fs where
    hreplicate :: a -> Hyper fs a
    hsize      :: Hyper fs a -> Int
  
  instance Shapely '[] where
    hreplicate a = Scalar a
    hsize        = const 1
  
  instance (Dimension f, Shapely fs) => Shapely (f : fs) where
    hreplicate a = Prism (hreplicate (pure a))
    hsize (Prism x) = size (first x) * hsize x
  
  instance Functor (Hyper fs) where
    fmap f (Scalar a) = Scalar (f a)
    fmap f (Prism x)  = Prism (fmap (fmap f) x)
  
  instance Shapely fs => Applicative (Hyper fs) where
    pure  = hreplicate
    (<*>) = hzipWith ($)
  
  hzipWith :: (a -> b -> c) -> Hyper fs a -> Hyper fs b -> Hyper fs c
  hzipWith f (Scalar a) (Scalar b) = Scalar (f a b)
  hzipWith f (Prism x)  (Prism y)  = Prism (hzipWith (azipWith f) x y)
  
  first :: Shapely fs => Hyper fs a -> a
  first (Scalar a) = a
  first (Prism x)  = head (toList (first x))
  
  -- | Generalized transposition over arbitrary-rank hypercuboids.
  transposeH :: Hyper (f : (g : fs)) a
             -> Hyper (g : (f : fs)) a
  transposeH (Prism (Prism x)) = Prism (Prism (fmap transpose x))
  
  transposeH' :: Hyper (f : (g : (h : fs))) a
              -> Hyper (f : (h : (g : fs))) a
  transposeH' (Prism (Prism (Prism x))) = Prism (Prism (Prism (fmap transpose x)))
  
  -- | Fold over a single dimension of a Hypercuboid.
  foldrH :: (a -> a -> a) -> a -> Hyper (f : fs) a -> Hyper fs a
  foldrH f z (Prism x) = fmap (foldr f z) x
  
  -- | Lift an unary function from values to hypercuboids of values.
  unary :: Shapely fs => (a -> b) -> (Hyper fs a -> Hyper fs b)
  unary = fmap
  
  -- | Lift a binary function from values to two sets of hypercuboids, which can
  -- be aligned properly.
  binary :: ( Compatible fs gs
            , Max fs gs ~ hs
            , Alignable fs hs
            , Alignable gs hs
            ) => (a -> b -> c)
              -> Hyper fs a
              -> Hyper gs b
              -> Hyper hs c
  binary f x y = hzipWith f (align x) (align y)
  
  up :: (Shapely fs, Dimension f) => Hyper fs a -> Hyper (f : fs) a
  up = Prism . fmap pure
  
  -- | Generalized, rank-polymorphic inner product.
  innerH :: ( Max fs gs ~  (f : hs)
            , Alignable fs (f : hs)
            , Alignable gs (f : hs)
            , Compatible fs gs
            , Num a
            ) => Hyper fs a
              -> Hyper gs a
              -> Hyper hs a
  innerH xs ys = foldrH (+) 0 (binary (*) xs ys)
  
  -- | Generalized, rank-polymorphic matrix product.
  matrixH :: ( Num a
             , Dimension f
             , Dimension g
             , Dimension h
             ) => Hyper '[ g, f ] a
               -> Hyper '[ h, g ] a
               -> Hyper '[ h, f ] a
  matrixH x y = case (crystal x, transposeH y) of
    (xs, Prism (Prism ys)) -> hzipWith inner (up xs) (Prism (up ys))
  
  --------------------------------------------------------------------------------
  -- Alignment
  
  class (Shapely fs, Shapely gs) => Alignable fs gs where
    align :: Hyper fs a -> Hyper gs a
  
  instance Alignable '[] '[] where
    align = id
  
  instance (Dimension f, Alignable fs gs) => Alignable (f : fs) (f : gs) where
    align (Prism x) = Prism (align x)
  
  instance (Dimension f, Shapely fs) => Alignable '[] (f : fs) where
    align (Scalar a) = hreplicate a
  
  type family Max (fs :: [Type -> Type]) (gs :: [Type -> Type]) :: [Type -> Type] where
    Max '[]      '[]      = '[]
    Max '[]      (f : gs) = f : gs
    Max (f : fs) '[]      = f : fs
    Max (f : fs) (f : gs) = f : Max fs gs
  
  type family Compatible (fs :: [Type -> Type]) (gs :: [Type -> Type]) :: Constraint where
    Compatible '[] '[]           = ()
    Compatible '[] (f : gs)      = ()
    Compatible (f : fs) '[]      = ()
    Compatible (f : fs) (f : gs) = Compatible fs gs
    Compatible a b               = TypeError (
           'Text "Mismatched dimensions!"
     ':$$: 'Text "The dimension " ':<>: 'ShowType a ':<>: 'Text " can't be aligned with"
     ':$$: 'Text "the dimension " ':<>: 'ShowType b)
  
  --------------------------------------------------------------------------------
  -- Flattened, sparse Hypercuboids
  
  elements :: Shapely fs => Hyper fs a -> [a]
  elements (Scalar a) = [a]
  elements (Prism a)  = concat (map toList (elements a))
  
  data Flat fs a where
    Flat :: Shapely fs => Vector.Vector a -> Flat fs a
  
  instance Functor (Flat fs) where
    fmap f (Flat v) = Flat (fmap f v)
  
  instance Show a => Show (Flat fs a) where
    show = showHyper . fmap show where
      showHyper :: Flat gs String -> String
      showHyper (Flat v) = showVector (toList v)
  
  flatten :: Shapely fs => Hyper fs a -> Flat fs a
  flatten hs = Flat (Vector.fromList (elements hs))
  
  data Sparse fs a where
    Sparse :: Shapely fs => a -> IntMap.IntMap a -> Sparse fs a
  
  unsparse :: forall fs a. Shapely fs => Sparse fs a -> Flat fs a
  unsparse (Sparse e xs) = Flat (Vector.unsafeAccum (flip const) vs as)
    where
      as     = IntMap.assocs xs
      vs     = Vector.replicate l e
      l      = hsize (hreplicate () :: Hyper fs ())
  
  --------------------------------------------------------------------------------
  -- Examples
  
  type Matrix n m v = Vector n (Vector m v)
  
  example1 :: Int
  example1 = inner v1 v2 where
    v1 = [ 1, 2, 3 ] :: Vector 3 Int
    v2 = [ 4, 5, 6 ] :: Vector 3 Int
  
  example2 :: Matrix 2 2 Int
  example2 = matrix m1 m2 where
    m1 = [ [ 1, 2, 3 ]
         , [ 4, 5, 6 ]
         ] :: Matrix 2 3 Int
  
    m2 = [ [ 9, 8 ]
         , [ 6, 5 ]
         , [ 3, 2 ]
         ] :: Matrix 3 2 Int
  
  example3 :: Hyper '[] Int
  example3 = innerH v1 v2 where
    v1 = Prism (Scalar [1, 2, 3]) :: Hyper '[Vector 3] Int
    v2 = Prism (Scalar [4, 5, 6]) :: Hyper '[Vector 3] Int
  
  example4 :: Hyper '[Vector 2, Vector 2] Int
  example4 = matrixH v1 v2 where
    x = [ [ 1, 2, 3 ]
        , [ 4, 5, 6 ]
        ] :: Matrix 2 3 Int
  
    y = [ [ 9, 8 ]
        , [ 6, 5 ]
        , [ 3, 2 ]
        ] :: Matrix 3 2 Int
  
    v1 = Prism (Prism (Scalar x)) :: Hyper '[Vector 3, Vector 2] Int
    v2 = Prism (Prism (Scalar y)) :: Hyper '[Vector 2, Vector 3] Int
  
  type Cuboid n m p v = Vector n (Vector m (Vector p v))
  
  uu = [ [ 1, 2, 3 ] ] :: Matrix 1 3 Int
  
  xx = [ [ 1, 2, 3 ]
      , [ 4, 5, 6 ]
      ] :: Matrix 2 3 Int
  
  yy = [ [ 9, 8 ]
      , [ 6, 5 ]
      , [ 3, 2 ]
      ] :: Matrix 3 2 Int
  
  zz :: Cuboid 2 2 3 Int
  zz = [[ [ 1, 2, 3 ]
        , [ 4, 5, 6 ]
        ]
       ,
        [ [ 1, 2, 3 ]
        , [ 4, 5, 6 ]
        ]
       ]
  
  zzH = Prism (Prism (Prism (Scalar zz)))
  