{-# LANGUAGE TypeOperators #-}
{- |
Do not import this module. It is only for demonstration purposes.
-}
module Numeric.LAPACK.Example.EconomicAllocation where

import qualified Numeric.LAPACK.Matrix.Square as Square
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import Numeric.LAPACK.Matrix (ZeroInt, (#-#), (#*|), (#\|), (|||))
import Numeric.LAPACK.Vector ((|-|))
import Numeric.LAPACK.Format ((##))

import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Shape ((:+:)((:+:)))


type ZeroInt2 = ZeroInt:+:ZeroInt
type Vector sh = Vector.Vector sh Double
type Matrix height width = Matrix.General height width Double
type SquareMatrix size = Square.Square size Double


balances0 :: Vector ZeroInt2
balances0 =
   Vector.fromList (Matrix.zeroInt 2 :+: Matrix.zeroInt 2)
      [100000, 90000, -50000, -120000]

expenses0 :: Matrix ZeroInt ZeroInt2
expenses0 =
   Matrix.fromList (Matrix.zeroInt 2) (Matrix.zeroInt 2 :+: Matrix.zeroInt 2) $
   [16000,  4000,  8000, 12000,
    10000, 30000, 40000, 20000]

normalize ::
   (Eq height, Shape.C height, Shape.C width) =>
   Matrix height width -> Matrix height width
normalize x = Matrix.scaleRows (Array.map recip (Matrix.rowSums x)) x


completeIdSquare ::
   (Shape.C sh0, Eq sh0, Shape.C sh1, Eq sh1) =>
   Matrix (sh0:+:sh1) sh1 -> SquareMatrix (sh0:+:sh1)
completeIdSquare x =
   Square.fromGeneral $
      (Matrix.takeLeft $ Matrix.fromFull $ Square.identityFromHeight x)
      |||
      x

iterated ::
   (Shape.C sh0, Eq sh0, Shape.C sh1, Eq sh1) =>
   Matrix sh1 (sh0:+:sh1) -> Vector (sh0:+:sh1) -> Vector (sh0:+:sh1)
iterated expenses =
   -- 'Stream.head' would be total
   head . dropWhile ((>=1e-5) . Vector.normInf . Vector.takeRight) .
   iterate (completeIdSquare (Matrix.transpose (normalize expenses)) #*|)


compensated ::
   (Shape.C sh0, Eq sh0, Shape.C sh1, Eq sh1) =>
   Matrix sh1 (sh0:+:sh1) -> Vector (sh0:+:sh1) -> Vector sh0
compensated expenses balances =
   let a = Matrix.transpose $ normalize expenses
       p = Matrix.takeTop a
       k = Square.fromGeneral $ Matrix.takeBottom a
       x = Vector.takeLeft balances
       y = Vector.takeRight balances
   in x |-| p #*| (k #-# Square.identityFrom k) #\| y


main :: IO ()
main = do
   iterated expenses0 balances0 ## "%10.2f"
   compensated expenses0 balances0 ## "%10.2f"