{- |
Module      : ConClusion.Array.Util
Description : Additional tools to work with numerical arrays
Copyright   : Phillip Seeber, 2023
License     : AGPL-3
Maintainer  : phillip.seeber@googlemail.com
Stability   : experimental
Portability : POSIX, Windows
-}
module ConClusion.Array.Util (
  IndexException (..),
  magnitude,
  normalise,
  angle,
  minDistAt,
  minDistAtVec,
  iMinimumM,
)
where

import Data.Massiv.Array as Massiv hiding (IndexException)
import RIO
import System.IO.Unsafe (unsafePerformIO)

-- | Exception regarding indexing in some kind of aaray.
newtype IndexException = IndexException String deriving (Show)

instance Exception IndexException

-- | Magnitude of a vector (length).
magnitude :: (Massiv.Numeric r e, Source r e, Floating e) => Massiv.Vector r e -> e
magnitude v = sqrt $ v !.! v

-- | Normalise a vector.
normalise :: (Massiv.Numeric r e, Source r e, Floating e) => Massiv.Vector r e -> Massiv.Vector r e
normalise v = v .* (1 / magnitude v)

-- | Angle between two vectors.
angle :: (Massiv.Numeric r e, Source r e, Floating e) => Massiv.Vector r e -> Massiv.Vector r e -> e
angle a b = acos $ a !.! b / (magnitude a * magnitude b)

-- | Find the minimal distance in a distance matrix, which is not the main diagonal.
{-# SCC minDistAt #-}
minDistAt ::
  ( Manifest r e
  , MonadThrow m
  , Ord e
  ) =>
  Massiv.Matrix r e ->
  m (e, Ix2)
minDistAt arr
  | isEmpty arr = throwM $ SizeEmptyException (Massiv.size arr)
  | otherwise = return . unsafePerformIO $ ifoldlP minFold start chFold start arr
 where
  ix0 = pureIndex 0
  e0 = arr Massiv.! ix0
  start = (e0, ix0)
  minFold acc@(eA, _) ix@(m :. n) e = if e < eA && m > n then (e, ix) else acc
  chFold acc@(eA, _) ch@(e, _) = if e < eA then ch else acc

-- | Find the minimal element of a vector, which is at a larger than the supplied index.
minDistAtVec ::
  ( Manifest r e
  , MonadThrow m
  , Ord e
  ) =>
  Ix1 ->
  Massiv.Vector r e ->
  m (e, Ix1)
minDistAtVec ixStart vec
  | isEmpty vec = throwM $ SizeEmptyException (Massiv.size vec)
  | ixStart >= nElems = throwM $ IndexOutOfBoundsException (Sz nElems) ixStart
  | otherwise = do
      let (minE, minIx) = unsafePerformIO $ ifoldlP minFold startAcc chFold startAcc searchVec
      return (minE, minIx + ixStart + 1)
 where
  Sz nElems = Massiv.size vec
  searchVec = Massiv.drop (Sz $ ixStart + 1) vec
  ix0 = 0
  e0 = searchVec Massiv.! ix0
  startAcc = (e0, ix0)
  minFold acc@(eA, _) ix e = if e < eA then (e, ix) else acc
  chFold acc ch = min acc ch

-- | Like 'Massiv.minimumM' but also returns the index of the minimal element.
iMinimumM ::
  ( Manifest r a
  , MonadThrow m
  , Index ix
  , Ord a
  ) =>
  Array r ix a ->
  m (a, ix)
iMinimumM arr
  | isEmpty arr = throwM $ SizeEmptyException (Massiv.size arr)
  | otherwise = return . unsafePerformIO $ ifoldlP minFold start chFold start arr
 where
  ix0 = pureIndex 0
  e0 = arr Massiv.! ix0
  start = (e0, ix0)

  minFold acc@(eA, _) ix e = if e < eA then (e, ix) else acc
  chFold acc@(eA, _) ch@(e, _) = if e < eA then ch else acc