{-# LANGUAGE DuplicateRecordFields #-}

module Data.Walker where

import           Data.Hashable     (hash)
import           Data.List
import           Data.SCAD
import           Data.Shared
import qualified Data.URDF         as URDF
import qualified Data.SRDF as SRDF
import           Linear.V3
import           Linear.Quaternion

type LegName = String
type CenterOfMass = Position
type PosOffset = Position

jointNames = ["alfa","beta","gama","delta","epsilon","zeta","eta","theta"]

data SCADModel = SCADModel {
        module_name :: SCADName,
        args        :: [Arg],
        imports     :: [Filename]
    }

scadModel :: SCADArg a => SCADName -> [(String, a)] -> [Filename] -> SCADModel
scadModel name args imps = SCADModel{
            module_name = name,
            args = map (uncurry toArg) args,
            imports = imps
        }


data Part = Part {
            model            :: SCADModel,
            collision_models :: [SCADModel],
            part_weight      :: Maybe Mass,
            center_of_mass   :: Maybe CenterOfMass
        };

data JointConfig = JointConfig {
            min_angle     :: Angle,
            default_angle :: Angle,
            max_angle     :: Angle,
            vel_lim       :: AngularVelocity,
            effor_lim     :: Effort,
            g_steps       :: Integer

}

data Link = Link {
            link_name     :: String,
            offset        :: PosOffset,
            part          :: Part,
            axis          :: Axis,
            joint_config ::JointConfig
        };
newtype Tip = Tip {
            point :: Position
        }

data Leg = Leg {
            leg_name :: LegName,
            chain    :: [Link],
            tip      :: Tip,
            pose     :: Pose
        };

data Body = Body {
        body_part :: Part
    };

data Walker = Walker {
            walker_name :: String,
            legs        :: [Leg],
            body        :: Body
    };


getLinkName leg name = "leg/" ++ leg_name leg ++ "/" ++ name

getJointName = getLinkName

poseToURDFOrigin :: Pose -> URDF.Origin 
poseToURDFOrigin pose = URDF.Origin (Just $ pos pose) ( Just $ quat_to_rpy $ ori pose )
    where

getLegJoints :: Leg -> [URDF.Joint]
getLegJoints leg = [ first_joint ] ++ map linkToURDFJoint link_triples ++ [ last_joint ]
    where
        link_names = map (getLinkName leg . link_name) $ chain leg
        link_triples = take ( length $ chain leg ) $ zip3 ( getLinkName leg "base" : link_names ) ( chain leg) (tail jointNames)

        jConfigToURDFLimit (JointConfig min_angle _ max_angle vel_lim effor_lim _) = URDF.Limit effor_lim vel_lim  (Just min_angle) (Just max_angle)

        first_joint = URDF.Joint {
                name = getJointName leg ( head jointNames ),
                jtype = URDF.Fixed,
                parent = "body",
                child = getLinkName leg "base",
                origin = Just $ URDF.Origin (Just $ pos $ pose leg) Nothing,
                axis = Nothing,
                calibration = Nothing,
                mimic = Nothing,
                safety_controller = Nothing
            }

        linkToURDFJoint :: (String, Link, String) -> URDF.Joint
        linkToURDFJoint (p_link_name, ch_link, name) = URDF.Joint {
                                name = getJointName leg name,
                                jtype = URDF.Revolute ( jConfigToURDFLimit $ joint_config ch_link ) Nothing,
                                parent = p_link_name,
                                child = getLinkName leg $ link_name ch_link,
                                origin = Just $ URDF.Origin ( Just $ offset ch_link ) Nothing,
                                axis = Just $ axis ch_link,
                                calibration = Nothing,
                                mimic = Nothing  ,
                                safety_controller = Nothing
                                }
        last_joint = URDF.Joint {
                name = getJointName leg ( jointNames !! (2 + length (chain leg)) ),
                jtype = URDF.Fixed,
                parent = last link_names,
                child = getLinkName leg "tip",
                origin = Just $ URDF.Origin ( Just $ point $ tip leg ) Nothing,
                axis = Nothing,
                calibration = Nothing,
                mimic = Nothing,
                safety_controller = Nothing
            }

getStlName leg name = "src/stl/" ++ getLinkName leg name ++ ".stl"

getURDFStlName leg name = "package://schpin_koke/" ++ getStlName leg name

getLinkVisual :: Leg -> Link -> URDF.Visual
getLinkVisual leg link = URDF.Visual {
                        v_geometry = URDF.Mesh ( getURDFStlName leg $ link_name link ) (Just $ V3 0.001 0.001 0.001),
                        name = Nothing,
                        origin = Nothing,
                        material = Just $ URDF.Material name (Just color ) Nothing
                    }
            where
                name = getLinkName leg $ link_name link
                color = URDF.Color {
                        r = fromIntegral ( hash (name ++ "red") `mod` 1024) / 1024.0,
                        g = fromIntegral ( hash (name ++ "green") `mod` 1024) / 1024.0,
                        b = fromIntegral ( hash (name ++ "blue") `mod` 1024) / 1024.0,
                        a = 1
                    }

getLinkCollision :: Leg -> Link -> [URDF.Collision]
getLinkCollision leg link = map f (collision_models $ part link)
            where
                f model =  URDF.Collision {
                        geometry = URDF.Mesh ( getURDFStlName leg ( module_name model ) ) (Just $ V3 0.001 0.001 0.001),
                        name = Nothing,
                        origin = Nothing
                    }

getLegLinks :: Leg -> [URDF.Link]
getLegLinks leg = [URDF.Link (getLinkName leg "base") [] Nothing []] ++ map linkToURDFLink ( chain leg) ++ [URDF.Link (getLinkName leg "tip") [] Nothing []]
    where
        linkToURDFLink leg_link = URDF.Link {
                                name = getLinkName leg $ link_name leg_link,
                                visuals = [ getLinkVisual leg leg_link  ],
                                inertial = Nothing,
                                collision = getLinkCollision leg leg_link
                                }


getBodyLink :: Body -> URDF.Link
getBodyLink body =  URDF.Link {
                        name = "body",
                        visuals =  [ body_vis ],
                        inertial =  Nothing,
                        collision = body_cols
                    }
    where
        body_cols = map f (collision_models $ body_part body )
        f model = URDF.Collision {
                    geometry = URDF.Mesh ("package://schpin_koke/src/stl/" ++ module_name model ++ ".stl") (Just $ V3 0.001 0.001 0.001),
                    name = Nothing,
                    origin = Nothing
            }
        body_vis = URDF.Visual {
                    v_geometry = URDF.Mesh  "package://schpin_koke/src/stl/body.stl" (Just $ V3 0.001 0.001 0.001),
                    name = Nothing,
                    origin = Nothing,
                    material = Nothing
                }

getURDF :: Walker -> URDF.URDF
getURDF walker = URDF.URDF (walker_name walker) joints links
    where
        joints = concatMap getLegJoints ( legs walker )
        links = getBodyLink ( body walker) : (concatMap getLegLinks $ legs walker)


getLegSRDFChain :: Leg -> SRDF.Chain
getLegSRDFChain leg = SRDF.Chain (link_name $ head urdf_chain) (link_name $ last urdf_chain )
    where
        urdf_chain = getLegLinks leg
        link_name :: URDF.Link -> UrdfName
        link_name (URDF.Link name _ _ _) = name

getLegSRDFJoints :: Leg -> [SRDF.Joint]
getLegSRDFJoints leg = zipWith convert_chain (chain leg) ( tail jointNames )
    where
        convert_chain (Link _ _ _ _ joint_config) joint_name = SRDF.Joint (getJointName leg joint_name) (build_joint_config joint_config)
        build_joint_config (JointConfig _ _ _ _ _ g_steps ) = [SRDF.GenericNode "graph_steps" [SRDF.Attribute "value" "16"] []]

getSRDF :: Walker -> SRDF.SRDF
getSRDF walker = SRDF.SRDF (walker_name walker) [leg_group] items [schpin_config]
    where
        leg_group = SRDF.Group "legs" [] (concatMap getLegSRDFJoints $ legs walker) (map getLegSRDFChain $ legs walker) [] [] 
        items = []
        schpin_config = SRDF.GenericNode "schpin_config" [] [
                            SRDF.GenericNode "environment" [SRDF.Attribute "resolution" "0.05"] [],
                            SRDF.GenericNode "leg_workarea" [SRDF.Attribute "resolution" "0.01"] [],
                            SRDF.GenericNode "leg_mapping" [SRDF.Attribute "resolution" "0.01"] []
                         ] 

getScadName leg name = "src/scad/gen/" ++ getLinkName leg name ++ ".scad"

getSTLRenderList :: Walker -> [(Filename, Filename)]
getSTLRenderList walker = ("src/scad/gen/body.scad", "src/stl/body.stl") : (concatMap convert ( legs walker ) ++ body_collisions)
        where
            body_collisions :: [(Filename, Filename)]
            body_collisions = map (\x->("src/scad/gen/" ++ module_name x ++ ".scad", "src/stl/" ++ module_name x ++ ".stl" ))  $ collision_models $ body_part $ body walker
            convert :: Leg -> [(Filename, Filename)]
            convert leg = zip ( map (getScadName leg) names ) (map (getStlName leg) names )
                where
                    link_names = map link_name $ chain leg
                    collision_names = concatMap ( \x -> map module_name $ collision_models $ part x ) $ chain leg
                    names = link_names ++ collision_names

getSCADModels :: Walker -> [(Filename, SCADModel)]
getSCADModels walkie = ("src/scad/gen/body.scad", model $ body_part $ body walkie) : (concatMap get_leg_models ( legs walkie ) ++ (body_collisions $ body walkie))
    where
        body_collisions :: Body -> [(Filename, SCADModel)]
        body_collisions body = map (\x -> ("src/scad/gen/" ++ (module_name x) ++ ".scad", x)) (collision_models $ body_part body)
        get_leg_models leg = concatMap (get_link_model leg) $ chain leg 
        get_link_model :: Leg -> Link -> [(Filename, SCADModel)]
        get_link_model leg (Link link_name _ part _ _) = (getScadName leg link_name, model part) : (map (\x -> (getScadName leg $ module_name x, x)) $ collision_models part) 

modelToSCAD :: (Filename, SCADModel) -> (Filename, [SCAD])
modelToSCAD (file, model) = (file, s)
    where
        s = map to_import_f ( imports model) ++ [ Module (module_name model) (args model) ]
        to_import_f :: Filename -> SCAD
        to_import_f filename = Use $ concat ( replicate (  length (filter (== '/') file) - 2 ) "../" ) ++ filename

getSCADSkeleton :: String -> Walker -> ( Part -> SCAD ) -> (Filename, [SCAD])
getSCADSkeleton filename (Walker _ legs body) part_to_skeleton = (filename, [Use "../body.scad"] ++ map Use ( concatMap get_imports legs) ++ map to_skeleton legs ++ [part_to_skeleton $ body_part body])
    where
        get_imports (Leg _ chain _ _) = map (\x -> "../" ++ x) $ concatMap (imports . model . part ) chain

        to_skeleton :: Leg -> SCAD
        to_skeleton (Leg _ chain tip pose) = Translate (pos pose) [ Rotate (ori pose) ( link_to_skeleton chain ) ]

        link_to_skeleton :: [Link] -> [SCAD]
        link_to_skeleton [] = []
        link_to_skeleton (link:rest) = [ Translate (offset link) [ Rotate (toQuat ( axis link) (default_angle $ joint_config link) ) ( part_to_skeleton ( part link ) :   link_to_skeleton rest)]]

partToVisual (Part model _ _ _) = Module (module_name model) (args model)
partToCollision (Part _ collision_models  _ _) = Debug $ map ( \x -> Module(module_name x) (args x) ) collision_models


getSCAD :: Walker -> [(Filename, [SCAD])]
getSCAD walkie = map modelToSCAD (getSCADModels walkie) ++ [getSCADSkeleton "src/scad/gen/skeleton.scad" walkie partToVisual] ++ [getSCADSkeleton "src/scad/gen/col_skeleton.scad" walkie partToCollision]
