import sqlite3


def get_data_for_episode(conn: sqlite3.Connection, table_name: str, ep_num: int):
    """
    :param conn: Sqlite3 connection object
    :param table_name: Name of table to be queried
    :param ep_num: Episode to get data from
    :return: action, state, reward, step_num, joint_pos, joint_vel, cum_reward
    """

    cur = conn.cursor()
    sqlite_statement = \
        f''' 
            SELECT action, state, reward, step_num, joint_pos, joint_vel, cum_reward
            FROM {table_name} where episode_num == {ep_num}
        '''

    cur.execute(sqlite_statement)
    data = cur.fetchall()
    data2 = [*zip(*data)]
    return tuple(data2)