{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeAbstractions #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
module Data.Array.Nested.Convert where

import Control.Category
import Data.Proxy
import Data.Type.Equality

import Data.Array.Mixed.Lemmas
import Data.Array.Mixed.Types
import Data.Array.Nested.Internal.Lemmas
import Data.Array.Nested.Mixed
import Data.Array.Nested.Mixed.Shape
import Data.Array.Nested.Ranked
import Data.Array.Nested.Shaped
import Data.Array.Nested.Shaped.Shape


stoRanked :: Elt a => Shaped sh a -> Ranked (Rank sh) a
stoRanked sarr@(Shaped arr)
  | Refl <- lemRankMapJust (sshape sarr)
  = mtoRanked arr

rcastToShaped :: Elt a => Ranked (Rank sh) a -> ShS sh -> Shaped sh a
rcastToShaped (Ranked arr) targetsh
  | Refl <- lemRankReplicate (shxRank (shCvtSX targetsh))
  , Refl <- lemRankMapJust targetsh
  = mcastToShaped arr targetsh

-- | The only constructor that performs runtime shape checking is 'CastXS''.
-- For the other construtors, the types ensure that the shapes are already
-- compatible. To convert between 'Ranked' and 'Shaped', go via 'Mixed'.
data Castable a b where
  CastId  :: Castable a a
  CastCmp :: Castable b c -> Castable a b -> Castable a c

  CastRX  :: Castable a b -> Castable (Ranked n a)           (Mixed (Replicate n Nothing) b)
  CastSX  :: Castable a b -> Castable (Shaped sh a)          (Mixed (MapJust sh) b)

  CastXR  :: Elt b
          => Castable a b -> Castable (Mixed sh a)           (Ranked (Rank sh) b)
  CastXS  :: Castable a b -> Castable (Mixed (MapJust sh) a) (Shaped sh b)
  CastXS' :: (Rank sh ~ Rank sh', Elt b) => ShS sh'
          -> Castable a b -> Castable (Mixed sh a)           (Shaped sh' b)

  CastRR  :: Castable a b -> Castable (Ranked n a)           (Ranked n b)
  CastSS  :: Castable a b -> Castable (Shaped sh a)          (Shaped sh b)
  CastXX  :: Castable a b -> Castable (Mixed sh a)           (Mixed sh b)

instance Category Castable where
  id = CastId
  (.) = CastCmp

castCastable :: (Elt a, Elt b) => Castable a b -> a -> b
castCastable = \c x -> munScalar (go c (mscalar x))
  where
    -- The 'esh' is the extension shape: the casting happens under a whole
    -- bunch of additional dimensions that it does not touch. These dimensions
    -- are 'esh'.
    -- The strategy is to unwind step-by-step to a large Mixed array, and to
    -- perform the required checks and castings when re-nesting back up.
    go :: Castable a b -> Mixed esh a -> Mixed esh b
    go CastId x = x
    go (CastCmp c1 c2) x = go c1 (go c2 x)
    go (CastRX c) (M_Ranked (M_Nest esh x)) = M_Nest esh (go c x)
    go (CastSX c) (M_Shaped (M_Nest esh x)) = M_Nest esh (go c x)
    go (CastXR @_ @_ @sh c) (M_Nest @esh esh x)
      | Refl <- lemRankAppRankEqRepNo (Proxy @esh) (Proxy @sh)
      = let x' = go c x
            ssx' = ssxAppend (ssxFromShape esh)
                             (ssxReplicate (shxRank (shxDropSSX @esh @sh (mshape x') (ssxFromShape esh))))
        in M_Ranked (M_Nest esh (mcast ssx' x'))
    go (CastXS c) (M_Nest esh x) = M_Shaped (M_Nest esh (go c x))
    go (CastXS' @sh @sh' sh' c) (M_Nest @esh esh x)
      | Refl <- lemRankAppRankEqMapJust (Proxy @esh) (Proxy @sh) (Proxy @sh')
      = M_Shaped (M_Nest esh (mcast (ssxFromShape (shxAppend esh (shCvtSX sh')))
                                    (go c x)))
    go (CastRR c) (M_Ranked (M_Nest esh x)) = M_Ranked (M_Nest esh (go c x))
    go (CastSS c) (M_Shaped (M_Nest esh x)) = M_Shaped (M_Nest esh (go c x))
    go (CastXX c) (M_Nest esh x) = M_Nest esh (go c x)

    lemRankAppRankEqRepNo :: Proxy esh -> Proxy sh
                          -> Rank (esh ++ sh) :~: Rank (esh ++ Replicate (Rank sh) Nothing)
    lemRankAppRankEqRepNo _ _ = unsafeCoerceRefl

    lemRankAppRankEqMapJust :: Rank sh ~ Rank sh'
                            => Proxy esh -> Proxy sh -> Proxy sh'
                            -> Rank (esh ++ sh) :~: Rank (esh ++ MapJust sh')
    lemRankAppRankEqMapJust _ _ _ = unsafeCoerceRefl