content: module Telescope.Data.Array where
  
  import Control.Exception (throw)
  import Control.Monad.Catch
  import Data.Binary.Get (ByteOffset, runGetOrFail)
  import Data.Binary.Put (Put, runPut)
  import Data.ByteString qualified as BS
  import Data.ByteString.Lazy qualified as BL
  import Data.Massiv.Array (Array, Comp (..), D, Index, Ix1, Ix2 (..), Ix3, Ix4, Ix5, IxN (..), Lower, Prim, Source, Stream, Sz (..), Vector)
  import Data.Massiv.Array qualified as M
  import Data.Massiv.Array.IO (Linearity (..), Pixel (..), SRGB)
  import Data.Massiv.Array.IO qualified as M
  import Data.Word (Word8)
  import Telescope.Data.Axes
  import Telescope.Data.Binary
  
  
  {- | Decode Big Endian encoded binary data into an Array of arbitrary dimensions 'ix' by specifying the 'Axes' 'Row'
  
  >>> decodeArray @Ix2 @Float [3, 2] input
  Array P Seq (Sz (2 :. 3))
    [ [ 1.0, 2.0, 3.0 ]
    , [ 4.0, 5.0, 6.0 ]
    ]
  -}
  decodeArray
    :: forall ix a m
     . (AxesIndex ix, Prim a, BinaryValue a, MonadThrow m, MonadCatch m)
    => Axes Row
    -> BS.ByteString
    -> m (Array D ix a)
  decodeArray = decodeArrayOrder BigEndian
  
  
  {- | Encode an Array as a Lazy ByteString based on the type of the element 'a'
  
  >>> myArray = decodeArray @Ix2 @Float BPFloat [3, 2] input
  >>> output = encodeArray myArray
  -}
  encodeArray
    :: (Source r a, Stream r Ix1 a, PutArray ix, BinaryValue a, Prim a)
    => Array r ix a
    -> BS.ByteString
  encodeArray = BL.toStrict . runPut . putArray BigEndian
  
  
  {- | Decode binary data into an Array of arbitrary dimensions 'ix' by specifying the 'Axes' 'Row' and the 'ByteOrder'
  
  >>> decodeArrayOrder @Ix2 @Float BigEndian [3, 2] input
  Array P Seq (Sz (2 :. 3))
    [ [ 1.0, 2.0, 3.0 ]
    , [ 4.0, 5.0, 6.0 ]
    ]
  -}
  decodeArrayOrder
    :: forall ix a m
     . (AxesIndex ix, BinaryValue a, MonadThrow m, MonadCatch m)
    => ByteOrder
    -> Axes Row
    -> BS.ByteString
    -> m (Array D ix a)
  decodeArrayOrder bo as inp = do
    fromVector as $ decodeVector @a Par bo inp
  
  
  -- | Decode binary data into as a 1d Array
  decodeVector
    :: forall a
     . (BinaryValue a)
    => Comp
    -> ByteOrder
    -> BS.ByteString
    -> Vector D a
  decodeVector c bo inp =
    let v = parseWordVector inp
     in M.makeArray c (arraySize v) (valueAt v)
   where
    numBytes = byteSize @a
  
    arraySize v =
      let Sz s = M.size v
       in Sz $ s `div` numBytes
  
    valueAt :: Vector D Word8 -> Ix1 -> a
    valueAt v ix =
      fromWords $ wordsAt v ix
  
    wordsAt :: Vector D Word8 -> Ix1 -> [Word8]
    wordsAt v ix =
      M.toList $ M.slice (ix * numBytes) (Sz numBytes) v
  
    parseWordVector :: BS.ByteString -> Vector D Word8
    parseWordVector = M.fromByteString c
  
    -- this should never actually fail
    fromWords :: [Word8] -> a
    fromWords ws =
      case runGetOrFail (get bo) (BL.pack ws) of
        Left (ip, byts, e) -> throw $ BinaryParseError byts (e <> " " <> show ip <> show ws)
        Right (_, _, a) -> a
  
  
  -- | Resize a Vector into an Array
  fromVector
    :: forall ix a m
     . (AxesIndex ix, MonadThrow m, MonadCatch m)
    => Axes Row
    -> Vector D a
    -> m (Array D ix a)
  fromVector as v = do
    ix <- axesIndex as
    ea <- try $ M.resizeM (Sz ix) v
    case ea of
      Left (e :: M.SizeException) -> throwM $ ResizeMismatch (show e)
      Right a -> pure a
  
  
  data ArrayError
    = BinaryParseError !ByteOffset String
    | AxesMismatch !(Axes Row)
    | ResizeMismatch String
    deriving (Show, Exception)
  
  
  -- | Convert between an 'Index' and 'Axes' 'Row'
  class (Index ix) => AxesIndex ix where
    axesIndex :: (MonadThrow m) => Axes Row -> m ix
    indexAxes :: ix -> Axes Row
  
  
  instance AxesIndex Ix1 where
    axesIndex (Axes [i]) = pure i
    axesIndex as = throwM $ AxesMismatch as
    indexAxes n = Axes [n]
  
  
  instance AxesIndex Ix2 where
    axesIndex (Axes [c, r]) = do
      ix1 <- axesIndex $ Axes [r]
      pure $ c :. ix1
    axesIndex as = throwM $ AxesMismatch as
    indexAxes (c :. r) = Axes [c, r]
  
  
  instance AxesIndex Ix3 where
    axesIndex = axesIndexN
    indexAxes = indexAxesN
  
  
  instance AxesIndex Ix4 where
    axesIndex = axesIndexN
    indexAxes = indexAxesN
  
  
  instance AxesIndex Ix5 where
    axesIndex = axesIndexN
    indexAxes = indexAxesN
  
  
  -- | Encode an 'Array' of dimensions up to 'Ix5'
  class PutArray ix where
    putArray
      :: (BinaryValue a, Source r a, Stream r Ix1 a, Prim a)
      => ByteOrder
      -> Array r ix a
      -> Put
  
  
  instance PutArray Ix1 where
    putArray bo = M.sfoldl (\b a -> b <> put bo a) mempty
  instance PutArray Ix2 where
    putArray = M.foldOuterSlice . putArray
  instance PutArray Ix3 where
    putArray = M.foldOuterSlice . putArray
  instance PutArray Ix4 where
    putArray = M.foldOuterSlice . putArray
  instance PutArray Ix5 where
    putArray = M.foldOuterSlice . putArray
  
  
  axesIndexN :: (AxesIndex (Lower (IxN n))) => (MonadThrow m) => Axes Row -> m (IxN n)
  axesIndexN (Axes (a : as)) = do
    ixl <- axesIndex (Axes as)
    pure $ a :> ixl
  axesIndexN as = throwM $ AxesMismatch as
  
  
  indexAxesN :: (AxesIndex (Lower (IxN n))) => IxN n -> Axes Row
  indexAxesN (d :> ix) =
    let Axes ax = indexAxes ix
     in Axes $ d : ax
  
  
  -- | Get the Axes for a given array size
  sizeAxes :: (AxesIndex ix, Index ix) => Sz ix -> Axes Column
  sizeAxes (Sz ix) = toColumnMajor $ indexAxes ix
  
  
  -- approximates python's viridis
  heatmap :: forall n. (Ord n, RealFrac n) => Array D Ix2 n -> Array D Ix2 (Pixel (SRGB 'NonLinear) Word8)
  heatmap arr =
    let maxn = maximum arr :: n
     in fmap (color . (/ maxn)) arr
   where
    color x =
      let r = floor (255 * abs (x - 0.3))
          g = floor (255 * x)
          b = floor (255 * (1 - x))
       in M.PixelSRGB r g b
  