{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE NamedFieldPuns        #-}

module Walker.URDF where

import           Data.Hashable                            ( hash )
import           Linear.V3
import qualified Numeric.Units.Dimensional.Prelude
                                               as U
import           Shared
import qualified URDF
import           Walker
import           Data.Maybe


poseToURDFOrigin :: Pose -> URDF.Origin
poseToURDFOrigin pose = URDF.Origin (Just $ pos pose) (Just $ orientation rpy)
  where
    rpy = (U.*~ U.radian) <$> quatToRPY (ori pose)
    orientation (V3 r p y) = URDF.Orientation r p y

getURDFLegJoints :: ExtendedBody -> ExtendedLeg -> [URDF.Joint]
getURDFLegJoints body leg = link_to_joint <$> triplets
  where
    joints :: [ExtendedJoint]
    joints = mapJoints id $ root_joint leg

    link_names :: [String]
    link_names = (urdf_name (body :: ExtendedBody)) : 
	(mapLinks (\x -> urdf_name (x :: ExtendedLink)) $ root_joint leg)

    triplets = zip3 link_names joints (tail link_names)

    link_to_joint :: (String, ExtendedJoint, String) -> URDF.Joint
    link_to_joint (p_link_name, ExtendedJoint { urdf_name, joint }, ch_name) =
        URDF.Joint
            { name              = urdf_name
            , jtype             = case jtype joint of
                RevoluteJoint { axis, min_angle, max_angle, vel_lim, effor_lim }
                    -> URDF.Revolute
                        { limit        = URDF.Limit effor_lim
                                                    vel_lim
                                                    (Just min_angle)
                                                    (Just max_angle)
                        , axis         = Just axis
                        , rev_dynamics = Nothing
                        }
                FixedJoint -> URDF.Fixed{}
            , parent            = p_link_name
            , child             = ch_name
            , origin            = Just $ poseToURDFOrigin $ offset joint
            , calibration       = Nothing
            , mimic             = Nothing
            , safety_controller = Nothing
            }

getURDFLegTransmissions :: ExtendedLeg -> [URDF.Transmission]
getURDFLegTransmissions leg = catMaybes $ mapJoints joint_to_trans $ root_joint leg
  where
    joint_to_trans :: ExtendedJoint -> Maybe URDF.Transmission
    joint_to_trans (ExtendedJoint j urdf_name) = case jtype j of
        RevoluteJoint {actuator_name} -> Just $ URDF.Transmission actuator_name "transmission_interface/SimpleTransmission" urdf_name "EffortJointInterface" actuator_name
        FixedJoint -> Nothing

getURDFSCADModelVisual :: String -> ExtendedSCADModel -> URDF.Visual
getURDFSCADModelVisual link_name ExtendedSCADModel {urdf_filename} = URDF.Visual
        { v_geometry = URDF.Mesh urdf_filename (Just $ V3 0.001 0.001 0.001)
        , name       = Nothing
        , origin     = Nothing
        , material   = Just $ URDF.Material link_name (Just color) Nothing
        }
  where
    color = URDF.Color { r = hash_color (link_name ++ "red")
                       , g = hash_color (link_name ++ "green")
                       , b = hash_color (link_name ++ "blue")
                       , a = 1
                       }

    hash_color :: String -> ColorChannel
    hash_color seed = fromIntegral (hash seed `mod` 1024) / 1024.0

getURDFAssemblyVisuals :: String -> ExtendedAssembly -> [URDF.Visual]
getURDFAssemblyVisuals link_name Assembly {amodel} = [getURDFSCADModelVisual link_name amodel]

getURDFAssemblyCollisions :: ExtendedAssembly -> [URDF.Collision]
getURDFAssemblyCollisions Assembly { collision_models } = f <$> collision_models
  where
    f ExtendedSCADModel { urdf_filename } = URDF.Collision
        { geometry = URDF.Mesh urdf_filename (Just $ V3 0.001 0.001 0.001)
        , name     = Nothing
        , origin   = Nothing
        }

getInertiaMatrixGuess :: Mass -> URDF.InertiaMatrix
getInertiaMatrixGuess m = URDF.InertiaMatrix v z z z v z
  where
    z :: Inertia
    z = 0.0 U.*~ (U.gram U.* U.metre U.* U.metre)

    r :: Radius
    r = 0.25 U.*~ U.metre 

    v :: Inertia
    v = ((2.0 / 3.0) U.*~ U.one) U.* (m U.* r U.* r )

getURDFLegLinks :: ExtendedLeg -> [URDF.Link]
getURDFLegLinks leg = mapLinks link_to_urdf $ root_joint leg
  where
    link_to_urdf ExtendedLink { urdf_name, link } = URDF.Link
        { name      = urdf_name
        , visuals   = case assembly link of
			(Just ass) -> getURDFAssemblyVisuals urdf_name ass
			Nothing -> []
        , inertial  = Just $ URDF.Inertial (weight link) (getInertiaMatrixGuess $ weight link) Nothing
        , collision = case assembly link of 
			(Just ass) -> getURDFAssemblyCollisions ass
			Nothing -> []
        }

getURDFBodyLink :: ExtendedBody -> URDF.Link
getURDFBodyLink ExtendedBody { base_body, urdf_name } = URDF.Link
    { name      = urdf_name
    , visuals   = getURDFAssemblyVisuals urdf_name $ body_ass base_body
    , inertial  = Just $ URDF.Inertial (body_weight base_body) (getInertiaMatrixGuess $ body_weight base_body) Nothing
    , collision = getURDFAssemblyCollisions $ body_ass base_body
    }

getURDF :: ExtendedWalker -> URDF.URDF
getURDF walker = URDF.URDF
    { name   = walker_name walker
    , joints = concatMap (getURDFLegJoints $ body walker) $ legs walker
    , links  = getURDFBodyLink (body walker)
                   : (concatMap getURDFLegLinks $ legs walker)
    , transmissions = concatMap getURDFLegTransmissions $ legs walker
    }
