{-# LANGUAGE TypeFamilies #-}
module LLVM.Core.UnaryVector (
   T(Cons), vector, cyclicVector, empty, cons, withEmpty, withHead, with, head,
   FixedList, Length, With,
   ) where

import qualified Type.Data.Num.Unary as Unary

import Control.Applicative (Applicative, pure, liftA2, (<*>))

import qualified Data.Traversable as Trav
import qualified Data.NonEmpty as NonEmpty
import qualified Data.Empty as Empty
import Data.Traversable (Traversable, foldMapDefault)
import Data.Foldable (Foldable, foldMap)

import Prelude hiding (replicate, map, head, unzip, zipWith)


newtype T n a = Cons (FixedList n a)

type family FixedList n :: * -> *
type instance FixedList Unary.Zero = Empty.T
type instance FixedList (Unary.Succ n) = NonEmpty.T (FixedList n)

type family Length (f :: * -> *)
type instance Length Empty.T = Unary.Zero
type instance Length (NonEmpty.T f) = Unary.Succ (Length f)


vector ::
   (Unary.Natural n, n ~ Length (FixedList n)) =>
   FixedList n a -> T n a
vector = Cons

cyclicVector ::
   (Unary.Natural n) =>
   NonEmpty.T [] a -> T n a
cyclicVector xt@(NonEmpty.Cons x xs) =
   runOp0 $
   Unary.switchNat
      (Op0 empty)
      (Op0 $ cons x $ cyclicVectorAppend xt xs)

cyclicVectorAppend ::
   (Unary.Natural n) =>
   NonEmpty.T [] a -> [a] -> T n a
cyclicVectorAppend ys xt =
   runOp0 $
   Unary.switchNat
      (Op0 empty)
      (Op0 $
       case xt of
          [] -> cyclicVector ys
          x:xs -> cons x $ cyclicVectorAppend ys xs)

empty :: T Unary.Zero a
empty = Cons Empty.Cons

cons :: a -> T n a -> T (Unary.Succ n) a
cons x (Cons xs) = Cons $ NonEmpty.Cons x xs


withEmpty :: b -> T Unary.Zero a -> b
withEmpty x (Cons Empty.Cons) = x

withHead ::
   (a -> T n a -> b) ->
   T (Unary.Succ n) a -> b
withHead f (Cons (NonEmpty.Cons x xs)) = f x (Cons xs)


newtype Head a n = Head {runHead :: T n a -> a}

head :: (Unary.Positive n) => T n a -> a
head =
   runHead $
   Unary.switchPos
      (Head $ \(Cons (NonEmpty.Cons a _)) -> a)


newtype
   WithVector a b n =
      WithVector {
         runWithVector :: WithRec a b n -> T n a -> b
      }

type family WithRec a b n
type instance WithRec a b Unary.Zero = b
type instance WithRec a b (Unary.Succ n) = a -> WithRec a b n

type With n a b = WithRec a b n

with :: (Unary.Natural n) => With n a b -> T n a -> b
with =
   runWithVector $
   Unary.switchNat
      (WithVector withEmpty)
      (WithVector $ \f v -> withHead (\x -> with (f x)) v)


newtype Op0 a n = Op0 {runOp0 :: T n a}

replicate :: (Unary.Natural n) => a -> T n a
replicate a =
   runOp0 $
   Unary.switchNat
      (Op0 empty)
      (Op0 $ cons a $ replicate a)


newtype Op1 a b n = Op1 {runOp1 :: T n a -> T n b}

map ::
   (Unary.Natural n) =>
   (a -> b) -> T n a -> T n b
map f =
   runOp1 $
   Unary.switchNat
      (Op1 $ withEmpty empty)
      (Op1 $ withHead $ \a -> cons (f a) . map f)


newtype Op2 a b c n = Op2 {runOp2 :: T n a -> T n b -> T n c}

zipWith ::
   (Unary.Natural n) =>
   (a -> b -> c) ->
   T n a -> T n b -> T n c
zipWith f =
   runOp2 $
   Unary.switchNat
      (Op2 $ const $ withEmpty empty)
      (Op2 $ \at bt ->
       withHead (\a as ->
          withHead (\b bs -> cons (f a b) $ zipWith f as bs) bt) at)


newtype
   Sequence f a n =
      Sequence {runSequence :: T n (f a) -> f (T n a)}


instance (Unary.Natural n) => Functor (T n) where
   fmap = map

instance (Unary.Natural n) => Applicative (T n) where
   pure = replicate
   f <*> a = zipWith ($) f a

instance (Unary.Natural n) => Foldable (T n) where
   foldMap = foldMapDefault

instance (Unary.Natural n) => Traversable (T n) where
   sequenceA =
      runSequence $
      Unary.switchNat
         (Sequence $ withEmpty $ pure empty)
         (Sequence $ withHead $ \x xs -> liftA2 cons x $ Trav.sequenceA xs)