Linear.hs

{-# LANGUAGE AllowAmbiguousTypes  #-}
{-# LANGUAGE PolyKinds            #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE UndecidableInstances #-}

module Math.Linear
  ( -- * Points and vectors
    Segment(..), Mat22(..)

  -- * Points and vectors (second version)
  , ℝ(..), T(.., V2, V3, V4)
  , Fin(..), MFin(..)
  , universe, coordinates, choose
  , RepDim, RepresentableQ(..)
  , Representable(..), set, injection, projection
  , Vec(..), (!), find, zipIndices

  , swap, rotate

  ) where

-- base
import Prelude hiding ( unzip )
import Control.Applicative
  ( ZipList(..) )
import Data.Coerce
  ( coerce )
import Data.Kind
  ( Type )
import GHC.Exts
  ( proxy# )
import GHC.Generics
  ( Generic, Generic1, Generically(..), Generically1(..) )
import GHC.Stack
  ( HasCallStack )
import GHC.TypeNats
  ( Nat, KnownNat, type (<=)
  , natVal'
  )

-- acts
import Data.Act
  ( Act((•)), Torsor((-->)) )

-- deepseq
import Control.DeepSeq
  ( NFData(..), NFData1 )

-- groups
import Data.Group
  ( Group(..) )

-- groups-generic
import Data.Group.Generics
  ( )

-- brush-strokes
import Math.Linear.Internal
import Math.Ring
  ( Ring )
import qualified Math.Ring as Ring

--------------------------------------------------------------------------------

data Mat22 = Mat22 !Double !Double !Double !Double
data Segment p =
  Segment
    { segmentStart :: !p
    , segmentEnd   :: !p
    }
  deriving stock ( Generic, Generic1, Functor, Foldable, Traversable )
  deriving ( Semigroup, Monoid, Group )
    via Generically ( Segment p )
  deriving Applicative
    via Generically1 Segment
  deriving anyclass ( NFData, NFData1 )

instance Show p => Show ( Segment p ) where
  show ( Segment s e ) = show s ++ " -> " ++ show e

--------------------------------------------------------------------------------

-- | Tangent space to Euclidean space.
type T :: Type -> Type
newtype T e = T { unT :: e }
  deriving stock ( Eq, Functor, Foldable, Traversable )
  deriving newtype NFData

instance Semigroup ( T Double ) where
  (<>) = coerce ( (+) @Double )
instance Monoid ( T Double ) where
  mempty = T 0

instance Group ( T Double ) where
  invert ( T x ) = T ( negate x )

instance Act ( T Double ) Double where
  T u • v = u + v
instance Torsor ( T Double ) Double where
  a --> b = T ( b - a )

instance {-# OVERLAPPING #-} Show ( ℝ n ) => Show ( T ( ℝ n ) ) where
  show ( T p ) = "V" ++ drop 1 ( show p )
instance {-# INCOHERENT #-} Show v => Show ( T v ) where
  showsPrec p ( T v )
    = showParen ( p >= 11 )
    $ showString "T"
    . showsPrec 0 v

instance Applicative T where
  pure = T
  T f <*> T a = T ( f a )

{-# COMPLETE V2 #-}
pattern V2 :: Double -> Double -> T ( ℝ 2 )
pattern V2 x y = T ( ℝ2 x y )

{-# COMPLETE V3 #-}
pattern V3 :: Double -> Double -> Double -> T ( ℝ 3 )
pattern V3 x y z = T ( ℝ3 x y z )

{-# COMPLETE V4 #-}
pattern V4 :: Double -> Double -> Double -> Double -> T ( ℝ 4 )
pattern V4 x y z w = T ( ℝ4 x y z w )

--------------------------------------------------------------------------------

type Vec :: Nat -> Type -> Type
newtype Vec n a = Vec { vecList :: [ a ] }
type role Vec nominal representational

deriving newtype instance Show a => Show ( Vec n a )
deriving newtype instance   Eq a =>   Eq ( Vec n a )
deriving newtype instance  Ord a =>  Ord ( Vec n a )
deriving newtype instance  Functor ( Vec n )
deriving newtype instance Foldable ( Vec n )
deriving via ZipList
  instance Applicative ( Vec n )
instance Traversable ( Vec n ) where
  traverse f ( Vec as ) = Vec <$> traverse f as
instance NFData a => NFData ( Vec n a ) where
  rnf ( Vec as ) = rnf as

universe :: forall n. KnownNat n => Vec n ( Fin n )
universe = Vec [ Fin i | i <- [ 1 .. fromIntegral ( natVal' @n proxy# ) ] ]
{-# INLINEABLE universe #-}

coordinates :: forall r u. ( Representable r u ) => u -> Vec ( RepDim u ) r
coordinates u = fmap ( index u ) $ universe @( RepDim u )
{-# INLINEABLE coordinates #-}

-- | Binomial coefficient: choose all subsets of size @k@ of the given set
-- of size @n@.
choose
  :: forall n k
   . ( KnownNat n, KnownNat k
     , 1 <= n, 1 <= k, k <= n
     )
  => [ Vec k ( Fin n ) ]
choose = coerce $ go ( fromIntegral $ natVal' @n proxy# )
                     ( fromIntegral $ natVal' @k proxy# )
  where
    go :: Word -> Word -> [ [ Word ] ]
    go n k
      | k == 1
      = [ [ i ] | i <- [ 1 .. n ] ]
      | n == k
      = [ [ 1 .. n ] ]
    go n k = go ( n - 1 ) k
          ++ ( map ( ++ [ n ] ) $ go ( n - 1 ) ( k - 1 ) )
{-# INLINEABLE choose #-}

infixl 9 !
(!) :: forall l a. HasCallStack => Vec l a -> Fin l -> a
( Vec l ) ! Fin i = l !! ( fromIntegral i - 1 )

find :: forall l a. ( a -> Bool ) -> Vec l a -> MFin l
find f ( Vec v ) = MFin ( find_ 1 v )
  where
    find_ :: Word -> [ a ] -> Word
    find_ j ( a : as )
      | f a
      = j
      | otherwise
      = find_ ( j + 1 ) as
    find_ _ [] = 0

zipIndices :: forall n a. Vec n a -> [ ( Fin n, a ) ]
zipIndices ( Vec v ) = zipIndices_ 1 v
  where
    zipIndices_ ::  Word -> [ a ] -> [ ( Fin n, a ) ]
    zipIndices_ _ [] = []
    zipIndices_ i (a : as) = ( Fin i, a ) : zipIndices_ ( i + 1 ) as

--------------------------------------------------------------------------------

-- | Rotate a vector by the given angle (counter-clockwise),
-- given the cosine and sine of the angle (in that order).
rotate :: ( Representable r m, RepDim m ~ 2, Ring r )
       => r -- \( \cos \theta \)
       -> r -- \( \sin \theta \)
       -> T m
       -> T m
rotate cosθ sinθ ( T xy ) =
  let x = xy `index` Fin 1
      y = xy `index` Fin 2
  in  T $ tabulate \ case
        Fin 1 -> x Ring.* cosθ Ring.- y Ring.* sinθ
        _     -> y Ring.* cosθ Ring.+ x Ring.* sinθ
{-# INLINEABLE rotate #-}

-- | Swap the two coordinates of a 2D vector.
swap :: ( Representable r m, RepDim m ~ 2, Ring r )
       => T m
       -> T m
swap ( T xy ) =
  let x = xy `index` Fin 1
      y = xy `index` Fin 2
  in  T $ tabulate \ case
        Fin 1 -> y
        _     -> x
{-# INLINEABLE swap #-}