import sqlite3


def get_highest_rew_for_each_episode(conn: sqlite3.Connection, table_name: str):
    """
    This function is to be used for plotting reward vs episode num graph.
    Should be called after determining the correct table to query,
    using get_all_tables_ep_num and some additional logic
    :param conn: Sqlite database connection object
    :param table_name: Name of table to be queried
    :return: max(cum_reward), episode_num, step_num
    """
    cur = conn.cursor()
    sqlite_statement = \
        f''' 
        SELECT MAX(cum_reward), episode_num, step_num FROM {table_name} GROUP BY episode_num;
        '''
    cur.execute(sqlite_statement)
    data = cur.fetchall()
    data2 = [*zip(*data)]
    return tuple(data2)