-- |
-- Module      : Math.Manifold.Homogeneous
-- Copyright   : (c) Justus Sagemüller 2022
-- License     : GPL v3
-- 
-- Maintainer  : (@) jsag $ hvl.no
-- Stability   : experimental
-- Portability : portable
-- 

{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE UndecidableInstances       #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE DefaultSignatures          #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ConstraintKinds            #-}
{-# LANGUAGE UnicodeSyntax              #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE AllowAmbiguousTypes        #-}
{-# LANGUAGE TypeApplications           #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE EmptyCase                  #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE TypeInType                 #-}
{-# LANGUAGE CPP                        #-}
{-# LANGUAGE TemplateHaskell            #-}


module Math.Manifold.Homogeneous
    ( LieGroup(..), LieAlgebra, ActsOn(..)
    , SO
    ) where

import Data.VectorSpace
import Data.AffineSpace
import Data.Basis

import Math.Manifold.Core.PseudoAffine
import Data.Manifold.PseudoAffine
import Math.Manifold.Core.Types
import Data.Manifold.Types.Primitive
import Math.Manifold.VectorSpace.ZeroDimensional
import Math.LinearMap.Category
import Math.VectorSpace.Dual
import Data.Complex as ℂ
import Linear (V0, V1, V2, V3(..), V4, Quaternion(..), cross)
import qualified Linear.Affine as LinAff
import Data.Monoid.Additive

import Prelude hiding (($), (^))
import Control.Arrow.Constrained ((<<<), ($))
import Control.Applicative

import Data.Semigroup hiding (Dual)

import qualified Test.QuickCheck as QC

import Data.Kind (Type)
import GHC.TypeLits (Nat)
import Data.Coerce
import Data.Type.Coercion


newtype LieAlgebra g
    = LieAlgebra { forall g. LieAlgebra g -> Needle g
getLieNeedle :: Needle g }

copyNewtypeInstances [t| ∀ g . (Semimanifold g) => LieAlgebra g |]
    [''AdditiveGroup]



-- | Manifolds with a continuous group structure, whose 'Needle' space
--   is isomorphic to the associated Lie algebra.
--
--   Laws:
--
--   @
--   expMap zeroV ≡ mempty
--   lieBracket w v ≡ negateV (lieBracket v w)
--   ...
--   @
class (Semimanifold g, Monoid g) => LieGroup g where
  expMap :: LieAlgebra g -> g
  lieBracket :: Bilinear (LieAlgebra g) (LieAlgebra g) (LieAlgebra g)


data family SO_ (n :: Nat) (r :: Type)

type SO n = SO_ n Double

data instance SO_ 1 r = SO1Identity
 deriving (SO_ 1 r -> SO_ 1 r -> Bool
forall r. SO_ 1 r -> SO_ 1 r -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: SO_ 1 r -> SO_ 1 r -> Bool
$c/= :: forall r. SO_ 1 r -> SO_ 1 r -> Bool
== :: SO_ 1 r -> SO_ 1 r -> Bool
$c== :: forall r. SO_ 1 r -> SO_ 1 r -> Bool
Eq, Int -> SO_ 1 r -> ShowS
forall r. Int -> SO_ 1 r -> ShowS
forall r. [SO_ 1 r] -> ShowS
forall r. SO_ 1 r -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [SO_ 1 r] -> ShowS
$cshowList :: forall r. [SO_ 1 r] -> ShowS
show :: SO_ 1 r -> String
$cshow :: forall r. SO_ 1 r -> String
showsPrec :: Int -> SO_ 1 r -> ShowS
$cshowsPrec :: forall r. Int -> SO_ 1 r -> ShowS
Show)

instance (QC.Arbitrary r, Floating r) => QC.Arbitrary (SO_ 1 r) where
  arbitrary :: Gen (SO_ 1 r)
arbitrary = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall r. SO_ 1 r
SO1Identity

instance Semigroup (SO_ 1 r) where
  SO_ 1 r
R:SO_1r r
SO1Identity <> :: SO_ 1 r -> SO_ 1 r -> SO_ 1 r
<> SO_ 1 r
R:SO_1r r
SO1Identity = forall r. SO_ 1 r
SO1Identity
instance Monoid (SO_ 1 r) where
  mempty :: SO_ 1 r
mempty = forall r. SO_ 1 r
SO1Identity
  mappend :: SO_ 1 r -> SO_ 1 r -> SO_ 1 r
mappend = forall a. Semigroup a => a -> a -> a
(<>)

instance RealFloat' r => Semimanifold (SO_ 1 r) where
  type Needle (SO_ 1 r) = ZeroDim r
  SO_ 1 r
R:SO_1r r
SO1Identity .+~^ :: SO_ 1 r -> Needle (SO_ 1 r) -> SO_ 1 r
.+~^ ZeroDim r
Needle (SO_ 1 r)
Origin = forall r. SO_ 1 r
SO1Identity
  semimanifoldWitness :: SemimanifoldWitness (SO_ 1 r)
semimanifoldWitness = case forall v. TensorSpace v => LinearManifoldWitness v
linearManifoldWitness @r of
    LinearManifoldWitness r
LinearManifoldWitness -> forall x.
(Semimanifold (Needle x), Needle (Needle x) ~ Needle x) =>
SemimanifoldWitness x
SemimanifoldWitness

newtype instance SO_ 2 r = SO2 { forall r. SO_ 2 r -> Complex r
unitReprSO2 :: Complex r }
 deriving (SO_ 2 r -> SO_ 2 r -> Bool
forall r. Eq r => SO_ 2 r -> SO_ 2 r -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: SO_ 2 r -> SO_ 2 r -> Bool
$c/= :: forall r. Eq r => SO_ 2 r -> SO_ 2 r -> Bool
== :: SO_ 2 r -> SO_ 2 r -> Bool
$c== :: forall r. Eq r => SO_ 2 r -> SO_ 2 r -> Bool
Eq, Int -> SO_ 2 r -> ShowS
forall r. Show r => Int -> SO_ 2 r -> ShowS
forall r. Show r => [SO_ 2 r] -> ShowS
forall r. Show r => SO_ 2 r -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [SO_ 2 r] -> ShowS
$cshowList :: forall r. Show r => [SO_ 2 r] -> ShowS
show :: SO_ 2 r -> String
$cshow :: forall r. Show r => SO_ 2 r -> String
showsPrec :: Int -> SO_ 2 r -> ShowS
$cshowsPrec :: forall r. Show r => Int -> SO_ 2 r -> ShowS
Show)

instance (QC.Arbitrary r, Floating r) => QC.Arbitrary (SO_ 2 r) where
  arbitrary :: Gen (SO_ 2 r)
arbitrary = forall r. Complex r -> SO_ 2 r
SO2 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> Complex a
ℂ.cis forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
QC.arbitrary

instance RealFloat r => Semigroup (SO_ 2 r) where
  SO2 Complex r
a <> :: SO_ 2 r -> SO_ 2 r -> SO_ 2 r
<> SO2 Complex r
b = forall r. Complex r -> SO_ 2 r
SO2 forall (f :: * -> * -> *) a b.
(Function f, Object f a, Object f b) =>
f a b -> a -> b
$ Complex r
aforall a. Num a => a -> a -> a
*Complex r
b  -- perhaps should normalize?
instance RealFloat r => Monoid (SO_ 2 r) where
  mempty :: SO_ 2 r
mempty = forall r. Complex r -> SO_ 2 r
SO2 Complex r
1
  mappend :: SO_ 2 r -> SO_ 2 r -> SO_ 2 r
mappend = forall a. Semigroup a => a -> a -> a
(<>)

instance RealFloat' r => Semimanifold (SO_ 2 r) where
  type Needle (SO_ 2 r) = r
  SO_ 2 r
p .+~^ :: SO_ 2 r -> Needle (SO_ 2 r) -> SO_ 2 r
.+~^ Needle (SO_ 2 r)
d = SO_ 2 r
p forall a. Semigroup a => a -> a -> a
<> forall g. LieGroup g => LieAlgebra g -> g
expMap (forall g. Needle g -> LieAlgebra g
LieAlgebra Needle (SO_ 2 r)
d)
  semimanifoldWitness :: SemimanifoldWitness (SO_ 2 r)
semimanifoldWitness = case forall v. TensorSpace v => LinearManifoldWitness v
linearManifoldWitness @r of
    LinearManifoldWitness r
LinearManifoldWitness -> forall x.
(Semimanifold (Needle x), Needle (Needle x) ~ Needle x) =>
SemimanifoldWitness x
SemimanifoldWitness

instance RealFloat' r => LieGroup (SO_ 2 r) where
  expMap :: LieAlgebra (SO_ 2 r) -> SO_ 2 r
expMap = forall r. Complex r -> SO_ 2 r
SO2 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> Complex a
cis forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall g. LieAlgebra g -> Needle g
getLieNeedle
  lieBracket :: Bilinear
  (LieAlgebra (SO_ 2 r))
  (LieAlgebra (SO_ 2 r))
  (LieAlgebra (SO_ 2 r))
lieBracket = forall v. AdditiveGroup v => v
zeroV


newtype instance SO_ 3 r = SO3 { forall r. SO_ 3 r -> Quaternion r
unitReprSO3 :: Quaternion r }
 deriving (SO_ 3 r -> SO_ 3 r -> Bool
forall r. Eq r => SO_ 3 r -> SO_ 3 r -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: SO_ 3 r -> SO_ 3 r -> Bool
$c/= :: forall r. Eq r => SO_ 3 r -> SO_ 3 r -> Bool
== :: SO_ 3 r -> SO_ 3 r -> Bool
$c== :: forall r. Eq r => SO_ 3 r -> SO_ 3 r -> Bool
Eq, Int -> SO_ 3 r -> ShowS
forall r. Show r => Int -> SO_ 3 r -> ShowS
forall r. Show r => [SO_ 3 r] -> ShowS
forall r. Show r => SO_ 3 r -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [SO_ 3 r] -> ShowS
$cshowList :: forall r. Show r => [SO_ 3 r] -> ShowS
show :: SO_ 3 r -> String
$cshow :: forall r. Show r => SO_ 3 r -> String
showsPrec :: Int -> SO_ 3 r -> ShowS
$cshowsPrec :: forall r. Show r => Int -> SO_ 3 r -> ShowS
Show)

instance (QC.Arbitrary r, RealFloat r) => QC.Arbitrary (SO_ 3 r) where
  arbitrary :: Gen (SO_ 3 r)
arbitrary = do
    (r
a,r
b,r
c,r
d) <- forall a. Arbitrary a => Gen a
QC.arbitrary
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r. Quaternion r -> SO_ 3 r
SO3 forall (f :: * -> * -> *) a b.
(Function f, Object f a, Object f b) =>
f a b -> a -> b
$ case forall a. Floating a => a -> a
sqrt forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall (f :: * -> * -> *) a b.
(Function f, Object f a, Object f b) =>
f a b -> a -> b
$ (forall a. Num a => a -> Int -> a
^Int
2)forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>[r
a,r
b,r
c,r
d] of
      r
l | r
lforall a. Ord a => a -> a -> Bool
>r
0       -> forall a. a -> V3 a -> Quaternion a
Quaternion (r
aforall a. Fractional a => a -> a -> a
/r
l) (forall a. a -> a -> a -> V3 a
V3 r
b r
c r
d forall v s.
(VectorSpace v, s ~ Scalar v, Fractional s) =>
v -> s -> v
^/ r
l)
        | Bool
otherwise -> Quaternion r
1

instance RealFloat r => Semigroup (SO_ 3 r) where
  SO3 Quaternion r
a <> :: SO_ 3 r -> SO_ 3 r -> SO_ 3 r
<> SO3 Quaternion r
b = forall r. Quaternion r -> SO_ 3 r
SO3 forall (f :: * -> * -> *) a b.
(Function f, Object f a, Object f b) =>
f a b -> a -> b
$ Quaternion r
aforall a. Num a => a -> a -> a
*Quaternion r
b  -- perhaps should normalize?
instance RealFloat r => Monoid (SO_ 3 r) where
  mempty :: SO_ 3 r
mempty = forall r. Quaternion r -> SO_ 3 r
SO3 Quaternion r
1
  mappend :: SO_ 3 r -> SO_ 3 r -> SO_ 3 r
mappend = forall a. Semigroup a => a -> a -> a
(<>)

instance RealFloat' r => Semimanifold (SO_ 3 r) where
  type Needle (SO_ 3 r) = V3 r
  SO_ 3 r
p .+~^ :: SO_ 3 r -> Needle (SO_ 3 r) -> SO_ 3 r
.+~^ Needle (SO_ 3 r)
d = SO_ 3 r
p forall a. Semigroup a => a -> a -> a
<> forall g. LieGroup g => LieAlgebra g -> g
expMap (forall g. Needle g -> LieAlgebra g
LieAlgebra Needle (SO_ 3 r)
d)
  semimanifoldWitness :: SemimanifoldWitness (SO_ 3 r)
semimanifoldWitness = case forall v. TensorSpace v => LinearManifoldWitness v
linearManifoldWitness @r of
    LinearManifoldWitness r
LinearManifoldWitness -> forall x.
(Semimanifold (Needle x), Needle (Needle x) ~ Needle x) =>
SemimanifoldWitness x
SemimanifoldWitness

instance ∀ r . RealFloat' r => LieGroup (SO_ 3 r) where
  expMap :: LieAlgebra (SO_ 3 r) -> SO_ 3 r
expMap (LieAlgebra Needle (SO_ 3 r)
a) = forall r. Quaternion r -> SO_ 3 r
SO3 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> a
exp forall (f :: * -> * -> *) a b.
(Function f, Object f a, Object f b) =>
f a b -> a -> b
$ forall a. a -> V3 a -> Quaternion a
Quaternion r
0 Needle (SO_ 3 r)
a
  lieBracket :: Bilinear
  (LieAlgebra (SO_ 3 r))
  (LieAlgebra (SO_ 3 r))
  (LieAlgebra (SO_ 3 r))
lieBracket = coerce :: forall a b. Coercible a b => a -> b
coerce (forall a. Num a => V3 a -> V3 a -> V3 a
cross :: V3 r -> V3 r -> V3 r)

embedPureImagUnitQuaternion :: RealFloat r => S²_ r -> Quaternion r
embedPureImagUnitQuaternion :: forall r. RealFloat r => S²_ r -> Quaternion r
embedPureImagUnitQuaternion = forall a. a -> V3 a -> Quaternion a
Quaternion r
0 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall m v. NaturallyEmbedded m v => m -> v
embed

projectPureImagUnitQuaternion :: RealFloat r => Quaternion r -> S²_ r
projectPureImagUnitQuaternion :: forall r. RealFloat r => Quaternion r -> S²_ r
projectPureImagUnitQuaternion (Quaternion r
_ V3 r
p) = forall m v. NaturallyEmbedded m v => v -> m
coEmbed V3 r
p

-- | Manifolds that are homogeneous with respect to action by a Lie group.
--   Laws:
--
--   @
--   action mempty ≡ id                  (Identity)
--   action (a<>b) ≡ action a . action b (Compatibility)
--   @
class (Semimanifold m, LieGroup g) => g `ActsOn` m where
  action :: g -> m -> m

instance RealFloat' r => SO_ 2 r`ActsOn`S¹_ r where
  action :: SO_ 2 r -> S¹_ r -> S¹_ r
action (SO2 Complex r
β) S¹_ r
p = S¹_ r
p forall x. Semimanifold x => x -> Needle x -> x
.+~^ forall a. RealFloat a => Complex a -> a
ℂ.phase Complex r
β

instance RealFloat' r => SO_ 3 r`ActsOn`S²_ r where
  action :: SO_ 3 r -> S²_ r -> S²_ r
action (SO3 Quaternion r
γ) S²_ r
p = forall r. RealFloat r => Quaternion r -> S²_ r
projectPureImagUnitQuaternion forall (f :: * -> * -> *) a b.
(Function f, Object f a, Object f b) =>
f a b -> a -> b
$ Quaternion r
γ forall a. Num a => a -> a -> a
* Quaternion r
α forall a. Num a => a -> a -> a
* forall a. Fractional a => a -> a
recip Quaternion r
γ
   where α :: Quaternion r
α = forall r. RealFloat r => S²_ r -> Quaternion r
embedPureImagUnitQuaternion S²_ r
p