{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE MultiWayIf #-}
module Physics.Constraint where

import Numeric.LinearAlgebra hiding (step)
import Data.Vec hiding (map,foldr,length,scale)
import qualified Data.Vec as V
import qualified Data.IntMap as IntMap
import           Data.IntMap (IntMap)

import Data.Vec.Quaternion
import Physics.Object
import Debug.Trace

-- Implement two methods:
-- 1. spring penalty
-- 2. lagrange multipliers

type ConstraintSolver = [Constraint] -> IntMap Object -> IntMap Object

data Constraint
    = FixedPoint ObjectId (Vec3 Double) (Vec3 Double)
    -- ^ @FixedPoint a s t@ fixes the point @s@ in the object-centric coordinates of @a@ to the position @t@ in space. Rotation is unconstrained. Sometimes known as a nail constraint.
    -- | FixedAngle ObjectId Quat
    -- ^ @FixedAngle a q@ fixes the angular position of @a@ to be @q@.
    | BallSocket ObjectId (Vec3 Double) ObjectId (Vec3 Double)
    -- ^ @BallSocket a s b t@ constraints the points @s@ and @t@, in the coordinates of @a@ and @b@ respectively, to coincide in space. Rotation is unconstrained.
    -- | Weld ObjectId (Vec3 Double) ObjectId (Vec3 Double)
    -- ^ Welds together two objects at a point specified in each object's coordinates. This prevents both relative motion and relative rotation.
    -- TODO: universal joint (1 restricted), hinge joint (2 restricted), weld (3 restricted)
    -- TODO: confinement to a plane (for vertex/face collision)
    -- TODO: edge/edge collision
    deriving Show

{-
springPenalty :: ConstraintSolver
springPenalty cs os = foldr (flip (foldr applyF)) (foldr (flip (foldr applyT)) os (map torques cs)) (map forces cs)
    where applyF (i,f,r) = IntMap.adjust (oModifyBody (addForceAt r f)) i
          applyT (i,t) = IntMap.adjust (oModifyBody  (addTorque t)) i
          forces (FixedPoint a s t) = [(a, V.map (*k a) (lspring 0 (bodyToSpace (body a) s) t), s)]
          forces (FixedAngle a q) = []
          forces (BallSocket a s b t) = error "TODO: springPenalty BallSocket"
          torques (FixedPoint _ _ _) = []
          torques (FixedAngle a q) = [(a, V.map (*kappa a) (aspring q (bAngularPosition (body a))))]
          torques (BallSocket a s b t) = []
          k i = kHeuristic (body i)
          kappa i = kappaHeuristic (body i)
          body i = oBody (os IntMap.! i)
          pos i = bPosition (body i)
          lspring :: Double -> Vec3 Double -> Vec3 Double -> Vec3 Double
          lspring resting p1 p2 = p2 - p1 - V.map (*resting) (normalize' (p2-p1))
          aspring :: Quat -> Quat -> Vec3 Double
          -- FIXME: a might mean the opposite of what i think now
          aspring resting a = let (t,n) = toAxisAngle (a * qconj resting)
                              in vec (-t) * n
          -- Set spring constant based on the mass of the object and current linear drag
          kHeuristic b = 3.0 * bMass b * (1.0 + bDrag b)
          -- TODO: find a heuristic for kappa
          kappaHeuristic = const 2.0
-}

lagrangian :: ConstraintSolver
lagrangian constraints objects = let
    -- parameters of constraint-restoring feedback force
    ks = 80
    kd = 2
    -- state vectors
    n = IntMap.size objects
    body i = oBody (objects IntMap.! i)
    mass i = bMass (body i)
    pos i = bPosition (body i)
    vel i = bVelocity (body i)
    force i = bForce (body i)
    -- TODO: handle zero-mass bodies
    accel i = force i / vec (mass i)
    amass i = unpackMat (bAngularMass (body i))
    amass_inertials = [multmm (rotm i) (multmm (amass i) (transpose (rotm i))) | i <- [0..n-1]]
    amass_inertial i = amass_inertials !! i
    amass_inertial_inverses = map invert3x3 amass_inertials
    amass_inertial_inverse i = amass_inertial_inverses !! i
    apos i = bAngularPosition (body i)
    rotms = [toRotationMatrix (apos i) | i <- [0..n-1]]
    rotm i = rotms !! i
    -- converts from body-local rotation frame to world frame
    wframe bi r = qrotate (apos bi) r
    avel i = bAngularVelocity (body i)
    aforce i = bAngularForce (body i)
    aaccels = [cramer'sRule (amass_inertial i) (aforce i + (phips!!i)) | i <- [0..n-1]]
    aaccel i = aaccels !! i
    -- (vel (body 0), avel (body 0), ...)
    xdot = buildVector (6*n) $ \i -> let l = i `div` 6 in
        case i `mod` 6 of
             0 -> get n0 (vel l)
             1 -> get n1 (vel l)
             2 -> get n2 (vel l)
             3 -> get n0 (avel l)
             4 -> get n1 (avel l)
             5 -> get n2 (avel l)
             _ -> error "impossible"
    xdotdot = buildVector (6*n) $ \i -> let l = i `div` 6 in
        case i `mod` 6 of
             0 -> get n0 (accel l)
             1 -> get n1 (accel l)
             2 -> get n2 (accel l)
             3 -> get n0 (aaccel l)
             4 -> get n1 (aaccel l)
             5 -> get n2 (aaccel l)
             _ -> error "impossible"
    -- global force vector, not considering precession
    phi = buildVector (6*n) $ \i -> let l = i `div` 6 in
        case i `mod` 6 of
             0 -> get n0 (force l)
             1 -> get n1 (force l)
             2 -> get n2 (force l)
             3 -> get n0 (aforce l)
             4 -> get n1 (aforce l)
             5 -> get n2 (aforce l)
             _ -> error "impossible"
    -- phip is the effects of free precession
    phips = [ - cross (avel i) (multmv (amass_inertial i) (avel i)) | i <- [0..n-1]]
    phip = buildVector (6*n) $ \i -> let l = i `div` 6 in
        case i `mod` 6 of
             0 -> 0
             1 -> 0
             2 -> 0
             3 -> get n0 (phips !! l)
             4 -> get n1 (phips !! l)
             5 -> get n2 (phips !! l)
             _ -> error "impossible"
    imass = buildMatrix (6*n) (6*n) $ \(i,j) ->
        if i `div` 6 /= j `div` 6 then 0 else
        let l = i `div` 6
            ix = i `mod` 6
            jx = j `mod` 6 in
        if | ix == jx && ix < 3 -> recip (mass l)
           | ix >= 3 && jx >= 3 -> getElem (jx-3) (getElem (ix-3) (amass_inertial_inverse l))
           | otherwise -> 0
    -- =================
    c_slice (FixedPoint bi (wframe bi->s) t) = toVector $ pos bi + s - t
    c_slice (BallSocket bi (wframe bi->s) bj (wframe bj->t)) = toVector $ pos bi + s - pos bj - t
    cdot_slice (FixedPoint bi (wframe bi->s) _) = toVector $ vel bi + cross (avel bi) s
    cdot_slice (BallSocket bi (wframe bi->s) bj (wframe bj->t)) = toVector $ vel bi + cross (avel bi) s - vel bj - cross (avel bj) t
    j_slice (FixedPoint bi (wframe bi->(s1:.s2:.s3:.())) _) =
        buildMatrix 3 (6*n) $ \(i,j) ->
            if j `div` 6 /= bi then 0 else
            [[1,0,0,0,s3,-s2]
            ,[0,1,0,-s3,0,s1]
            ,[0,0,1,s2,-s1,0]] !! i !! mod j 6
    j_slice (BallSocket bi (wframe bi->(s1:.s2:.s3:.())) bj (wframe bj->(t1:.t2:.t3:.()))) =
        buildMatrix 3 (6*n) $ \(i,j) ->
            if | j `div` 6 == bi ->
                [[1,0,0,0,s3,-s2]
                ,[0,1,0,-s3,0,s1]
                ,[0,0,1,s2,-s1,0]] !! i !! mod j 6
               | j `div` 6 == bj -> negate $
                [[1,0,0,0,t3,-t2]
                ,[0,1,0,-t3,0,t1]
                ,[0,0,1,t2,-t1,0]] !! i !! mod j 6
               | otherwise -> 0
    jdot_slice (FixedPoint bi (wframe bi->(s1:.s2:.s3:.())) _) =
        let (w1:.w2:.w3:.()) = avel bi in
        buildMatrix 3 (6*n) $ \(i,j) ->
            if j `div` 6 /= bi then 0 else
            [[0,0,0,0,          w1*s2-w2*s1,w1*s3-w3*s1]
            ,[0,0,0,w2*s1-w1*s2,0,          w2*s3-w3*s2]
            ,[0,0,0,w3*s1-w1*s3,w3*s2-w2*s3,0]] !! i !! mod j 6
    jdot_slice (BallSocket bi (wframe bi->(s1:.s2:.s3:.())) bj (wframe bj->(t1:.t2:.t3:.()))) =
        let (v1:.v2:.v3:.()) = avel bi
            (w1:.w2:.w3:.()) = avel bj in
        buildMatrix 3 (6*n) $ \(i,j) ->
            if | j `div` 6 == bi ->
                [[0,0,0,0,          v1*s2-v2*s1,v1*s3-v3*s1]
                ,[0,0,0,v2*s1-v1*s2,0,          v2*s3-v3*s2]
                ,[0,0,0,v3*s1-v1*s3,v3*s2-v2*s3,0]] !! i !! mod j 6
               | j `div` 6 == bj -> negate $
                [[0,0,0,0,          w1*s2-w2*s1,w1*s3-w3*s1]
                ,[0,0,0,w2*s1-w1*s2,0,          w2*s3-w3*s2]
                ,[0,0,0,w3*s1-w1*s3,w3*s2-w2*s3,0]] !! i !! mod j 6
               | otherwise -> 0
    -- =================
    c = join (map c_slice constraints)
    cdot = join (map cdot_slice constraints)
    j = fromBlocks (map (\c -> [j_slice c]) constraints)
    jdot = fromBlocks (map (\c -> [jdot_slice c]) constraints)
    -- lambda are the lagrangian multipliers
    lambda = (- j <> imass <> trans j) <\> (jdot <> xdot + j <> imass <> (phi + phip) + scale ks c + scale kd cdot)
    phic = trans j <> lambda
    extract3 k = (phic@>k):.(phic@>(k+1)):.(phic@>(k+2)):.()
    update i b = b {
        bForce = bForce b + extract3 (6*i),
        bAngularForce = bAngularForce b + extract3 (6*i+3)
        }
    in IntMap.mapWithKey (\i -> oModifyBody (update i)) objects

toVector (x:.y:.z:.()) = 3|>[x,y,z]

-- Two-dimensional constraints for reference
{-

jacobian :: ConstraintSolver
jacobian cs os = let
    -- state vectors
    n = Seq.length os
    pos i = oPosition (Seq.index os i)
    vel i = oVelocity (Seq.index os i)
    stateVec f = buildVector (2*n) (\k ->
        let Vector2 x y = f $ Seq.index os (div k 2)
        in if even k then x else y)
    velocities = stateVec oVelocity
    forces = stateVec oForce
    -- Mass matrix and its inverse
    -- TODO: make it a vector instead
    m = diag (stateVec (\o -> Vector2 (oMass o) (oMass o)))
    w = diag (stateVec (\o -> 1 /| Vector2 (oMass o) (oMass o)))
    -- Jacobian. dimension (constraints)x(2*objects)
    mkJ f = buildMatrix (length cs) (2*n) $ \(ic,io) ->
        case cs !! ic of
             FixedDistance r i j ->
                 let Vector2 x y =
                        if i == div io 2
                           then f (FixedDistance r i j)
                        else if j == div io 2
                           then f (FixedDistance r j i)
                        else 0
                 in if even io then x else y
             Fixed i p ->
                 let Vector2 x y =
                        if i == div io 2
                           then f (Fixed i p)
                        else 0
                 in if even io then x else y
    jay = mkJ c_dq
    jaydot = mkJ c_dqdt
    -- Constraints!
    c (FixedDistance d i j) = vdot (pos i - pos j) (pos i - pos j) - d*d
    c (Fixed i p) = vdot (pos i - p) (pos i - p)
    c_dt (FixedDistance d i j) = 2 * vdot (vel i - vel j) (pos i - pos j)
    c_dt (Fixed i p) = 2 * vdot (vel i) (pos i - p)
    c_dq (FixedDistance d i j) = 2 *| (pos i - pos j)
    c_dq (Fixed i p) = 2 *| (pos i - p)
    c_dqdt (FixedDistance _ i j) = 0
    c_dqdt (Fixed i p) = 0
    constraintVec f = fromList (map f cs)
    -- feedback
    ks = 20
    kd = 5
    -- result
    lambda = (jay <> w <> ctrans jay) <\> (- jaydot <> asColumn velocities - jay <> w <> asColumn forces - asColumn (scale ks (constraintVec c) + scale kd (constraintVec c_dt)))
    constraintForces = head $ toColumns $ ctrans jay <> lambda
    apply o fv = addForce (Vector2 (fv@>0) (fv@>1)) o
    in
    -- apply constraint forces
    zipWithTF apply os (takesV (replicate n 2) constraintForces)

-}

dualTensor :: Num a => Vec3 a -> Mat33 a
dualTensor (a1:.a2:.a3:.()) = ( (   0 :.(-a3):.  a2 :.())
                              :.(  a3 :.   0 :.(-a1):.())
                              :.((-a2):.  a1 :.   0 :.()):.())

invert3x3 :: (Num a, Fractional a) => Mat33 a -> Mat33 a
invert3x3 a@(transpose->(x0:.x1:.x2:.())) = V.map (* vec (recip (V.det a))) (cross x1 x2 :. cross x2 x0 :. cross x0 x1 :. ())

