{-
   Haskell Dynamics Engine, Copyright (C) 2007 Ruben Henner Zilibowitz
   All rights reserved. Email: rubenz@cse.unsw.edu.au Web: www.cse.unsw.edu.au/~rubenz
   
   This is free software; you can redistribute it and/or modify it
   under the terms of the GNU General Public License that is include with this
   package in the file LICENSE-GPL.TXT.
-}

{-
Code for solving the Linear Complementarity Problem.

Given a matrix A and a vector b find a vector f such that:
a = Af + b
a >= 0
f >= 0
a.f = 0

This algorithm should succeed provided A is positive definite. This may not be a necessary condition however.
The algorithm comes from the paper "Fast Contact Force Computation for Nonpenetrating Rigid Bodies" by David Baraff.

Author: Ruben Henner Zilibowitz
-}

module LCP where

import Data.Set hiding (null,elems,map)
import Gausselim
import Data.Array

{-
Given A and b find an f such that:
a = Af + b
a >= 0
f >= 0
a.f = 0
-}

{- The difference between 1 and the least value greater than 1 that is
   representable in the given floating point type, b**1-p.  -}
epsilon :: Double
epsilon = 1e-9

{- Maximum representable finite floating-point number,

	(1 - b**-p) * b**emax
-}
infinity :: Double
infinity = 1e9

-- max number of runs through while_d_drive_to_zero_d or drive_to_zero
max_depth :: Int
max_depth = 25

-- compute_forces
compute_forces :: Matrix Double -> Vector Double -> Vector Double
compute_forces matA vecb = let (f,_,_,_) = solve_lcp matA vecb in f

-- compute_forces_with_holonomic
compute_forces_with_holonomic :: Int -> Matrix Double -> Vector Double -> Vector Double
compute_forces_with_holonomic n matA vecb = let (f,_,_,_) = while_d_drive_to_zero_d max_depth (vecf ++ (replicate m 0),(replicate n 0) ++ (drop n vecb),fromList [1..n],empty) matA vecb in f
   where
      matA' = take n (map (take n) matA)
      vecb' = take n vecb
      vecf = gauss matA' vecb'
      m = (length vecb) - n

-- solve_lcp
solve_lcp :: Matrix Double -> Vector Double -> (Vector Double, Vector Double, Set Int, Set Int)
solve_lcp matA vecb = while_d_drive_to_zero_d max_depth (replicate n 0,vecb,empty,empty) matA vecb
   where n = length vecb

while_d_drive_to_zero_d :: Int -> (Vector Double,Vector Double,Set Int,Set Int) -> Matrix Double -> Vector Double -> (Vector Double,Vector Double,Set Int,Set Int)
while_d_drive_to_zero_d 0 _ matA vecb = error ("LCP solver: while_d_drive_to_zero_d max depth exceeded with: A=" ++ (show matA) ++ " b=" ++ (show vecb))
while_d_drive_to_zero_d n dat@(f,a,c,nc) matA vecb
   | (null ds) = dat
   | otherwise = while_d_drive_to_zero_d (n-1) (drive_to_zero max_depth (head ds) dat matA vecb) matA vecb
   where ds = [d | (d,x) <- zip [0..] a, -x > epsilon]

-- drive_to_zero
drive_to_zero :: Int -> Int -> (Vector Double,Vector Double,Set Int,Set Int) -> Matrix Double -> Vector Double -> (Vector Double,Vector Double,Set Int,Set Int)
drive_to_zero 0 _ _ matA vecb = error ("LCP solver: drive_to_zero max depth exceeded with: A=" ++ (show matA) ++ " b=" ++ (show vecb))
drive_to_zero n d (f,a,c,nc) matA vecb
   | (j == (-1)) = error ("LCP solver: drive_to_zero infinite force with: A=" ++ (show matA) ++ " b=" ++ (show vecb) ++ " d=" ++ (show d))
   | (j `member` c) = drive_to_zero (n-1) d (f',a',delete j c,insert j nc) matA vecb
   | (j `member` nc) = drive_to_zero (n-1) d (f',a',insert j c,delete j nc) matA vecb
   | otherwise = (f',a',insert j c,nc)
   where
      df = fdirection d c nc matA vecb
      da = matA `matXvec` df
      (s,j) = maxstep (f,a,c,nc) df da d
      f' = f `vecsum` (s `scalarXvec` df)
      a' = a `vecsum` (s `scalarXvec` da)

-- fdirection
fdirection :: Int -> Set Int -> Set Int -> Matrix Double -> Vector Double -> Vector Double
fdirection d c nc matA vecb = df
   where
      cList = toList c
      matA11 = [[(matA!!j)!!k | k <- cList] | j <- cList]
      vecv1 = [-((matA!!j)!!d) | j <- cList]
      x = gauss matA11 vecv1
      dfArr = accumArray (flip const) 0 (0,n-1) ((d,1):(zip cList x))
      n = length vecb
      df = elems dfArr

-- maxstep
maxstep :: (Vector Double,Vector Double,Set Int,Set Int) -> Vector Double -> Vector Double -> Int -> (Double,Int)
maxstep (f,a,c,nc) df da d = result
   where
      (s0,j0) = if (da!!d > 0) then (-(a!!d)/(da!!d),d) else (infinity,-1)
      cList = toList c
      ncList = toList nc
      tests = (s0,j0) : ([(-(f!!i)/(df!!i),i) | i <- cList, df!!i < 0] ++
                         [(-(a!!i)/(da!!i),i) | i <- ncList, df!!i < 0])
      result = minimum tests

-- matrix and vector functions
matXvec :: Matrix Double -> Vector Double -> Vector Double
matXvec mat vec = map (sum.(zipWith (*) vec)) mat

scalarXvec :: Double -> Vector Double -> Vector Double
scalarXvec x = map (x*)

vecsum :: Vector Double -> Vector Double -> Vector Double
vecsum = zipWith (+)

{-
--- tests
matM0 = [[2,1],[1,2 :: Double]]
q0 = [-5,-6 :: Double]

matM1 = [[1,0,0],[2,1,0],[2,2,1]] :: [[Double]]
q1 = [-8,-12,-14] :: [Double]

-- cycling occurs on this example
matM2 = [[1,2,0],[0,1,2],[2,0,1]] :: [[Double]]
q2 = [-1,-1,-1] :: [Double]

matM3 = [[2,1,1],[1,2,1],[1,1,2]] :: [[Double]]
q3 = [-4,-5,-1] :: [Double]

matM4 = [[1,-1,0,1],[1,0,0,1],[0,0,1,-3],[1,0,1,1]] :: [[Double]]
q4 = [-1,-2,-3,-4] :: [Double]

matM5 = [[10,0,-2],[2,0.1,-0.4],[0,0.2,0.1]] :: [[Double]]
q5 = [10,1,-1] :: [Double]

matM6 = [[1,0,2],[-2,1,4],[-4,2,9]] :: [[Double]]
q6 = [1,-1,3] :: [Double]

matM7 = [[1,2,-1,2],[-1,1,2,-1],[2,-1,1,2],[-1,2,-1,1]] :: [[Double]]
q7 = [-2,1,-1,2] :: [Double]
-}
{-
matM1 = [[34,16,4,1],
         [16,34,16,1],
         [4,16,8,1],
         [-1,-1,-1,0]] :: [[Double]]
q1 = [-66,-54,-20,1] :: [Double]

matM2 = [[0,0,2,2,1],
         [0,0,1,2,2],
         [1,2,0,0,0],
         [3,1,0,0,0],
         [2,3,0,0,0]] :: [[Double]]
q2 = [-1,-1,-1,-1,-1] :: [Double]

matM3 = [[-2,1],[1,-2]] :: [[Double]]
q3 = [4,-1] :: [Double]
-}