{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}

-- | Harpie-backed finite indices for @Circuit.Mat@.
--
-- Harpie shapes are rectangular products of type-level naturals; there is no
-- native shape constructor for coproducts.  This module provides the minimal
-- bridge: a @KnownNat@-indexed wrapper @F n@ around @Harpie.Shape.Fin n@, with
-- a 'Circuit.Mat.Finite' instance, so that matrices can use harpie-style
-- type-level sizes while reusing the existing 'Mat' machinery.
--
-- The central isomorphism is the standard disjoint-union-as-initial-segment:
--
-- @
-- Either (F i) (F j)  ≅  F (i + j)
-- @
--
-- realised by 'eitherToF' and 'fToEither'.  This lets the biproduct 'parH'
-- embed block-diagonally into a single rectangular harpie dimension, and lets
-- 'traceF' reuse 'traceMat' on a feedback channel whose size is known at the
-- type level.
module Circuit.Mat.Harpie
  ( -- * KnownNat wrapper
    F (..),
    finF,
    safeFinF,
    fromF,
    toF,

    -- * Coproduct isomorphism
    eitherToF,
    fToEither,

    -- * Matrices indexed by F
    matF,
    runMatF,
    parH,
    traceF,
  )
where

import Circuit.Mat (Finite (..), Mat (..), runMat, traceMat)
import GHC.TypeNats (KnownNat, Nat, type (+))
import Harpie.Shape (Fin (..), fin, safeFin, valueOf)
import NumHask.Algebra.Additive (Additive (..))
import NumHask.Algebra.Multiplicative (Multiplicative (..))
import NumHask.Algebra.Ring (StarSemiring (..))
import Prelude hiding (id, sum, (*), (+))

-- $setup
--
-- >>> :set -XDataKinds
-- >>> :set -XTypeApplications
-- >>> :m -Prelude
-- >>> :set -XRebindableSyntax
-- >>> import NumHask.Prelude
-- >>> import Harpie.Shape (Fin (..))
-- >>> import Circuit.Mat
-- >>> import Circuit.Mat.Harpie

-- | A 'KnownNat'-backed finite index type.
--
-- @F n@ is exactly @Fin n@, but carries a 'Circuit.Mat.Finite' instance
-- derived from the type-level natural.  This is the smallest wrapper that lets
-- 'Mat' treat harpie-style indices as enumerable objects.
newtype F (n :: Nat) = F {forall (n :: Nat). F n -> Fin n
unF :: Fin n}
  deriving stock (F n -> F n -> Bool
(F n -> F n -> Bool) -> (F n -> F n -> Bool) -> Eq (F n)
forall (n :: Nat). F n -> F n -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall (n :: Nat). F n -> F n -> Bool
== :: F n -> F n -> Bool
$c/= :: forall (n :: Nat). F n -> F n -> Bool
/= :: F n -> F n -> Bool
Eq, Eq (F n)
Eq (F n) =>
(F n -> F n -> Ordering)
-> (F n -> F n -> Bool)
-> (F n -> F n -> Bool)
-> (F n -> F n -> Bool)
-> (F n -> F n -> Bool)
-> (F n -> F n -> F n)
-> (F n -> F n -> F n)
-> Ord (F n)
F n -> F n -> Bool
F n -> F n -> Ordering
F n -> F n -> F n
forall (n :: Nat). Eq (F n)
forall (n :: Nat). F n -> F n -> Bool
forall (n :: Nat). F n -> F n -> Ordering
forall (n :: Nat). F n -> F n -> F n
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: forall (n :: Nat). F n -> F n -> Ordering
compare :: F n -> F n -> Ordering
$c< :: forall (n :: Nat). F n -> F n -> Bool
< :: F n -> F n -> Bool
$c<= :: forall (n :: Nat). F n -> F n -> Bool
<= :: F n -> F n -> Bool
$c> :: forall (n :: Nat). F n -> F n -> Bool
> :: F n -> F n -> Bool
$c>= :: forall (n :: Nat). F n -> F n -> Bool
>= :: F n -> F n -> Bool
$cmax :: forall (n :: Nat). F n -> F n -> F n
max :: F n -> F n -> F n
$cmin :: forall (n :: Nat). F n -> F n -> F n
min :: F n -> F n -> F n
Ord)
  deriving newtype (Int -> F n -> ShowS
[F n] -> ShowS
F n -> String
(Int -> F n -> ShowS)
-> (F n -> String) -> ([F n] -> ShowS) -> Show (F n)
forall (n :: Nat). Int -> F n -> ShowS
forall (n :: Nat). [F n] -> ShowS
forall (n :: Nat). F n -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall (n :: Nat). Int -> F n -> ShowS
showsPrec :: Int -> F n -> ShowS
$cshow :: forall (n :: Nat). F n -> String
show :: F n -> String
$cshowList :: forall (n :: Nat). [F n] -> ShowS
showList :: [F n] -> ShowS
Show)

-- | Construct an @F n@ from an 'Int'.  Errors if out of bounds.
--
-- >>> finF @3 2
-- 2
finF :: forall n. (KnownNat n) => Int -> F n
finF :: forall (n :: Nat). KnownNat n => Int -> F n
finF = Fin n -> F n
forall (n :: Nat). Fin n -> F n
F (Fin n -> F n) -> (Int -> Fin n) -> Int -> F n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat). KnownNat n => Int -> Fin n
fin @n

-- | Construct an @F n@ from an 'Int' safely.
--
-- >>> safeFinF @3 2
-- Just 2
-- >>> safeFinF @3 3
-- Nothing
safeFinF :: forall n. (KnownNat n) => Int -> Maybe (F n)
safeFinF :: forall (n :: Nat). KnownNat n => Int -> Maybe (F n)
safeFinF = (Fin n -> F n) -> Maybe (Fin n) -> Maybe (F n)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Fin n -> F n
forall (n :: Nat). Fin n -> F n
F (Maybe (Fin n) -> Maybe (F n))
-> (Int -> Maybe (Fin n)) -> Int -> Maybe (F n)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat). KnownNat n => Int -> Maybe (Fin n)
safeFin @n

-- | Coerce to the underlying @Fin n@.
fromF :: F n -> Fin n
fromF :: forall (n :: Nat). F n -> Fin n
fromF = F n -> Fin n
forall (n :: Nat). F n -> Fin n
unF

-- | Coerce from a @Fin n@.
toF :: Fin n -> F n
toF :: forall (n :: Nat). Fin n -> F n
toF = Fin n -> F n
forall (n :: Nat). Fin n -> F n
F

-- | Enumerate all @F n@ values.
--
-- >>> universe :: [F 3]
-- [0,1,2]
instance (KnownNat n) => Finite (F n) where
  universe :: [F n]
universe = [Fin n -> F n
forall (n :: Nat). Fin n -> F n
F (forall (n :: Nat). KnownNat n => Int -> Fin n
fin @n Int
i) | Int
i <- [Int
0 .. forall (n :: Nat). KnownNat n => Int
valueOf @n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]

-- | Disjoint-union embedding into @F (i + j)@.
--
-- Left summand occupies indices @[0, i)@; right summand occupies @[i, i+j)@.
--
-- >>> eitherToF @2 @3 (Left (finF 1))
-- 1
-- >>> eitherToF @2 @3 (Right (finF 1))
-- 3
eitherToF :: forall i j. (KnownNat i, KnownNat j) => Either (F i) (F j) -> F (i + j)
eitherToF :: forall (i :: Nat) (j :: Nat).
(KnownNat i, KnownNat j) =>
Either (F i) (F j) -> F (i + j)
eitherToF (Left (F (UnsafeFin Int
a))) = Fin (i + j) -> F (i + j)
forall (n :: Nat). Fin n -> F n
F (Int -> Fin (i + j)
forall {k} (s :: k). Int -> Fin s
UnsafeFin Int
a)
eitherToF (Right (F (UnsafeFin Int
b))) = Fin (i + j) -> F (i + j)
forall (n :: Nat). Fin n -> F n
F (Int -> Fin (i + j)
forall {k} (s :: k). Int -> Fin s
UnsafeFin (forall (n :: Nat). KnownNat n => Int
valueOf @i Int -> Int -> Int
forall a. Additive a => a -> a -> a
+ Int
b))

-- | Split an @F (i + j)@ into the left or right summand.
--
-- >>> fToEither @2 @3 (finF @5 1)
-- Left 1
-- >>> fToEither @2 @3 (finF @5 3)
-- Right 1
fToEither :: forall i j. (KnownNat i, KnownNat j) => F (i + j) -> Either (F i) (F j)
fToEither :: forall (i :: Nat) (j :: Nat).
(KnownNat i, KnownNat j) =>
F (i + j) -> Either (F i) (F j)
fToEither (F (UnsafeFin Int
x))
  | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< forall (n :: Nat). KnownNat n => Int
valueOf @i = F i -> Either (F i) (F j)
forall a b. a -> Either a b
Left (Fin i -> F i
forall (n :: Nat). Fin n -> F n
F (Int -> Fin i
forall {k} (s :: k). Int -> Fin s
UnsafeFin Int
x))
  | Bool
otherwise = F j -> Either (F i) (F j)
forall a b. b -> Either a b
Right (Fin j -> F j
forall (n :: Nat). Fin n -> F n
F (Int -> Fin j
forall {k} (s :: k). Int -> Fin s
UnsafeFin (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- forall (n :: Nat). KnownNat n => Int
valueOf @i)))

-- | Build a matrix from a function on underlying @Fin@ indices.
matF ::
  (KnownNat i, KnownNat j) =>
  (Fin i -> Fin j -> s) ->
  Mat s (F i) (F j)
matF :: forall (i :: Nat) (j :: Nat) s.
(KnownNat i, KnownNat j) =>
(Fin i -> Fin j -> s) -> Mat s (F i) (F j)
matF Fin i -> Fin j -> s
f = (F i -> F j -> s) -> Mat s (F i) (F j)
forall i j s. (Finite i, Finite j) => (i -> j -> s) -> Mat s i j
Mat ((F i -> F j -> s) -> Mat s (F i) (F j))
-> (F i -> F j -> s) -> Mat s (F i) (F j)
forall a b. (a -> b) -> a -> b
$ \(F Fin i
i) (F Fin j
j) -> Fin i -> Fin j -> s
f Fin i
i Fin j
j

-- | Run a matrix at underlying @Fin@ indices.
runMatF ::
  (Additive s, Multiplicative s) =>
  Mat s (F i) (F j) ->
  Fin i ->
  Fin j ->
  s
runMatF :: forall s (i :: Nat) (j :: Nat).
(Additive s, Multiplicative s) =>
Mat s (F i) (F j) -> Fin i -> Fin j -> s
runMatF Mat s (F i) (F j)
m Fin i
i Fin j
j = Mat s (F i) (F j) -> F i -> F j -> s
forall s j i.
(Additive s, Multiplicative s, Eq j) =>
Mat s i j -> i -> j -> s
runMat Mat s (F i) (F j)
m (Fin i -> F i
forall (n :: Nat). Fin n -> F n
F Fin i
i) (Fin j -> F j
forall (n :: Nat). Fin n -> F n
F Fin j
j)

-- | Block-diagonal biproduct for harpie-backed indices.
--
-- The result lives on a single rectangular dimension @F (i + k) × F (j + l)@,
-- matching what a harpie @Array @[i + k, j + l] s@ expects.
--
-- >>> let m = matF @2 @2 (\i j -> if i == j then (1 :: Int) else 0)
-- >>> let n = matF @1 @1 (\_ _ -> (7 :: Int))
-- >>> runMat (parH m n) (eitherToF (Left (finF @2 1))) (eitherToF (Left (finF @2 1)))
-- 1
-- >>> runMat (parH m n) (eitherToF (Right (finF @1 0))) (eitherToF (Right (finF @1 0)))
-- 7
-- >>> runMat (parH m n) (eitherToF (Left (finF @2 0))) (eitherToF (Right (finF @1 0)))
-- 0
parH ::
  ( KnownNat i,
    KnownNat j,
    KnownNat k,
    KnownNat l,
    KnownNat (i + k),
    KnownNat (j + l),
    Additive s,
    Multiplicative s
  ) =>
  Mat s (F i) (F j) ->
  Mat s (F k) (F l) ->
  Mat s (F (i + k)) (F (j + l))
parH :: forall (i :: Nat) (j :: Nat) (k :: Nat) (l :: Nat) s.
(KnownNat i, KnownNat j, KnownNat k, KnownNat l, KnownNat (i + k),
 KnownNat (j + l), Additive s, Multiplicative s) =>
Mat s (F i) (F j)
-> Mat s (F k) (F l) -> Mat s (F (i + k)) (F (j + l))
parH Mat s (F i) (F j)
m Mat s (F k) (F l)
n = (F (i + k) -> F (j + l) -> s) -> Mat s (F (i + k)) (F (j + l))
forall i j s. (Finite i, Finite j) => (i -> j -> s) -> Mat s i j
Mat ((F (i + k) -> F (j + l) -> s) -> Mat s (F (i + k)) (F (j + l)))
-> (F (i + k) -> F (j + l) -> s) -> Mat s (F (i + k)) (F (j + l))
forall a b. (a -> b) -> a -> b
$ \F (i + k)
r F (j + l)
c ->
  case (F (i + k) -> Either (F i) (F k)
forall (i :: Nat) (j :: Nat).
(KnownNat i, KnownNat j) =>
F (i + j) -> Either (F i) (F j)
fToEither F (i + k)
r, F (j + l) -> Either (F j) (F l)
forall (i :: Nat) (j :: Nat).
(KnownNat i, KnownNat j) =>
F (i + j) -> Either (F i) (F j)
fToEither F (j + l)
c) of
    (Left F i
r', Left F j
c') -> Mat s (F i) (F j) -> F i -> F j -> s
forall s j i.
(Additive s, Multiplicative s, Eq j) =>
Mat s i j -> i -> j -> s
runMat Mat s (F i) (F j)
m F i
r' F j
c'
    (Right F k
r', Right F l
c') -> Mat s (F k) (F l) -> F k -> F l -> s
forall s j i.
(Additive s, Multiplicative s, Eq j) =>
Mat s i j -> i -> j -> s
runMat Mat s (F k) (F l)
n F k
r' F l
c'
    (Either (F i) (F k), Either (F j) (F l))
_ -> s
forall a. Additive a => a
zero

-- | Trace over a feedback channel of known size.
--
-- This is just 'traceMat' with the 'Finite' dictionary supplied by the
-- @KnownNat n@ constraint on @F n@.
--
-- >>> let fFun (Left (F (UnsafeFin 0))) (Left (F (UnsafeFin 0))) = False; fFun (Left (F (UnsafeFin 0))) (Right (F (UnsafeFin 0))) = True; fFun (Right (F (UnsafeFin 0))) (Left (F (UnsafeFin 0))) = True; fFun _ _ = False
-- >>> let f = Mat fFun :: Mat Bool (Either (F 2) (F 1)) (Either (F 2) (F 1))
-- >>> runMat (traceF f) (finF @1 0) (finF @1 0)
-- True
traceF ::
  ( KnownNat n,
    StarSemiring s,
    Additive s,
    Multiplicative s,
    Finite b,
    Finite c
  ) =>
  Mat s (Either (F n) b) (Either (F n) c) ->
  Mat s b c
traceF :: forall (n :: Nat) s b c.
(KnownNat n, StarSemiring s, Additive s, Multiplicative s,
 Finite b, Finite c) =>
Mat s (Either (F n) b) (Either (F n) c) -> Mat s b c
traceF = Mat s (Either (F n) b) (Either (F n) c) -> Mat s b c
forall s a b c.
(StarSemiring s, Additive s, Multiplicative s, Finite a, Finite b,
 Finite c) =>
Mat s (Either a b) (Either a c) -> Mat s b c
traceMat