{-# OPTIONS -XParallelListComp #-}
module Physics.Hpysics.Utils where

import Physics.Hpysics.Types
import Data.List
import Data.Ord

-- * General Haskell
rotate :: [a] -> [a]
rotate [] = []
rotate (x:xs) = xs++[x]

swap :: (a,b) -> (b,a)
swap (x,y) = (y,x)

swapInList :: Int -> Int -> [a] -> [a]
swapInList i j list = [ if k == i then list !! j
                   else if k == j then list !! i
                   else l | k <- [0..] | l <- list ]
updateInList :: (a->a) -> Int -> [a] -> [a]
updateInList f i list = [ if k == i then f l else l | k <- [0..] | l <- list ]

infix 4 `equal`
equal :: FloatType -> FloatType -> Bool
x`equal`y = abs(x-y)<floatEpsilon

-- * Linear algebra

-- | solves system of two linear equations
--
-- > {a11 x1 + a12 x2 = b1
-- > {a21 x1 + a22 x2 = b2
solveLinear2 ::
    FloatType -> FloatType -> FloatType -> 
    FloatType -> FloatType -> FloatType -> 
    Maybe (FloatType, FloatType)
solveLinear2
    a11 a12 b1
    a21 a22 b2 = let
    det = a11*a22 - a12 * a21 in
    if (det /= 0) -- TODO better comparison
        then Just ((b1*a22-b2*a12)/det, (a11*b2-a21*b1)/det)
        else Nothing

-- * Vectors
infix  4 `vecEqual`
infixl 6 `add`, `sub`
infixl 7 `mul`, <.>, `scale`

vecEqual :: Vec -> Vec -> Bool
Vec x1 y1 z1 `vecEqual` Vec x2 y2 z2 = x1`equal`x2 && y1`equal`y2 && z1`equal`z2
vecMap :: (FloatType -> FloatType) -> Vec -> Vec
f `vecMap` (Vec x y z) = Vec (f x) (f y) (f z)
zipVecWith :: (FloatType -> FloatType -> FloatType) -> Vec -> Vec -> Vec
zipVecWith f (Vec x1 y1 z1) (Vec x2 y2 z2) = Vec (f x1 x2) (f y1 y2) (f z1 z2)
scale :: FloatType -> Vec -> Vec
scale a = vecMap (*a)
vecSum :: Vec -> FloatType
vecSum (Vec x y z) = x+y+z
addConstant :: FloatType -> Vec -> Vec
addConstant a = vecMap (+a)
add :: Vec -> Vec -> Vec
add = zipVecWith (+)
sub :: Vec -> Vec -> Vec
sub = zipVecWith (-)
mul :: Vec -> Vec -> Vec
mul = zipVecWith (*)
divide :: Vec -> Vec -> Vec
divide = zipVecWith (/)
neg :: Vec -> Vec
neg = vecMap negate
cross :: Vec -> Vec -> Vec
cross (Vec x1 y1 z1) (Vec x2 y2 z2) = Vec (y1*z2-y2*z1) (z1*x2-z2*x1) (x1*y2-x2*y1)
(<.>) :: Vec -> Vec -> FloatType
(<.>) x y = vecSum $ x `mul` y
dotSquare x = x <.> x
norm :: Vec -> FloatType
norm = sqrt . dotSquare
normalize :: Vec -> Vec
normalize x = 1/norm x `scale` x
zeroVector :: Vec
zeroVector = Vec 0 0 0
-- | Linear (convex) combination
lc :: Vec -> Vec -> FloatType -> Vec
lc a b lambda = lambda `scale` a `add` (1-lambda) `scale` b

-- | Find a vector, orthogonal to given /unit/ one
orthogonal :: Vec -> Vec
orthogonal n@(Vec x y z) = let
    -- some vector not collinear to n. Why 10? Dunno...
    v = if abs x `equal` 1 then Vec 0 1 0 else Vec (x+10) 0 0
    in normalize $ v`sub`(v<.>n`scale`n)

-- | Build orthonormal basis of three vectors, with third
-- being parallel to given vector.
makeOrthoBasis :: Vec -> (Vec, Vec, Vec)
makeOrthoBasis k = let
    k' = normalize k
    i = orthogonal k'
    j = cross i k'
    in (i,j,k')

-- * Vec4
infixl 6 `add4`, `sub4`
infixl 7 `mul4`, `dot4`, `scale4`

vec4Map :: (FloatType -> FloatType) -> Vec4 -> Vec4
f `vec4Map` (Vec4 w x y z) = Vec4 (f w) (f x) (f y) (f z)
zipVec4With :: (FloatType -> FloatType -> FloatType) -> Vec4 -> Vec4 -> Vec4
zipVec4With f (Vec4 w1 x1 y1 z1) (Vec4 w2 x2 y2 z2) = Vec4 (f w1 w2) (f x1 x2) (f y1 y2) (f z1 z2)
scale4 :: FloatType -> Vec4 -> Vec4
scale4 a = vec4Map (*a)
vec4Sum :: Vec4 -> FloatType
vec4Sum (Vec4 w x y z) = w+x+y+z
addConstant4 :: FloatType -> Vec4 -> Vec4
addConstant4 a = vec4Map (+a)
add4 :: Vec4 -> Vec4 -> Vec4
add4 = zipVec4With (+)
sub4 :: Vec4 -> Vec4 -> Vec4
sub4 = zipVec4With (-)
mul4 :: Vec4 -> Vec4 -> Vec4
mul4 = zipVec4With (*)
divide4 :: Vec4 -> Vec4 -> Vec4
divide4 = zipVec4With (/)
neg4 :: Vec4 -> Vec4
neg4 = vec4Map negate
dot4 :: Vec4 -> Vec4 -> FloatType
dot4 x y = vec4Sum $ x `mul4` y
zeroVec4 :: Vec4
zeroVec4 = Vec4 0 0 0 0
unitVec4 :: Vec4
unitVec4 = Vec4 1 0 0 0

-- * Quaternions
qnorm :: Quaternion -> FloatType
qnorm (Vec4 w x y z) = sqrt $ w^2+x^2+y^2+z^2
qnormalize q = 1/qnorm q `scale4` q

-- * Matrices
mat43xVec :: Mat43 -> Vec -> Vec4
mat43xVec
    (Mat43
        a11 a12 a13
        a21 a22 a23
        a31 a32 a33
        a41 a42 a43)
    (Vec b1 b2 b3) =
    (Vec4
        (a11*b1+a12*b2+a13*b3)
        (a21*b1+a22*b2+a23*b3)
        (a31*b1+a32*b2+a33*b3)
        (a41*b1+a42*b2+a43*b3))
mat33xVec :: Mat33 -> Vec -> Vec
mat33xVec
    (Mat33
        a11 a12 a13
        a21 a22 a23
        a31 a32 a33)
    (Vec b1 b2 b3) =
    (Vec
        (a11*b1+a12*b2+a13*b3)
        (a21*b1+a22*b2+a23*b3)
        (a31*b1+a32*b2+a33*b3))
matScale :: FloatType -> Mat33 -> Mat33
matScale l 
    (Mat33
        a11 a12 a13
        a21 a22 a23
        a31 a32 a33)
        =
    (Mat33
        (l*a11) (l*a12) (l*a13)
        (l*a21) (l*a22) (l*a23)
        (l*a31) (l*a32) (l*a33))
unit33 = Mat33 1 0 0 0 1 0 0 0 1

-- * Geometry
distance :: Vec -> Vec -> FloatType
distance x y = norm $ x `sub` y
distToPlane, distToPlaneSigned :: Vec -> Plane -> FloatType
distToPlane point plane = abs $ distToPlaneSigned point plane
distToPlaneSigned point (Plane n s) = point <.> n - s
projectToPlane :: Plane -> Vec -> Vec
projectToPlane plane@(Plane n _) x = x `sub` distToPlaneSigned x plane `scale` n
-- point on a segment [a, b] is parametrized by lambda:
-- v2=lambda*a+(1-lambda)*b, lambda is from [0,1]
-- square of the distance between v1 and v2 is quadratic polynomial
-- which has the form
-- (a-b)^2 lambda^2+2(a-b)(b-v1)lambda+(b-v1)^2
-- This function finds lambda which minimises the distance
projectToLineLambda :: Vec -> Vec -> Vec -> FloatType
projectToLineLambda a b v1 = let a_b = a`sub`b in
    - a_b <.> (b`sub`v1) / dotSquare a_b
projectToLine :: Vec -> Vec -> Vec -> Vec
projectToLine a b v1 = lc a b lambda
    where lambda = projectToLineLambda a b v1
distToLine :: Vec -> Vec -> Vec -> FloatType
distToLine a b v1 = distance v1 $ projectToLine a b v1
planeLineIntersectionLambda :: Plane -> Vec -> Vec -> Maybe FloatType
planeLineIntersectionLambda (Plane n s) a b = let denom = (a`sub`b)<.>n
    in if (denom /= 0) then Just$(s - b<.>n)/denom else
        if a<.>n == s then Just 1 else Nothing
planeLineIntersection :: Plane -> Vec -> Vec -> Maybe Vec
planeLineIntersection plane a b = do
    lambda <- planeLineIntersectionLambda plane a b
    return $ lc a b lambda
rotationQ :: Vec -> FloatType -> Quaternion
rotationQ axis@(Vec x y z) angle = let
    a2 = angle/2
    sin_a2 = sin a2
    in Vec4 (cos a2) (x*sin_a2) (y*sin_a2) (z*sin_a2)
-- TODO : maybe need to transpose it
-- | <http://en.wikipedia.org/wiki/Rotation_matrix#Quaternion>
quaternionToMatrix :: Quaternion -> Mat33
quaternionToMatrix q@(Vec4 w x y z) = let
    nq = w^2+x^2+y^2+z^2
    s = if nq > 0.0 then 2/nq else 0
    (x', y', z') = (x*s,  y*s,  z*s)
    (wx, wy, wz) = (w*x', w*y', w*z')
    (xx, xy, xz) = (x*x', x*y', x*z')
    (yy, yz, zz) = (y*y', y*z', z*z')
    in Mat33
        (1.0-(yy+zz)) (xy-wz)       (xz+wy)
        (xy+wz)       (1.0-(xx+zz)) (yz-wx)
        (xz-wy)       (yz+wx)       (1.0-(xx+yy))
rotatePoint :: Quaternion -> Vec -> Vec
rotatePoint = mat33xVec . quaternionToMatrix
rotatePointAroundAxis :: Vec -- ^ axis
            -> FloatType -- ^ angle
            -> Vec -- ^ point
            -> Vec
rotatePointAroundAxis axis angle = rotatePoint $ rotationQ axis angle

-- | Dot product of arbitrary-dimension vectors represented by lists
dot :: [FloatType] -> [FloatType] -> FloatType
dot x y = sum $ zipWith (*) x y

-- | Solve system of linear equations by Gaussian elimination with partial
-- pivoting. Assumes non-degeneracy.
gaussSolve :: [[FloatType]] -> [FloatType] -> [FloatType]
gaussSolve a b = let
    (ta, tb) = transform 0 a b
    in utSolve ta tb
    where
    n = length a
    -- Recursively transform matrix to upper-triangular, starting from row k
    transform k a b = let
        -- pivot row is row (starting from k) with maximum absolute value of k-th element
        pRowIndex = flip maximumBy [k .. (n-1)] $ comparing $ \i -> abs $ a !! i !! k
        -- swap k-th and pivot rows
        a1 = swapInList k pRowIndex a
        b1 = swapInList k pRowIndex b
        -- pivot element
        pivot = a1 !! k !! k
        -- divide pivot row and free coef. by pivot
        a2 = updateInList (map (/pivot)) k a1
        b2 = updateInList (/pivot) k b1
        pRow = a2 !! k
        pb = b2 !! k
        -- transform lower rows
        ab3 = zipWith3
            (\i row b ->
                if i <= k
                    then (row, b)
                    else let kth = row !! k in (zipWith (-) row $ (kth*) `map` pRow, b-kth*pb))
            [0..] a2 b2
        a3 = map fst ab3
        b3 = map snd ab3
        in if k == n-1 then (a3, b3) else transform (k+1) a3 b3
-- | Solve upper-triangular system
utSolve :: [[FloatType]] -> [FloatType] -> [FloatType]
utSolve a b = solution where
    n = length a
    solution = [ (bi - drop (i+1) solution `dot` drop (i+1) vi)/(vi!!i)
        | vi <- a | bi <- b | i <- [0..n-1]]

-- | For vector a, gives matrix A such that a x b = A b. Matrix A is
-- skew-symmetric.
crossAsMatrix :: Vec -> Mat33
crossAsMatrix (Vec x y z) = Mat33
    0    (-z) y
    z    0    (-x)
    (-y) x    0