{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE DataKinds, KindSignatures, TypeOperators, FlexibleContexts #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Data.YAP.CliffordAlgebra
-- Copyright   :  (c) Ross Paterson 2021
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  R.Paterson@city.ac.uk
-- Stability   :  provisional
-- Portability :  type-level literals
--
-- An example instance of the algebraic classes: Clifford algebras (aka
-- geometric algebras) over a vector space with a given orthonormal basis.
--
-----------------------------------------------------------------------------
module Data.YAP.CliffordAlgebra (
    -- * Multivectors
    Multivector,
    -- ** Example algebras
    Complex, SplitComplex, Dual,
    Quaternion, SplitQuaternion,
    VGA, STA,
    -- * Basis vectors
    -- | The algebra is generated by scalars and a family of vectors
    -- comprising an orthonormal basis.
    --
    -- The first @p@ basis vectors square to @1@, the next @q@ to @-1@
    -- and the rest to @0@.
    -- If @u@ and @v@ are distinct basis vectors, @u*v = - v*u@.
    BasisID,
    basic,
    e0, e1, e2, e3, e4, e5,
    -- * Construction
    scalar,
    fromVector,
    pseudoscalar,
    standardBlade,
    fromBlades,
    -- * Extraction
    scalarPart,
    vectorPart,
    toBlades,
    -- * Operations
    maybeRecip,
    magnitude,
    -- ** Involutions
    involution,
    gradeInvolution,
    reversion,
    conjugate,
    ) where

import Prelude.YAP
import Data.YAP.Algebra
import Data.YAP.Vector
import Data.YAP.Utilities.List

import Control.Monad (guard)
import Data.Bits
import Data.List (intersperse)
import Data.Map (Map)
import qualified Data.Map as Map
import GHC.TypeLits

-- | Identifiers of basis vectors.
type BasisID = Int

-- A set of basis vectors of the vector space, defining a basis element
-- of the Clifford algebra.  Each bit represents a basis vector.
newtype BasisSet = BasisSet { basisElements :: Int }
    deriving (Eq, Ord, Show)

emptyBasisSet :: BasisSet
emptyBasisSet = BasisSet 0

singletonBasisSet :: BasisID -> BasisSet
singletonBasisSet i = BasisSet (bit i)

-- set of i..j
setRange :: Int -> Int -> BasisSet
setRange i j = BasisSet ((bit (j+1-i) - 1) `shiftL` i)

intersectionBasisSet :: BasisSet -> BasisSet -> BasisSet
intersectionBasisSet (BasisSet x) (BasisSet y) = BasisSet (x .&. y)

isEmptyBasisSet :: BasisSet -> Bool
isEmptyBasisSet (BasisSet s) = s == 0

-- | The number of linearly independent vectors in the blade
setSize :: BasisSet -> Int
setSize (BasisSet s) = popCount s

-- identifiers of basis vectors in the set
elems :: BasisSet -> [BasisID]
elems (BasisSet s) = [i | (i, b) <- zip [0..] (bits s), b]

-- bits of x, least significant first
bits :: (Bits a) => a -> [Bool]
bits x =
    [testBit b 0 | b <- takeWhile (/= zeroBits) (iterate (flip shiftR 1) x)]

-- | A type that includes scalars of type @a@, vectors over @a@ and values
-- that can be formed from these by multiplication and addition.
--
-- The basis consists of @p@ vectors that square to @1@, @q@ vectors that
-- square to @-1@ and @r@ vectors that square to @0@.
newtype Multivector (p::Nat) (q::Nat) (r::Nat) a = MV (Map BasisSet a)
    -- Blades with zero coefficients are omitted.
    deriving (Eq, Ord)

-- Examples

-- | Complex numbers: @i = 'e0'@ squares to @-1@
type Complex = Multivector 0 1 0

-- | Split-complex numbers (equivalent to pairs of numbers): @'e0'^2 = 1@
type SplitComplex = Multivector 1 0 0

-- | Dual numbers: @'e0'^2 = 0@
type Dual = Multivector 0 0 1

-- | Quaternions: @i = 'e0'@, @j = 'e1'@, @k = i*j@
-- (or other encodings suggested by symmetry)
type Quaternion = Multivector 0 2 0

-- | Split-quaternions (equivalent to pairs of quaternions)
type SplitQuaternion = Multivector 0 3 0

-- | Algebra generated by Euclidean vector algebra in \(n\) spatial dimensions
type VGA n = Multivector n 0 0

-- | Minkowski spacetime algebra (time and 3 spatial dimensions)
type STA = Multivector 1 3 0

makeMultivector :: (KnownNat p, KnownNat q, KnownNat r) =>
    (Int -> Int -> Int -> Multivector p q r a) -> Multivector p q r a
makeMultivector k =
    makeMultivectorSing (\ p q r -> k (intVal p) (intVal q) (intVal r))

makeMultivectorSing :: (KnownNat p, KnownNat q, KnownNat r) =>
    (SNat p -> SNat q -> SNat r -> Multivector p q r a) -> Multivector p q r a
makeMultivectorSing k = k natSing natSing natSing

-- The signature of a multivector
signature :: (KnownNat p, KnownNat q, KnownNat r) =>
    Multivector p q r a -> (Int, Int, Int)
signature = signatureSing natSing natSing natSing

-- helper function to get the singletons for p, q and r
signatureSing :: (KnownNat p, KnownNat q, KnownNat r) =>
    SNat p -> SNat q -> SNat r -> Multivector p q r a -> (Int, Int, Int)
signatureSing sing_p sing_q sing_r _ =
    (intVal sing_p, intVal sing_q, intVal sing_r)

intVal :: (KnownNat n) => SNat n -> Int
intVal sing = fromInteger (natVal sing)

-- set of indices of basis vectors that square to -1
negBasisSet :: (KnownNat p, KnownNat q, KnownNat r) =>
    Multivector p q r a -> BasisSet
negBasisSet mv = setRange p (p+q-1)
  where
    (p, q, _) = signature mv

-- set of indices of basis vectors that square to zero
nullBasisSet :: (KnownNat p, KnownNat q, KnownNat r) =>
    Multivector p q r a -> BasisSet
nullBasisSet mv = setRange (p+q) (p+q+r-1)
  where
    (p, q, r) = signature mv

-- | Embedding of a scalar value.
-- A semiring morphism satisfying
--
-- prop> scalar a * m = m * scalar a
--
scalar :: (Eq a, AdditiveMonoid a) => a -> Multivector p q r a
scalar v
  | v == zero = zero
  | otherwise = MV (Map.singleton emptyBasisSet v)

-- | Unit pseudoscalar (product of all basis vectors)
pseudoscalar :: (KnownNat p, KnownNat q, KnownNat r, Semiring a) =>
    Multivector p q r a
pseudoscalar = makeMultivector $ \ p q r ->
    MV (Map.singleton (setRange 0 (p+q+r-1)) one)

-- | A product of linearly independent vectors
type Blade a = (BasisSet, a)

-- | The number of linearly independent vectors in the blade
grade :: Blade a -> Int
grade (s, _) = setSize s

maxGrade :: Multivector p q r a -> Int
maxGrade (MV mv) = maximum (0:map grade (Map.assocs mv))

-- | A distinct basis vector for each identifier satisfying
--
-- prop> basic i * scalar a = scalar a * basic i
-- prop> i /= j => basic i * basic j = - basic j * basic i
-- prop> i < p => basic i^2 = 1
-- prop> p <= i && i < p+q => basic i^2 = -1
-- prop> p+q <= i && i < p+q+r => basic i^2 = 0
basic :: (Ring a) => BasisID -> Multivector p q r a
basic i = MV (Map.singleton (singletonBasisSet i) 1)

-- | Basis vector @'basic' 0@
e0 :: (Ring a) => Multivector p q r a
e0 = basic 0

-- | Basis vector @'basic' 1@
e1 :: (Ring a) => Multivector p q r a
e1 = basic 1

-- | Basis vector @'basic' 2@
e2 :: (Ring a) => Multivector p q r a
e2 = basic 2

-- | Basis vector @'basic' 3@
e3 :: (Ring a) => Multivector p q r a
e3 = basic 3

-- | Basis vector @'basic' 4@
e4 :: (Ring a) => Multivector p q r a
e4 = basic 4

-- | Basis vector @'basic' 5@
e5 :: (Ring a) => Multivector p q r a
e5 = basic 5

-- Given a list in ascending order of keys, with no repeated keys, but
-- possibly with zero coefficients, convert to Multivector
multivector :: (Eq a, AdditiveMonoid a) => [Blade a] -> Multivector p q r a
multivector bs = MV (Map.filter (/= zero) (Map.fromListWith (+) bs))

instance (Ord a, Show a, AdditiveMonoid a) => Show (Multivector p q r a) where
    showsPrec d mv = case toBlades mv of
        [] -> showChar '0'
        [b] -> showBlade d b
        bs -> showParen (d > 6) $
            compose $ intersperse (showString " + ") $ map (showBlade 7) bs

showBlade :: (Show a) => Int -> ([BasisID], a) -> ShowS
showBlade d ([], v) =  showParen (d > 10) $
        showString "scalar " . showsPrec 11 v
showBlade d (is, v) = showParen (d > 7) $
    compose $ intersperse (showChar '*') $
        showString "scalar " . showsPrec 11 v:[showChar 'e' . shows i | i <- is]

-- | A basis multivector, formed from a product of the indicated basis vectors.
standardBlade :: (KnownNat p, KnownNat q, KnownNat r, Eq a, Ring a) =>
    [BasisID] -> Multivector p q r a
standardBlade is = product (map basic is)

-- | Embedding of a vector
fromVector :: (KnownNat p, KnownNat q, KnownNat r, Eq a, Ring a) =>
    Vector (p+q+r) a -> Multivector p q r a
fromVector v = sum [scalar a * basic i | (i, a) <- zip [0..] (coordinates v)]

-- | Construct a multivector as the sum of standard blades.
-- @
-- fromBlades bs = sum [scalar a * standardBlade is | (is, a) <- bs]
-- @
fromBlades :: (KnownNat p, KnownNat q, KnownNat r, Eq a, Ring a) =>
    [([BasisID], a)] -> Multivector p q r a
fromBlades ias = sum [scalar a * standardBlade is | (is, a) <- ias]

instance (Eq a, AdditiveMonoid a) => AdditiveMonoid (Multivector p q r a) where
    MV xs + MV ys = MV (Map.filter (/= zero) (Map.unionWith (+) xs ys))
    zero = MV Map.empty
    atimes n (MV xs)
      | n == zero = zero
      | otherwise = MV (fmap (atimes n) xs)

instance (Eq a, AbelianGroup a) => AbelianGroup (Multivector p q r a) where
    negate (MV xs) = MV (fmap negate xs)
    gtimes n (MV xs)
      | n == zero = zero
      | otherwise = MV (fmap (gtimes n) xs)

instance (KnownNat p, KnownNat q, KnownNat r, Eq a, Ring a) =>
        Semiring (Multivector p q r a) where
    (*) = multMV
    one = scalar one
    fromNatural i = scalar (fromNatural i)

multMV :: (KnownNat p, KnownNat q, KnownNat r, Eq a, Ring a) =>
    Multivector p q r a -> Multivector p q r a -> Multivector p q r a
multMV mv@(MV xs) (MV ys) =
    sum [multivector [multBlade negs nulls x y] |
        x <- Map.assocs xs, y <- Map.assocs ys]
  where
    negs = negBasisSet mv
    nulls = nullBasisSet mv

multBlade :: (Ring a) => BasisSet -> BasisSet -> Blade a -> Blade a -> Blade a
multBlade negs nulls (BasisSet x, xv) (BasisSet y, yv) =
    (BasisSet (xor x y),
        quad negs nulls (BasisSet (x .&. y)) * sign (oddMerge x y) * xv * yv)

-- Quadratic form of the vector space on a set of basis elements.
quad :: (Ring a) => BasisSet -> BasisSet -> BasisSet -> a
quad negs nulls x
  | isEmptyBasisSet (intersectionBasisSet nulls x) =
    sign (odd (setSize (intersectionBasisSet negs x)))
  | otherwise = 0

-- sign corresponding to a given parity
sign :: (Ring a) => Bool -> a
sign odd_parity
  | odd_parity = -1
  | otherwise = 1

-- Returns True if merging the set positions of x (least significant
-- first) with those of y would require an odd number of swaps.
oddMerge :: (FiniteBits a) => a -> a -> Bool
oddMerge x y = oddParity (upwardParity x .&. y)

-- the value contains an odd number of 1s
oddParity :: (FiniteBits a) => a -> Bool
oddParity = odd . popCount

-- set each bit to the parity of the bits above that position in x
upwardParity :: (Bits a) => a -> a
upwardParity x =
    foldr (.|.) zeroBits
        [bit i | (i, b) <- zip [0..] (tail (scanr xor False (bits x))), b]

instance (KnownNat p, KnownNat q, KnownNat r, Eq a, Ring a) =>
        Ring (Multivector p q r a) where
    fromInteger i = scalar (fromInteger i)

instance (KnownNat p, KnownNat q, KnownNat r, Eq a, FromRational a) =>
        FromRational (Multivector p q r a) where
    fromRational x = scalar (fromRational x)

-- | The scalar part of a multivector
scalarPart :: (AdditiveMonoid a) => Multivector p q r a -> a
scalarPart (MV bs) = case Map.assocs bs of
    (BasisSet 0, c):_ -> c
    _ -> zero

-- | The vector part of a multivector
vectorPart ::
    -- need to add KnownNat (p + q + r), because GHC cannot deduce it
    (KnownNat p, KnownNat q, KnownNat r, KnownNat (p + q + r),
     AdditiveMonoid a) =>
    Multivector p q r a -> Vector (p+q+r) a
vectorPart (MV vs) = vector (pad 0 vp)
  where
    vp = [(bv, c) | b@(i, c) <- Map.assocs vs, grade b == 1, bv <- elems i]

pad :: (AdditiveMonoid a) => Int -> [(Int, a)] -> [a]
pad _ [] = []
pad p ((k, v):vs) = replicate (k-p) zero ++ v:pad (k+1) vs

-- | Decompose the multivector as a set of standard blades
-- (the inverse of 'fromBlades').
toBlades :: (AdditiveMonoid a) => Multivector p q r a -> [([BasisID], a)]
toBlades (MV vs) = [(elems i, c) | (i, c) <- Map.assocs vs]

-- | General grade-based involution that negates grades satisfying
-- the predicate.
involution :: (Ring a) =>
    (Int -> Bool) -> Multivector p q r a -> Multivector p q r a
involution p (MV vs) = MV (Map.mapWithKey sign_flip vs)
  where
    sign_flip i v
      | p (setSize i) = - v
      | otherwise = v

-- | Automorphism that negates odd grades.
gradeInvolution :: (Ring a) => Multivector p q r a -> Multivector p q r a
gradeInvolution = involution odd

-- | Anti-automorphism equivalent to reversing the order of multiplication
-- of basis vectors.
reversion :: (Ring a) => Multivector p q r a -> Multivector p q r a
reversion = involution (\ n -> odd (n `div` 2))

-- | Clifford conjugate, an anti-automorphism equivalent to the
-- composition of 'gradeInvolution' and 'reversion'.
conjugate :: (Ring a) => Multivector p q r a -> Multivector p q r a
conjugate = involution (\ n -> odd ((n+1) `div` 2))

-- | Square of the magnitude of a multi-vector, equal to
-- @'scalarPart' ('reversion' x * x)@.
-- This may be negative.
magnitude :: (KnownNat p, KnownNat q, KnownNat r, Ring a) =>
    Multivector p q r a -> a
magnitude mv@(MV xs) =
    sum [quad negs nulls x * xv * xv | (x, xv) <- Map.assocs xs]
  where
    negs = negBasisSet mv
    nulls = nullBasisSet mv

-- https://core.ac.uk/download/pdf/74374477.pdf
-- "Multivector and multivector matrix inverses in real Clifford algebras"
-- Eckhard Hitzer and Stephen Sangwine
-- Applied Mathematics and Computation 311(2017):375-389

-- | The reciprocal of a multivector, if such exists and is unique.
-- 'maybeRecip' returns 'Nothing' if the multivector is a factor of zero
-- (so no reciprocal exists).
--
-- This implementation is incomplete, and fails for values containing
-- more than 5 different basis vectors.
maybeRecip :: (KnownNat p, KnownNat q, Eq a, Field a) =>
    Multivector p q 0 a -> Maybe (Multivector p q 0 a)
maybeRecip v
  | g <= 5 = do
    let vv = v * v'
    let s = scalarPart vv
    guard (s /= 0)
    return (v' * scalar (recip s))
  -- add more?
  | otherwise = error $ "maybeRecip: not defined for grade " ++ show g
  where
    g = maxGrade v
    -- value such that v*v' is scalar
    v'
      | g == 0 = 1
      | g == 1 = c1
      | g == 2 = c1
      | g == 3 = c3
      | g == 4 = c1 * involution (>= 3) (v * c1)
      | g == 5 = c3 * involution (\ n -> n `mod` 3 == 1) (v * c3)
      | otherwise = error $ "not defined for grade " ++ show g
    c1 = conjugate v
    c3 = c1 * gradeInvolution v * reversion v