{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | QR factorization via Givens rotations.
--
module Linear.QR (
  givens_rotator,
  qr )
where

import qualified Algebra.Ring as Ring ( C )
import qualified Algebra.Algebraic as Algebraic ( C )
import Data.Vector.Fixed ( ifoldl )
import Data.Vector.Fixed.Cont ( Arity )
import NumericPrelude hiding ( (*) )

import Linear.Matrix (
  Mat(..),
  (*),
  (!!!),
  construct,
  identity_matrix,
  transpose )


-- | Construct a givens rotation matrix that will operate on row @i@
--   and column @j@. This is done to create zeros in some column of
--   the target matrix. You must also supply that column's @i@th and
--   @j@th entries as arguments.
--
--   Examples (Watkins, p. 193):
--
--   >>> import Linear.Matrix ( Mat2, fromList )
--   >>> let m  = givens_rotator 0 1 1 1 :: Mat2 Double
--   >>> let m2 = fromList [[1, -1],[1, 1]] :: Mat2 Double
--   >>> m == (1 / (sqrt 2) :: Double) *> m2
--   True
--
givens_rotator :: forall m a. (Ring.C a, Algebraic.C a, Arity m)
               => Int -> Int -> a -> a -> Mat m m a
givens_rotator i j xi xj =
  construct f
  where
    xnorm = sqrt $ xi^2 + xj^2
    c = xi / xnorm
    s = xj / xnorm

    f :: Int -> Int -> a
    f y z
      | y == i && z == i = c
      | y == j && z == j = c
      | y == i && z == j = negate s
      | y == j && z == i = s
      | y == z = fromInteger 1
      | otherwise = fromInteger 0


-- | Compute the QR factorization of a matrix (using Givens
--   rotations). This is accomplished with two folds: the first one
--   traverses the columns of the matrix from left to right, and the
--   second traverses the entries of the column from top to
--   bottom.
--
--   The state that is passed through the fold is the current (q,r)
--   factorization. We keep the pair updated by multiplying @q@ and
--   @r@ by the new rotator (or its transpose).
--
qr :: forall m n a. (Arity m, Arity n, Algebraic.C a, Ring.C a)
   => Mat m n a -> (Mat m m a, Mat m n a)
qr matrix =
  ifoldl col_function initial_qr columns
  where
    Mat columns = transpose matrix
    initial_qr = (identity_matrix, matrix)

    -- | Process the column and spit out the current QR
    --   factorization. In the first column, we want to get rotators
    --   Q12, Q13, Q14,... In the second column, we want rotators Q23,
    --   Q24, Q25,...
    col_function (q,r) col_idx col =
      ifoldl (f col_idx) (q,r) col

    -- | Process the entries in a column, doing basically the same
    --   thing as col_dunction does. It updates the QR factorization,
    --   maybe, and returns the current one.
    f col_idx (q,r) idx x
      | idx <= col_idx = (q,r) -- leave it alone.
      | otherwise =
          (q*rotator, (transpose rotator)*r)
          where
            rotator :: Mat m m a
            rotator = givens_rotator col_idx idx (r !!! (idx, col_idx)) x