import torch
import pickle


def load_model_training(path, epoch, ac, ac_targ, pi_optimizer, q_optimizer):
  """Load the model and all things necessary for continueing the training.
  
  Args:
    path (str):     Path to the directory of checkpoints to save
    epoch (int):    Current epoch number
    ac:             Actor-critic model
    ac_targ:        Target actor-critic model
    pi_optimizer:   Adam optimizer for policy
    q_optimizer:    Adam optimizer for q-value function approximator
  
  Returns:
    ac:             Actor-critic model
    ac_targ:        Target actor-critic model
    replay_buffer:  Buffer that contains all transitions
    pi_optimizer:   Adam optimizer for policy
    q_optimizer:    Adam optimizer for q-value function approximator
    t_total (int):  Total number of timesteps
  """
  checkpoint_path = path + '/checkpoint_' + str(epoch) + '.pt'
  checkpoint = torch.load(checkpoint_path)
  ac.load_state_dict(checkpoint['state_dict'])
  ac_targ.load_state_dict(checkpoint['target_state_dict'])
  pi_optimizer.load_state_dict(checkpoint['pi_optimizer'])
  q_optimizer.load_state_dict(checkpoint['q_optimizer'])
  # replay buffer
  rb_path = path + '/replay_buffer.pkl'
  with open(rb_path, 'rb') as inp:
    replay_buffer = pickle.load(inp)
  return ac, ac_targ, replay_buffer, pi_optimizer, q_optimizer, checkpoint['t_total']