-- SPDX-License-Identifier: MPL-2.0

{- | List-based vector operations. Inputs must be finite lists. Checked binary
operations return a 'MathError' for unequal lengths; legacy 'dot' and 'add'
retain their original exception behavior for that case.
-}
module HMath.Linear (
    dot,
    dotChecked,
    magnitude,
    normalize,
    add,
    addChecked,
    subtractChecked,
    scale,
) where

import HMath.Error (MathError (..))

{- | Inner product. Empty vectors have inner product zero.
Unequal lengths raise an exception; use 'dotChecked' to handle that input.
-}
dot :: (Num a) => [a] -> [a] -> a
dot :: forall a. Num a => [a] -> [a] -> a
dot [a]
xs [a]
ys = (MathError -> a) -> (a -> a) -> Either MathError a -> a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (a -> MathError -> a
forall a b. a -> b -> a
const ([Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"dot: dimension mismatch")) a -> a
forall a. a -> a
id ([a] -> [a] -> Either MathError a
forall a. Num a => [a] -> [a] -> Either MathError a
dotChecked [a]
xs [a]
ys)

-- | Inner product, or 'VectorLengthMismatch' with the left and right lengths.
dotChecked :: (Num a) => [a] -> [a] -> Either MathError a
dotChecked :: forall a. Num a => [a] -> [a] -> Either MathError a
dotChecked [a]
xs [a]
ys = [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([a] -> a) -> Either MathError [a] -> Either MathError a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> a -> a) -> [a] -> [a] -> Either MathError [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> Either MathError [c]
zipWithChecked a -> a -> a
forall a. Num a => a -> a -> a
(*) [a]
xs [a]
ys

-- | Euclidean magnitude. The empty vector has magnitude zero.
magnitude :: (Floating a) => [a] -> a
magnitude :: forall a. Floating a => [a] -> a
magnitude [a]
xs = a -> a
forall a. Floating a => a -> a
sqrt ([a] -> [a] -> a
forall a. Num a => [a] -> [a] -> a
dot [a]
xs [a]
xs)

{- | Normalize by Euclidean magnitude. A vector whose computed magnitude is
zero is returned unchanged, including the empty vector. Floating-point
overflow and non-finite values retain the underlying type's behavior.
-}
normalize :: (Eq a, Floating a) => [a] -> [a]
normalize :: forall a. (Eq a, Floating a) => [a] -> [a]
normalize [a]
xs
    | a
m a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 = [a]
xs
    | Bool
otherwise = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
m) [a]
xs
  where
    m :: a
m = [a] -> a
forall a. Floating a => [a] -> a
magnitude [a]
xs

{- | Elementwise addition. Unequal lengths raise an exception; use 'addChecked'
to handle that input.
-}
add :: (Num a) => [a] -> [a] -> [a]
add :: forall a. Num a => [a] -> [a] -> [a]
add [a]
xs [a]
ys = (MathError -> [a]) -> ([a] -> [a]) -> Either MathError [a] -> [a]
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either ([a] -> MathError -> [a]
forall a b. a -> b -> a
const ([Char] -> [a]
forall a. HasCallStack => [Char] -> a
error [Char]
"add: dimension mismatch")) [a] -> [a]
forall a. a -> a
id ([a] -> [a] -> Either MathError [a]
forall a. Num a => [a] -> [a] -> Either MathError [a]
addChecked [a]
xs [a]
ys)

-- | Elementwise addition, requiring equal lengths. Two empty lists succeed.
addChecked :: (Num a) => [a] -> [a] -> Either MathError [a]
addChecked :: forall a. Num a => [a] -> [a] -> Either MathError [a]
addChecked = (a -> a -> a) -> [a] -> [a] -> Either MathError [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> Either MathError [c]
zipWithChecked a -> a -> a
forall a. Num a => a -> a -> a
(+)

-- | Elementwise left-minus-right subtraction, requiring equal lengths.
subtractChecked :: (Num a) => [a] -> [a] -> Either MathError [a]
subtractChecked :: forall a. Num a => [a] -> [a] -> Either MathError [a]
subtractChecked = (a -> a -> a) -> [a] -> [a] -> Either MathError [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> Either MathError [c]
zipWithChecked (-)

-- | Multiply every component by a scalar, preserving length.
scale :: (Num a) => a -> [a] -> [a]
scale :: forall a. Num a => a -> [a] -> [a]
scale a
s = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a -> a -> a
forall a. Num a => a -> a -> a
* a
s)

zipWithChecked :: (a -> b -> c) -> [a] -> [b] -> Either MathError [c]
zipWithChecked :: forall a b c. (a -> b -> c) -> [a] -> [b] -> Either MathError [c]
zipWithChecked a -> b -> c
operation [a]
xs [b]
ys
    | Int
leftLength Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
rightLength = MathError -> Either MathError [c]
forall a b. a -> Either a b
Left (Int -> Int -> MathError
VectorLengthMismatch Int
leftLength Int
rightLength)
    | Bool
otherwise = [c] -> Either MathError [c]
forall a b. b -> Either a b
Right ((a -> b -> c) -> [a] -> [b] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> b -> c
operation [a]
xs [b]
ys)
  where
    leftLength :: Int
leftLength = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
    rightLength :: Int
rightLength = [b] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [b]
ys