import copy


def train_qtab(agent, environment, episodes):

    """
    Trains a tabular Q-learning agent.

    Parameters
    ----------
    :param agent : AgentQtab instance
           Tabular Q-learning agent to be trained.
    :param environment : Env instance
           Environment instance specifying and simulating the dynamics of
           the environment the agent is in.
    :param episodes : int
           Number of episodes to train the agent.

    Returns
    -------
    :returns agent : AgentQtab instance
             Trained agent.
    :returns states : list of lists
             States experienced in each episode.
    :returns actions : list of lists
             Actions taken in each episode.
    :returns rewards : list of lists
             Rewards received in each episode.
    :returns new_states : list of lists
             New states experienced in each episode.
    :returns pred : list of ndarrays
             Q-value predictions made during the training.
    """

    print("Training tabular Q-learning agent.")

    agent = copy.deepcopy(agent)
    env = copy.deepcopy(environment)

    # create lists to save variables:
    states = [[] for _ in range(episodes)]
    actions = [[] for _ in range(episodes)]
    rewards = [[] for _ in range(episodes)]
    new_states = [[] for _ in range(episodes)]
    pred = []

    for episode in range(episodes):

        # reset environment to start state:
        env.reset()

        # increase iteration count in agent:
        agent.iter += 1

        # train agent:
        while not env.done:

            s = env.get_state()
            a = agent.choose_action()
            r, ss, done, _ = env.take_action(a)

            states[episode].append(s)
            actions[episode].append(a)
            rewards[episode].append(r)
            new_states[episode].append(ss)

            q_tab = agent.update(a, r)
            pred.append(q_tab)

    return agent, states, actions, rewards, new_states, pred