```{-# LANGUAGE TypeOperators, FlexibleContexts, TypeFamilies #-}
{-# OPTIONS_GHC -Wall -fno-warn-orphans #-}
-- {-# OPTIONS_GHC -funbox-strict-fields #-}
-- {-# OPTIONS_GHC -ddump-simpl-stats -ddump-simpl #-}
----------------------------------------------------------------------
-- |
-- Module      :  Data.LinearMap
-- Copyright   :  (c) Conal Elliott 2008
--
-- Maintainer  :  conal@conal.net
-- Stability   :  experimental
--
-- Linear maps
----------------------------------------------------------------------

module Data.LinearMap
( (:-*) , linear, lapply, atBasis, idL, (*.*)
, liftMS, liftMS2, liftMS3
, liftL, liftL2, liftL3
) where

import Control.Applicative ((<\$>),Applicative,liftA2,liftA3)
import Control.Arrow       (first)

import Data.MemoTrie      ((:->:)(..))
import Data.VectorSpace   (VectorSpace(..))
import Data.Basis         (HasBasis(..), linearCombo)

-- Linear maps are almost but not quite a Control.Category.  The type
-- class constraints interfere.  They're almost an Arrow also, but for the
-- constraints and the generality of arr.

-- | An optional additive value
type MSum a = Maybe (Sum a)

-- nsum :: MSum a
-- nsum = Nothing

jsum :: a -> MSum a
jsum = Just . Sum

-- | Linear map, represented as an optional memo-trie from basis to
-- values, where 'Nothing' means the zero map (an optimization).
type u :-* v = MSum (Basis u :->: v)

-- TODO: Try a partial trie instead, excluding (known) zero elements.
-- Then 'lapply' could be much faster for sparse situations.  Make sure to
-- correctly sum them.  It'd be more like Jason Foutz's formulation
-- <http://metavar.blogspot.com/2008/02/higher-order-multivariate-automatic.html>
-- which uses in @IntMap@.

-- PROBLEM: u :-* v is a type synonym, and Basis is an associated type synonym, resulting in a subtle
-- ambiguity: u:-*v == u':-*v' does not imply that u==u', since Basis
-- might map different types to the same basis (e.g., Float & Double).
--
-- Work in progress.  See NewLinearMap.hs

-- | Function (assumed linear) as linear map.
linear :: (HasBasis u, HasTrie (Basis u)) =>
(u -> v) -> (u :-* v)
linear f = jsum (trie (f . basisValue))

atZ :: AdditiveGroup b => (a -> b) -> (MSum a -> b)
atZ f = maybe zeroV (f . getSum)

-- atZ :: AdditiveGroup b => (a -> b) -> (a -> b)
-- atZ = id

-- | Evaluate a linear map on a basis element.  I've loosened the type to
-- work around a typing problem in 'derivAtBasis'.
-- atBasis :: (AdditiveGroup v, HasTrie (Basis u)) =>
--            (u :-* v) -> Basis u -> v
atBasis :: (HasTrie a, AdditiveGroup b) => MSum (a :->: b) -> a -> b
m `atBasis` b = atZ (`untrie` b) m

-- | Apply a linear map to a vector.
lapply :: ( VectorSpace v, Scalar u ~ Scalar v
, HasBasis u, HasTrie (Basis u) ) =>
(u :-* v) -> (u -> v)
lapply = atZ lapply'

-- Handy for 'lapply' and '(*.*)'.
lapply' :: ( VectorSpace v, Scalar u ~ Scalar v
, HasBasis u, HasTrie (Basis u) ) =>
(Basis u :->: v) -> (u -> v)
lapply' tr = linearCombo . fmap (first (untrie tr)) . decompose

-- Identity linear map
idL :: (HasBasis u, HasTrie (Basis u)) =>
u :-* u
idL = linear id

infixr 9 *.*
-- | Compose linear maps
(*.*) :: ( HasBasis u, HasTrie (Basis u)
, HasBasis v, HasTrie (Basis v)
, VectorSpace w
, Scalar v ~ Scalar w ) =>
(v :-* w) -> (u :-* v) -> (u :-* w)

-- Simple definition, but only optimizes out uv == zero
--
-- (*.*) vw = (fmap.fmap) (lapply vw)

-- Instead, use Nothing/zero if /either/ map is zeroV (exploiting linearity
-- when uv == zeroV.)

-- Nothing       *.* _             = Nothing
-- _             *.* Nothing       = Nothing
-- Just (Sum vw) *.* Just (Sum uv) = Just (Sum (lapply' vw <\$> uv))

-- (*.*) = liftA2 (\ (Sum vw) (Sum uv) -> Sum (lapply' vw <\$> uv))

-- (*.*) = (liftA2.inSum2) (\ vw uv -> lapply' vw <\$> uv)
(*.*) = (liftA2.inSum2) (\ vw uv -> lapply' vw <\$> uv)

-- (*.*) = (liftA2.inSum2) (\ vw -> fmap (lapply' vw))

-- (*.*) = (liftA2.inSum2) (fmap . lapply')

-- It may be helpful that @lapply vw@ is evaluated just once and not
-- once per uv.  'untrie' can strip off all of its trie constructors.

-- Less efficient definition:
--
--   vw `compL` uv = linear (lapply vw . lapply uv)
--
--   i.e., compL = inL2 (.)
--
-- The problem with these definitions is that basis elements get converted
-- to values and then decomposed, followed by recombination of the
-- results.

(a -> b)
-> (MSum a -> MSum b)
-- liftMS _ Nothing = Nothing
-- liftMS h ma = Just (Sum (h (z ma)))

liftMS = fmap.fmap

(a -> b -> c) ->
(MSum a -> MSum b -> MSum c)
liftMS2 _ Nothing Nothing = Nothing
liftMS2 h ma mb = Just (Sum (h (fromMS ma) (fromMS mb)))

(a -> b -> c -> d) ->
(MSum a -> MSum b -> MSum c -> MSum d)
liftMS3 _ Nothing Nothing Nothing = Nothing
liftMS3 h ma mb mc = Just (Sum (h (fromMS ma) (fromMS mb) (fromMS mc)))

fromMS :: AdditiveGroup u => MSum u -> u
fromMS Nothing        = zeroV
fromMS (Just (Sum u)) = u

-- | Apply a linear function to each element of a linear map.
-- @liftL f l == linear f *.* l@, but works more efficiently.
liftL :: (Functor f, AdditiveGroup (f a)) =>
(a -> b) -> MSum (f a) -> MSum (f b)
liftL = liftMS . fmap

-- | Apply a linear binary function (not to be confused with a bilinear
-- function) to each element of a linear map.