{-# LANGUAGE FlexibleInstances, FlexibleContexts, TypeSynonymInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.LinearAlgebra.Tensor
-- Copyright   :  (c) Alberto Ruiz 2009
-- License     :  GPL
--
-- Maintainer  :  Alberto Ruiz <aruiz@um.es>
-- Stability   :  experimental
--
-- Tensor computations. Indices can only be contracted if they are of different 'Variant' type.
--
-----------------------------------------------------------------------------


module Numeric.LinearAlgebra.Tensor (
    -- * The Tensor type
    Tensor, Variant(..),
    listTensor,
    -- * Tensor creation utilities
    superindex, subindex,
    vector, covector, transf,
    -- * Index manipulation
    switch, cov, contrav, forget,
    -- * General array operations
    module Numeric.LinearAlgebra.Array
) where

import Numeric.LinearAlgebra.Array.Internal
import Numeric.LinearAlgebra hiding (rank)
import Numeric.LinearAlgebra.Array

type Tensor t = NArray Variant t

data Variant = Co | Contra deriving (Eq)

instance Compat Variant where
    compat d1 d2 = iDim d1 == iDim d2 && iType d1 /= iType d2

instance Show (Idx Variant) where
    show (Idx n s Co)     = show n ++ "_" ++ s
    show (Idx n s Contra) = show n ++ "^" ++ s

instance (Coord t) => Show (Tensor t) where
    show t | null (dims t) = show (coords t @>0)
           | otherwise = "listTensor " ++ show (dims t) ++ " "++ show (toList (coords t))

flipV Co = Contra
flipV Contra = Co

-- | Creates a tensor from a list of dimensions and a list of coordinates.
-- A positive dimension means that the index is assumed to be contravariant (vector-like), and
-- a negative dimension means that the index is assumed to be covariant (like a linear function, or covector). Contractions can only be performed between indices of different type.
listTensor :: Coord t
           => [Int] -- ^ dimensions
           -> [t]   -- ^ coordinates
           -> Tensor t
listTensor ds cs = mkNArray dms (product ds' |> (cs ++ repeat 0))
    where dms = zipWith3 Idx ds' (map show [1::Int ..]) (map f ds)
          ds' = map abs ds
          f n | n>0       = Contra
              | otherwise = Co

-- | Create an 'Tensor' from a list of parts with a contravariant index (@superindex = 'newIndex' 'Contra'@).
superindex :: Coord t => Name -> [Tensor t] -> Tensor t
superindex = newIndex Contra

-- | Create an 'Tensor' from a list of parts with a covariant index (@subindex = 'newIndex' 'Co'@).
subindex :: Coord t => Name -> [Tensor t] -> Tensor t
subindex   = newIndex Co



-- | Change the 'Variant' nature of all dimensions to the opposite ones.
switch :: Tensor t -> Tensor t
switch = mapTypes flipV

-- | Make all dimensions covariant.
cov :: NArray i t -> Tensor t
cov     = mapTypes (const Co)

-- | Make all dimensions contravariant.
contrav :: NArray i t -> Tensor t
contrav = mapTypes (const Contra)

-- | Remove the 'Variant' nature of coordinates.
forget :: NArray i t -> Array t
forget = mapTypes (const None)

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

-- | Create a contravariant rank-1 tensor from a list of coordinates.
vector :: [Double] -> Tensor Double
vector   = fromVector Contra . fromList

-- | Create a covariant rank-1 tensor from a list of coordinates.
covector :: [Double] -> Tensor Double
covector = fromVector Co . fromList

-- | Create a 1-contravariant, 1-covariant rank-2 tensor from list of lists of coordinates.
transf :: [[Double]] -> Tensor Double
transf   = fromMatrix Contra Co . fromLists