{-# LANGUAGE NoImplicitPrelude,RebindableSyntax #-}
{-# LANGUAGE ScopedTypeVariables,FlexibleContexts #-}
module Algebra.Linear
  (
    Relation(..)
  , satisfies
  , solve
  , dependencies
  , equations
  , inverseImage
  
  , Matrix
  , Vector
  , matrixProduct
  , matrixVector
  , innerProduct
  
  , identity
  , matrixFromFunction
  , affine
  , permute
  , rowSwap
  
  , invert
  , determinant
  , adjoint
  , diagonal
  ) where


import           Auxiliary (δ,findAmong,headE,adorn)

import qualified Algebra.Field
import qualified Algebra.Lattice
import qualified Algebra.Module
import qualified Algebra.Ring
import           NumericPrelude

import           Control.Applicative ((<$>))
import           Control.Arrow       (first,second)
import           Data.List
  (
    partition
  , find
  , transpose
  , genericLength
  , genericTake
  , genericDrop
  , genericReplicate
  )
import qualified Data.Map as Map
import           Data.Proxy          (Proxy(Proxy))
import           Data.Reflection     (Reifies,reflect)


type Vector k
  = [k]

type Matrix k
  = [Vector k]

-- A relation among vectors of degree 'd' over field 'k'
data Relation k
  = Relation { getRelation :: [k] }
  deriving (Show)

satisfies :: (Algebra.Ring.C k,Eq k) => Vector k -> Relation k -> Bool
satisfies v (Relation r) = innerProduct v r == 0

-- | Calculate all dependencies among the given vectors of degree d.
dependencies :: (Algebra.Field.C k,Eq k) => Integer -> [Vector k] -> [Relation k]
dependencies d = map (Relation . genericDrop d) . filter (all (== zero) . genericTake d) . (\ (r,_,_) -> r) . reduce . adorn

-- | Calculate the equations satisfied by the subspace spanned by the given vectors of degree d.
equations :: (Algebra.Field.C k,Eq k) => Integer -> [Vector k] -> [Relation k]
equations d vs = dependencies (genericLength vs + 1) . transpose . (:) (genericReplicate d zero) $ vs

-- | Solve the given equations, in the form of a basis of vectors of degree d.
solve :: (Algebra.Field.C k,Eq k) => Integer -> [Relation k] -> [Vector k]
solve d = map getRelation . equations d . map getRelation

inverseImage :: (Algebra.Field.C k,Eq k) => Matrix k -> Vector k -> Vector k
inverseImage a = solveUpperTriangular u . matrixVector b where
  (u,b,_) = reduce a

invert :: (Algebra.Field.C k,Eq k) => Matrix k -> Maybe (Matrix k)
invert m = fmap strip . process . (\ (r,_,_) -> r) . reduce . adorn $ m where
  n = length m
  process x = go x where
    go [v] = Just [v]
    go (r@(pivot : r') : rest)
      | pivot == 1 = do
        m' <- go (map tail rest)
        rowsToAdd <- sequence . zipWith3 f [0 ..] r' $ m'
        return $ (1 : foldr (zipWith (+)) r' rowsToAdd) : map (0 :) m'
      | otherwise = Nothing
    f i c v
      | v₀ == 0   = Nothing
      | otherwise = Just $ map ((*) (negate c / v₀)) v
      where
        v₀ = v !! i
  strip = map (drop n)

determinant :: (Algebra.Field.C k,Eq k) => Matrix k -> k
determinant x = case reduce x of (_,_,σ) -> σ

adjoint :: (Algebra.Lattice.C k,Algebra.Field.C k,Eq k) => Matrix k -> Maybe (Matrix k)
adjoint a = (abs (determinant a) *>) <$> invert a

diagonal :: Matrix k -> [k]
diagonal []               = []
diagonal ((d : _) : rest) = d : diagonal (map tail rest)

matrixProduct :: (Algebra.Ring.C k) => Matrix k -> Matrix k -> Matrix k
matrixProduct a b = map (($ transpose b) . map . innerProduct) a 

matrixVector :: (Algebra.Ring.C k) => Matrix k -> Vector k -> Vector k
matrixVector a = map (\ [x] -> x) . matrixProduct a . map (: [])

innerProduct :: (Algebra.Ring.C k) => Vector k -> Vector k -> k
innerProduct a b = sum (zipWith (*) a b)

matrixFromFunction :: Int -> Int -> (Int -> Int -> k) -> Matrix k
matrixFromFunction m n f = [[f i j | j <- [1 .. n]] | i <- [1 .. m]]

identity :: (Algebra.Ring.C k) => Int -> Matrix k
identity n = matrixFromFunction n n δ


-- Triangular decomposition

solveUpperTriangular :: (Algebra.Field.C k,Eq k) => Matrix k -> Vector k -> Vector k
solveUpperTriangular u x = reverse $ solveLowerTriangular (reverse . map reverse $ u) (reverse x)

solveLowerTriangular :: (Algebra.Field.C k,Eq k) => Matrix k -> Vector k -> Vector k
solveLowerTriangular l x = go l x [] where
  go []        []        q = q
  go (lᵢ : l') (xᵢ : x') q = go l' x' (q ++ [qᵢ]) where
    (lFirst,lᵢᵢ : _) = splitAt (length q) lᵢ
    y = xᵢ - innerProduct q lFirst
    qᵢ
      | lᵢᵢ == 0  = error "Linear.solveLowerTriangular: zero on diagonal"
      | otherwise = y / lᵢᵢ

-- Compute the row echelon form of the matrix,
-- together with the basis transformation matrix,
-- and its determinant.
reduce :: (Algebra.Field.C k,Eq k) => Matrix k -> (Matrix k,Matrix k,k)
reduce []          = ([],[],1)
reduce xs@([] : _) = (xs,identity (length xs),1)
reduce vs          = case nonZero of
  []                    -> (\ (x,u,_σ) -> (map (0 :) x,u,0)) $ reduce (map tail vs)
  (v@(v₀ : _),i) : []   -> let
    (h,u,σ) = reduce (map (tail . fst) startZero)
    sign = if i == 1 then 1 else -1
   in
    ( (map (/ v₀) v :) . map (0 :) $ h
    , normalisation v₀ `matrixProduct` rowSwap n (1,i) `matrixProduct` shift u
    ,sign * v₀ * σ
    )
  (v@(v₀ : _),i) : rest -> let
    (reduced,translates) = unzip . flip map rest $ \ (x@(x₀ : _),j) -> let
      c = x₀ / v₀ in ((zipWith (\ vᵢ xᵢ -> xᵢ - c * vᵢ) v x,j),(j,c))
    (h,u,σ) = reduce $ v : map fst reduced ++ map fst startZero
    permutation = i : map snd reduced ++ map snd startZero
   in
    ( h
    , u
        `matrixProduct` permute permutation
        `matrixProduct` affine n i (map (second negate) translates)
    , permutationSign permutation * σ
    )
 where
  (startZero,nonZero) = partition ((==) 0 . head . fst) . flip zip [1 ..] $ vs
  normalisation v₀ = matrixFromFunction n n f where
    f 1 1 = 1 / v₀
    f i j = δ i j
  shift u = (1 : replicate (n - 1) 0) : map (0 :) u
  n = length vs

rowSwap :: (Algebra.Ring.C k) => Int -> (Int,Int) -> Matrix k
rowSwap n (k,l) = matrixFromFunction n n f where
  f i j
    | i == k && j == l           = one
    | i == l && j == k           = one
    | i == j && i /= k && i /= l = one
    | otherwise                  = zero

affine :: (Algebra.Ring.C k) => Int -> Int -> [(Int,k)] -> Matrix k
affine n k ps = matrixFromFunction n n $ \ i j -> δ i j + c i j where
  m = Map.fromList ps
  c i j = case (Map.lookup i m,j == k) of
    (Just cᵢ,True) -> cᵢ
    _              -> zero

permute :: (Algebra.Ring.C k) => [Int] -> Matrix k
permute is = matrixFromFunction n n (\ i j -> δ (is !! (i - 1)) j) where
  n = length is

permutationSign :: (Algebra.Ring.C k) => [Int] -> k
permutationSign [] = one
permutationSign (x : xs) = ε * permutationSign xs where
  inversions = length $ filter (x >) xs
  ε
    | even inversions = one
    | otherwise       = negate one

-- Example
x :: Matrix Rational
x = 
  [
    [0, 3,-6, 6,4,5 ]
  , [3,-7, 8,-5,8,9 ]
  , [3,-9,12,-9,6,15]
  ]