{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE NamedFieldPuns        #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE FlexibleInstances #-}
{-# OPTIONS_GHC -Wall #-}

module Walker where

import           Data.Maybe
import           Control.Monad.Reader
import           Language.OpenSCAD
import           Shared


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

data PackageConfig = PackageConfig {
        pkg_name :: String,
        scad_dir :: Filename,
        pkg_dir  :: Filename,
        src_dir  :: Filename
    };

type PackageMonad a = Reader PackageConfig a

instance Eq (Argument Expr) where
  (Argument lh          ) == (Argument rh          ) = lh == rh
  (NamedArgument lhk lhv) == (NamedArgument rhk rhv) = lhk == rhk && lhv == rhv
  _                       == _                       = False

instance Ord (Argument Expr) where
  compare a b = compare (show a) (show b)

instance Eq (Expr) where
  lh == rh = show lh == show rh

data SCADModel = SCADModel {
        module_name   :: String,
        args :: [Argument Expr],
        imports :: [Filename]
    } deriving (Show, Eq, Ord);


data ExtendedSCADModel = ExtendedSCADModel {
        smodel           :: SCADModel,
        scad_filename    :: Filename,
        stl_filename     :: Filename,
        urdf_filename    :: Filename,
        absolute_imports :: [Filename]
    } deriving (Show, Eq);

data Part m = Part {
            part_name        :: String,
            model            :: m,
            part_weight      :: Maybe Mass,
            center_of_mass   :: Maybe CenterOfMass
        } deriving (Show, Eq, Ord);

type SimplePart = Part SCADModel
type ExtendedPart = Part ExtendedSCADModel

data Assembly m = Assembly {
	    amodel :: m,
            collision_models :: [m]
	} deriving (Show);

type SimpleAssembly = Assembly SCADModel
type ExtendedAssembly = Assembly ExtendedSCADModel

data Link a = Link {
            link_name :: String,
	    weight :: Mass,
            assembly :: Maybe a
        } deriving (Show);

type SimpleLink = Link SimpleAssembly
data ExtendedLink = ExtendedLink {
        link      :: Link ExtendedAssembly,
        urdf_name :: String
    } deriving (Show);

data JointType = RevoluteJoint {
            axis          :: Axis,
            min_angle     :: Angle,
            default_angle :: Angle,
            max_angle     :: Angle,
            vel_lim       :: AngularVelocity,
            effor_lim     :: Effort,
            g_steps       :: Integer,
            actuator_name :: String
    } | FixedJoint {} deriving (Show);

data Joint l sub_j= Joint {
            joint_name :: String,
            offset     :: Pose,
            jtype      :: JointType,
            sub_link   :: l,
            sub_joint  :: Maybe sub_j
    } deriving (Show);

newtype SimpleJoint = SimpleJoint {
        j :: Joint SimpleLink SimpleJoint
    } deriving (Show);
data ExtendedJoint = ExtendedJoint {
            joint     :: Joint ExtendedLink ExtendedJoint,
            urdf_name :: String
    } deriving (Show);

data Leg j = Leg {
            leg_name   :: LegName,
            root_joint :: j
        } deriving (Show);

type SimpleLeg = Leg SimpleJoint
type ExtendedLeg = Leg ExtendedJoint

data Body a= Body {
	body_ass :: a,
        body_name :: String,
        body_weight :: Mass
    } deriving (Show);

type SimpleBody = Body SimpleAssembly
data ExtendedBody = ExtendedBody {
    base_body :: Body ExtendedAssembly,
    urdf_name :: String
} deriving (Show);

data Walker l b = Walker {
            walker_name :: String,
            legs        :: [l],
            body        :: b
    } deriving (Show);

type SimpleWalker = Walker SimpleLeg SimpleBody
type ExtendedWalker = Walker ExtendedLeg ExtendedBody


extendWalker :: SimpleWalker -> PackageMonad ExtendedWalker
extendWalker w = do
  elegs <- traverse extend_leg $ legs w
  ebody <- extend_body $ body w
  return $ w { legs = elegs, body = ebody }
 where
  extend_model :: String -> SCADModel -> PackageMonad ExtendedSCADModel
  extend_model prefix m = do
    scad  <- asks scad_dir
    pkg_d <- asks pkg_dir
    src   <- asks src_dir
    pkg   <- asks pkg_name
    let stl = "stl/" ++ prefix ++ "/" ++ module_name m ++ ".stl"
    return $ ExtendedSCADModel
      { smodel           = m
      , scad_filename    = scad ++ "/" ++ prefix ++ "/" ++ module_name m ++ ".scad"
      , stl_filename     = pkg_d ++ "/" ++ stl
      , urdf_filename    = "package://" ++ pkg ++ "/" ++ stl
      , absolute_imports = (\x -> src ++ "/" ++ x) <$> imports m
      }

  extend_part :: String -> SimplePart -> PackageMonad ExtendedPart
  extend_part prefix p = do
    model     <- extend_model prefix (model p)
    return $ p { model = model }

  extend_link :: SimpleLeg -> SimpleLink -> PackageMonad ExtendedLink
  extend_link leg l = do
    eass <- traverse (extend_assembly urdf_name) $ assembly l 
    return $ ExtendedLink { link = l { assembly = eass }, urdf_name = urdf_name }
    where urdf_name = "leg/" ++ leg_name leg ++ "/" ++ link_name l

  extend_joint :: SimpleLeg -> SimpleJoint -> PackageMonad ExtendedJoint
  extend_joint leg SimpleJoint { j } = do
    sub_l <- extend_link leg $ sub_link j
    sub_jnt <- case sub_joint j of
      (Just child) -> do
        echild <- extend_joint leg child
        return $ Just echild
      Nothing -> return Nothing
    return $ ExtendedJoint
      { joint     = j { sub_link = sub_l, sub_joint = sub_jnt }
      , urdf_name = "leg/" ++ leg_name leg ++ "/" ++ joint_name j
      }

  extend_leg :: SimpleLeg -> PackageMonad ExtendedLeg
  extend_leg l = do
    rjoint <- extend_joint l $ root_joint l
    return $ l { root_joint = rjoint }

  extend_body :: SimpleBody -> PackageMonad ExtendedBody
  extend_body b = do
    eass <- extend_assembly (body_name b) $ body_ass b
    return $ ExtendedBody { base_body = b { body_ass = eass}
                          , urdf_name = body_name b
                          }

  extend_assembly :: String -> SimpleAssembly -> PackageMonad ExtendedAssembly
  extend_assembly prefix (Assembly{amodel, collision_models}) = do
	  emodel <- extend_model prefix amodel
	  ecols <- traverse (extend_model prefix) collision_models
	  return $ Assembly { amodel = emodel, collision_models = ecols }


foldJoints :: (ExtendedJoint -> a -> a) -> a -> ExtendedJoint -> a
foldJoints f base ej = case sub_joint $ joint ej of
  Just child -> f ej $ foldJoints f base child
  Nothing    -> f ej base

mapJoints :: (ExtendedJoint -> a) -> ExtendedJoint -> [a]
mapJoints f = foldJoints (\w x -> f w : x) []

mapLinks :: (ExtendedLink -> a) -> ExtendedJoint -> [a]
mapLinks f rj = f <$> ( mapJoints (\x -> sub_link $ joint x) rj) 

mapAssemblies :: (ExtendedAssembly -> a) -> ExtendedWalker -> [a]
mapAssemblies f (Walker {legs, body}) = 
	(f $ body_ass $ base_body body ) : ( concatMap (\Leg {root_joint} -> catMaybes $ mapLinks ass_map root_joint) legs )
	where
		ass_map l = f <$> (assembly $ link l)	

mapModels :: forall a . (ExtendedSCADModel -> a) -> ExtendedWalker -> [a]
mapModels f w = concat $ mapAssemblies f_ass w
  where
    f_ass :: ExtendedAssembly -> [a]
    f_ass Assembly {amodel, collision_models} = ( f amodel ) : (f <$> collision_models)
    
