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

module URDF.Xml where

import           Data.List
import           Data.Maybe
import           Shared
import           URDF
import           System.Directory
import           Filesystem.Path
import           Filesystem.Path.CurrentOS
import           Linear.V3
import           Numeric.Units.Dimensional.Prelude hiding ( Angle
                                                          , Length
                                                          , Mass
                                                          )
import           Text.XML.HXT.Core                 hiding ( second )

mapMaybeToList :: (a -> b) -> Maybe a -> [b]
mapMaybeToList f m = map f $ maybeToList m

optionalAttrToList :: (ArrowXml a, Show b) => String -> Maybe b -> [a n XmlTree]
optionalAttrToList name maybe = map (sattr name . show) (maybeToList maybe)

cMassToXML :: ArrowXml a => Mass -> a XmlTree XmlTree
cMassToXML m = aelem "mass" [sattr "value" $ show $ m /~ (kilo gram)]

cPositionToAttr :: ArrowXml a => Position -> a n XmlTree
cPositionToAttr (V3 x y z) =
    sattr "xyz" (unwords $ map show [x /~ meter, y /~ meter, z /~ meter])

cAxisToXML :: ArrowXml a => Axis -> a XmlTree XmlTree
cAxisToXML ax = aelem "axis" [cPositionToAttr ax]

cOrientationToAttr :: ArrowXml a => Orientation -> a n XmlTree
cOrientationToAttr (Orientation r p y) =
    sattr "rpy" (unwords $ map show [r /~ radian, p /~ radian, y /~ radian])

cScaleToAttr :: ArrowXml a => Scale -> a n XmlTree
cScaleToAttr (V3 x y z) = sattr "scale" (unwords $ map show [x, y, z])

cColorToXML :: ArrowXml a => Color -> a XmlTree XmlTree
cColorToXML (Color r g b a) =
    aelem "color" [sattr "rgba" (unwords $ map show [r, g, b, a])]

cTextureToXML :: ArrowXml a => Texture -> a XmlTree XmlTree
cTextureToXML (Texture filename) = aelem "texture" [sattr "filename" filename]


cInertiaMatrixToXML :: ArrowXml a => InertiaMatrix -> a XmlTree XmlTree
cInertiaMatrixToXML m = aelem
    "inertia"
    [ sattr "ixx" $ conv (ixx m)
    , sattr "ixy" $ conv (ixy m)
    , sattr "ixz" $ conv (ixz m)
    , sattr "iyy" $ conv (iyy m)
    , sattr "iyz" $ conv (iyz m)
    , sattr "izz" $ conv (izz m)
    ]
  where
    conv :: Inertia -> String
    conv a = show $ a /~ (meter * meter * gram)

cGeometryToXML :: ArrowXml a => Geometry -> a XmlTree XmlTree
cGeometryToXML gem = mkelem "geometry" [] [inner gem]
  where
    inner (Box size) = aelem "box" [sattr "size" $ show $ size /~ meter]
    inner (Cylinder radius length) = aelem
        "cylinder"
        [ sattr "radius" $ show $ radius /~ meter
        , sattr "length" $ show $ length /~ meter
        ]
    inner (Sphere radius) =
        aelem "sphere" [sattr "radius" $ show $ radius /~ meter]
    inner (Mesh filename scale) = aelem
        "mesh"
        (sattr "filename" filename : mapMaybeToList cScaleToAttr scale)

cPrismaticDynamicsToXML :: ArrowXml a => PrismaticDynamics -> a XmlTree XmlTree
cPrismaticDynamicsToXML (PrismaticDynamics damping friction) = aelem
    "dynamics"
    (optionalAttrToList "damping" (fmap (/~ (newton * second / meter)) damping)
    ++ optionalAttrToList "friction" (fmap (/~ newton) friction)
    )

cRevoluteDynamicsToXML :: ArrowXml a => RevoluteDynamics -> a XmlTree XmlTree
cRevoluteDynamicsToXML (RevoluteDynamics damping friction) = aelem
    "dynamics"
    (  optionalAttrToList "damping" (fmap (/~ (newton * meter)) damping)
    ++ optionalAttrToList
           "damping"
           (fmap (/~ (newton * second / radian)) friction)
    )

cOriginToXML :: ArrowXml a => Origin -> a XmlTree XmlTree
cOriginToXML origin = aelem
    "origin"
    (  mapMaybeToList cPositionToAttr    (xyz origin)
    ++ mapMaybeToList cOrientationToAttr (rpy origin)
    )

cMaterialToXML :: ArrowXml a => Material -> a XmlTree XmlTree
cMaterialToXML (Material name color texture) = mkelem
    "material"
    [sattr "name" name]
    (mapMaybeToList cColorToXML color ++ mapMaybeToList cTextureToXML texture)

cInertialToXML :: ArrowXml a => Inertial -> a XmlTree XmlTree
cInertialToXML (Inertial mass inertia origin) = mkelem
    "inertial"
    []
    (  [cMassToXML mass, cInertiaMatrixToXML inertia]
    ++ mapMaybeToList cOriginToXML origin
    )

cVisualToXML :: ArrowXml a => Visual -> a XmlTree XmlTree
cVisualToXML (Visual v_geometry name origin material) = mkelem
    "visual"
    (optionalAttrToList "name" name)
    (  [cGeometryToXML v_geometry]
    ++ mapMaybeToList cOriginToXML   origin
    ++ mapMaybeToList cMaterialToXML material
    )

cCollisionToXML :: ArrowXml a => Collision -> a XmlTree XmlTree
cCollisionToXML (Collision geometry name origin) = mkelem
    "collision"
    (optionalAttrToList "name" name)
    (cGeometryToXML geometry : mapMaybeToList cOriginToXML origin)

cCalibrationToXML :: ArrowXml a => Calibration -> a XmlTree XmlTree
cCalibrationToXML calib = aelem
    "calibration"
    (  optionalAttrToList "rising"  (rising calib)
    ++ optionalAttrToList "falling" (falling calib)
    )

cLimitToXML :: ArrowXml a => Limit -> a XmlTree XmlTree
cLimitToXML (Limit effort velocity lower upper) = aelem
    "limit"
    (  [ sattr "effort"   (show effort)
       , sattr "velocity" (show $ velocity /~ (radian / second))
       ]
    ++ optionalAttrToList "lower" (fmap (/~ radian) lower)
    ++ optionalAttrToList "upper" (fmap (/~ radian) upper)
    )

cMimicToXML :: ArrowXml a => Mimic -> a XmlTree XmlTree
cMimicToXML (Mimic joint multiplier offset) = aelem
    "mimic"
    (  [sattr "joint" (show joint)]
    ++ optionalAttrToList "multiplier" multiplier
    ++ optionalAttrToList "offset"     offset
    )

cSafetyControllerToXML :: ArrowXml a => SafetyController -> a n XmlTree
cSafetyControllerToXML (SafetyController k_velocity k_position soft_lower_limit soft_upper_limit)
    = aelem
        "safety_controller"
        (  [sattr "k_velocity" $ show k_velocity]
        ++ optionalAttrToList "k_position" k_position
        ++ optionalAttrToList "soft_lower_limit"
                              (fmap (/~ radian) soft_lower_limit)
        ++ optionalAttrToList "soft_upper_limit"
                              (fmap (/~ radian) soft_upper_limit)
        )

cTypeToAttr :: ArrowXml a => JointType -> a n XmlTree
cTypeToAttr Revolute{}   = sattr "type" "revolute"
cTypeToAttr Continuous{} = sattr "type" "continuous"
cTypeToAttr Prismatic{}  = sattr "type" "prismatic"
cTypeToAttr Fixed        = sattr "type" "fixed"
cTypeToAttr Floating     = sattr "type" "floating"
cTypeToAttr Planar{}     = sattr "type" "planar"

cTypeToExtra :: ArrowXml a => JointType -> [a XmlTree XmlTree]
cTypeToExtra Revolute { axis, limit, rev_dynamics } =
    cLimitToXML limit
        :  mapMaybeToList cRevoluteDynamicsToXML rev_dynamics
        ++ mapMaybeToList cAxisToXML             axis
cTypeToExtra Continuous { axis, rev_dynamics } =
    mapMaybeToList cRevoluteDynamicsToXML rev_dynamics
        ++ mapMaybeToList cAxisToXML axis
cTypeToExtra Prismatic { axis, limit, pri_dynamics } =
    cLimitToXML limit
        :  mapMaybeToList cPrismaticDynamicsToXML pri_dynamics
        ++ mapMaybeToList cAxisToXML              axis
cTypeToExtra _ = []

cJointToXML :: ArrowXml a => Joint -> a XmlTree XmlTree
cJointToXML Joint { name, jtype, parent, child, origin, calibration, mimic, safety_controller }
    = mkelem
        "joint"
        [sattr "name" name, cTypeToAttr jtype]
        (  [ aelem "parent" [sattr "link" parent]
           , aelem "child"  [sattr "link" child]
           ]
        ++ mapMaybeToList cOriginToXML           origin
        ++ mapMaybeToList cCalibrationToXML      calibration
        ++ mapMaybeToList cMimicToXML            mimic
        ++ mapMaybeToList cSafetyControllerToXML safety_controller
        ++ cTypeToExtra jtype
        )

cLinkToXML :: ArrowXml a => Link -> a XmlTree XmlTree
cLinkToXML (Link name visuals inertial collision) = mkelem
    "link"
    [sattr "name" name]
    (  map cVisualToXML visuals
    ++ mapMaybeToList cInertialToXML inertial
    ++ map cCollisionToXML collision
    )

cTransmissionToXML :: ArrowXml a => Transmission -> a XmlTree XmlTree
cTransmissionToXML (Transmission name ttype joint interface actuator) = mkelem
    "transmission"
    [sattr "name" name ]
    [ selem "type" [ txt ttype ],
      mkelem "joint" [ sattr "name" joint ] [ selem "hardwareInterface" [ txt interface ] ],
      aelem "actuator" [ sattr "name" actuator ]
    ]
    

cURDFToXML :: ArrowXml a => URDF -> a XmlTree XmlTree
cURDFToXML (URDF name joints links trans) = mkelem
    "robot"
    [sattr "name" name]
    ((cJointToXML <$> joints) ++ (cLinkToXML <$> links) ++ (cTransmissionToXML <$> trans)) 

writeURDFXML :: URDF -> Filename -> IO ()
writeURDFXML urdf filename = do
    createDirectoryIfMissing True $ encodeString $ directory $ decodeString
        filename
    runX
        $   Text.XML.HXT.Core.root [] [cURDFToXML urdf]
        >>> writeDocument [withIndent yes] filename
    return ()
