-- |
-- Module      : ConClusion.Numeric.Data
-- Description : Type castings, conversions and utilities
-- Copyright   : Phillip Seeber, 2021
-- License     : AGPL-3
-- Maintainer  : phillip.seeber@googlemail.com
-- Stability   : experimental
-- Portability : POSIX, Windows
module ConClusion.Numeric.Data
  ( -- * Conversion of array types.

    -- Conversion between HMatrix and Massiv's array types
    IndexException (..),
    vecH2M,
    vecM2H,
    matH2M,
    matM2H,

    -- * Array Processing
    magnitude,
    normalise,
    angle,
    minDistAt,
    minDistAtVec,
    iMinimumM,

    -- * Utilities
    printMat,

    -- * Binary Trees
    BinTree (..),
    root,
    takeBranchesWhile,
    takeLeafyBranchesWhile,
  )
where

import Data.Aeson hiding (Array)
import Data.Massiv.Array as Massiv hiding (IndexException)
import Data.Massiv.Array.Manifest.Vector as Massiv
import Formatting
import Numeric.LinearAlgebra as LA hiding (magnitude, (<>))
import RIO
import qualified RIO.Vector.Storable as VectorS
import System.IO.Unsafe (unsafePerformIO)

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

instance Exception IndexException

-- | Converts a vector from the HMatrix package to the Massiv representation.
{-# SCC vecH2M #-}
vecH2M :: (Element e, Mutable r Ix1 e) => VectorS.Vector e -> Massiv.Vector r e
vecH2M hVec = fromVector' Seq (Sz $ VectorS.length hVec) hVec

-- | Converts a vector from the Massiv representation to the HMatrix representation.
{-# SCC vecM2H #-}
vecM2H :: (Manifest r Ix1 e, Element e) => Massiv.Vector r e -> LA.Vector e
vecM2H = Massiv.toVector

-- | Converts a matrix from the HMatrix representation to the Massiv representation.
{-# SCC matH2M #-}
matH2M :: (Mutable r Ix1 e, Element e) => LA.Matrix e -> Massiv.Matrix r e
matH2M hMat = Massiv.resize' (Sz $ nRows :. nCols) . vecH2M . LA.flatten $ hMat
  where
    nRows = LA.rows hMat
    nCols = LA.cols hMat

-- | Converts a matrix from Massiv to HMatrix representation.
{-# SCC matM2H #-}
matM2H :: (Manifest r Ix1 e, Element e, Resize r Ix2, Load r Ix2 e) => Massiv.Matrix r e -> LA.Matrix e
matM2H mMat = LA.reshape nCols . vecM2H . Massiv.flatten $ mMat
  where
    Sz (_nRows :. nCols) = Massiv.size mMat

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

-- | Normalise a vector.
normalise :: (Massiv.Numeric r e, Source r Ix1 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 Ix1 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 Ix2 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 Ix1 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 ix a,
    MonadThrow m,
    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

-- | Quickly print a matrix with numerical values
printMat :: (Source r Ix2 e, Real e) => Massiv.Matrix r e -> Massiv.Matrix D Text
printMat mat = Massiv.map (sformat (left 4 ' ' %. fixed 2)) mat

----------------------------------------------------------------------------------------------------
-- Binary Trees.

-- | A binary tree.
data BinTree e = Leaf e | Node e (BinTree e) (BinTree e)
  deriving (Eq, Show, Generic)

instance (FromJSON e) => FromJSON (BinTree e)

instance (ToJSON e) => ToJSON (BinTree e)

instance Functor BinTree where
  fmap f (Leaf a) = Leaf (f a)
  fmap f (Node a l r) = Node (f a) (fmap f l) (fmap f r)

-- | Look at the root of a binary tree.
root :: BinTree e -> e
root (Leaf e) = e
root (Node e _ _) = e

-- | Steps down each branch of a tree until some criterion is satisfied or the end of the branch is
-- reached. Each end of the branch is added to a result.
takeBranchesWhile :: (a -> Bool) -> BinTree a -> Massiv.Vector DL a
takeBranchesWhile chk tree = go tree (Massiv.empty @DL)
  where
    go (Leaf v) acc = if chk v then acc `snoc` v else acc
    go (Node v l r) acc =
      let vAcc = if chk v then acc `snoc` v else acc
          lAcc = go l vAcc
          rAcc = go r lAcc
       in if chk v then rAcc else vAcc

-- | Takes the first value in each branch, that does not fullfill the criterion anymore and adds it
-- to the result. Terminal leafes of the branches are always taken.
takeLeafyBranchesWhile :: (a -> Bool) -> BinTree a -> Massiv.Vector DL a
takeLeafyBranchesWhile chk tree = go tree (Massiv.empty @DL)
  where
    go (Leaf v) acc = acc `snoc` v
    go (Node v l r) acc =
      let vAcc = if chk v then acc else acc `snoc` v
          lAcc = go l vAcc
          rAcc = go r lAcc
       in if chk v then rAcc else vAcc