from SAC.RL_Framework_Mujoco import Muscle_Env    #Muscle_Env is the musculoskeletal environment

from simulate import Simulate

import config

import pdb

def main():

    """ Train an agent to control a musculoskeletal system using DRL and RNNs

        ---------------------------------------------------------------------

    * Simulate object runs training

    * Specify whether testing or not in config file, along with the preferred file to save testing statistics

    * To ensure model is saved properly, specify your preferred directories for storing state_dict

    * A list of possible commands with their functions is given in the config.py file as well as the simulate.py file

    * Specify kinematics by choosing kinematics_path in config file, monkey data is currently provided by default

    """

    ### PARAMETERS ###

    parser = config.config_parser()

    args = parser.parse_args()

    ### TRAINING OBJECT ###

    trainer = Simulate(Muscle_Env, args)

    # trainer = Simulate(

    #     Muscle_Env,

    #     args.model,

    #     args.gamma,

    #     args.tau,

    #     args.lr,

    #     args.alpha,

    #     args.automatic_entropy_tuning,

    #     args.seed,

    #     args.policy_batch_size,

    #     args.hidden_size,

    #     args.policy_replay_size,

    #     args.multi_policy_loss,

    #     args.batch_iters,

    #     args.cuda,

    #     args.visualize,

    #     args.root_dir,

    #     args.checkpoint_file,

    #     args.checkpoint_folder,

    #     args.statistics_folder,

    #     args.total_episodes,

    #     args.load_saved_nets_for_training,

    #     args.save_iter,

    #     args.mode,

    #     args.musculoskeletal_model_path,

    #     args.initial_pose_path,

    #     args.kinematics_path,

    #     args.nusim_data_path,

    #     args.condition_selection_strategy,

    #     args.sim_dt,

    #     args.frame_repeat,

    #     args.n_fixedsteps,

    #     args.timestep_limit,

    #     args.trajectory_scaling,

    #     args.center

    # )

    ### TRAIN OR TEST ###

    if args.mode == "train":

        trainer.train()

    elif args.mode in ["test", "SFE", "sensory_pert", "neural_pert", "musculo_properties"]:

        trainer.test(args.test_data_filename)

    else:

        raise NotImplementedError

if __name__ == '__main__':

    main()