"""Plot results from ray-based simulations

This method accepts as input the progress file generated by ray
(usually stored at ~/ray_results/.../progress.csv)
as well as the column(s) to be plotted.

If no column is specified, all existing columns will be printed.

Example usage
-----
::
    python ray_plot_progress.py </path/to/file>.csv mean_reward max_reward
"""

import csv
import argparse
import matplotlib.pyplot as plt
from collections import defaultdict


EXAMPLE_USAGE = 'ray_plot_progress.py ' + \
    '~/ray_results/experiment-tag/experiment-name/seed-id/progress.csv ' + \
    'evaluation/return-average training/return-average'


def plot_progress(filepath, columns):
    """
    Plot the values contained in the csv file at <filepath> for each column
    in the list of string columns.
    """
    data = defaultdict(list)

    with open(filepath) as f:
        # if columns list is empty, print a list of all columns and return
        if not columns:
            reader = csv.reader(f)
            print('Columns are: ' + ', '.join(next(reader)))
            return

        try:
            reader = csv.DictReader(f)
            for row in reader:
                for col in columns:
                    data[col].append(float(row[col]))
        except KeyError:
            print(f'Error: {__file__} was called with an unknown column name '
                  f'"{col}". Run `python {__file__} {filepath}` to get a list '
                  f'of all the existing columns.')
            raise
        except ValueError:
            print(f'Error: {__file__} was called with an invalid column name '
                  f'"{col}". This column contains values that are not '
                  f'convertible to floats.')
            raise

    for col_name, values in data.items():
        plt.plot(values, label=col_name)
    plt.legend()
    plt.show()


def create_parser():

    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description='[Flow] Plots progress.csv file generated by ray.',
        epilog='Example usage:\n\t' + EXAMPLE_USAGE)

    parser.add_argument('file', type=str, help='Path to the csv file.')
    parser.add_argument(
        'columns', type=str, nargs='*', help='Names of the columns to plot.')

    return parser


if __name__ == '__main__':
    parser = create_parser()
    args = parser.parse_args()
    plot_progress(args.file, args.columns)
"""Plot results from ray-based simulations

This method accepts as input the progress file generated by ray
(usually stored at ~/ray_results/.../progress.csv)
as well as the column(s) to be plotted.

If no column is specified, all existing columns will be printed.

Example usage
-----
::
    python plot_ray_results.py </path/to/file>.csv mean_reward max_reward
"""

import csv
import argparse
import matplotlib.pyplot as plt
from collections import defaultdict


EXAMPLE_USAGE = 'plot_ray_results.py ' + \
    '~/ray_results/experiment-tag/experiment-name/seed-id/progress.csv ' + \
    'evaluation/return-average training/return-average'


def plot_progress(filepath, columns):
    """
    Plot the values contained in the csv file at <filepath> for each column
    in the list of string columns.
    """
    data = defaultdict(list)

    with open(filepath) as f:
        # if columns list is empty, print a list of all columns and return
        if not columns:
            reader = csv.reader(f)
            print('Columns are: ' + ', '.join(next(reader)))
            return

        try:
            reader = csv.DictReader(f)
            for row in reader:
                for col in columns:
                    data[col].append(float(row[col]))
        except KeyError:
            print(f'Error: {__file__} was called with an unknown column name '
                  f'"{col}". Run `python {__file__} {filepath}` to get a list '
                  f'of all the existing columns.')
            raise
        except ValueError:
            print(f'Error: {__file__} was called with an invalid column name '
                  f'"{col}". This column contains values that are not '
                  f'convertible to floats.')
            raise

    for col_name, values in data.items():
        plt.plot(values, label=col_name)
    plt.legend()
    plt.show()


def create_parser():

    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description='[Flow] Plots progress.csv file generated by ray.',
        epilog='Example usage:\n\t' + EXAMPLE_USAGE)

    parser.add_argument('file', type=str, help='Path to the csv file.')
    parser.add_argument(
        'columns', type=str, nargs='*', help='Names of the columns to plot.')

    return parser


if __name__ == '__main__':
    parser = create_parser()
    args = parser.parse_args()
    plot_progress(args.file, args.columns)
fix filename in plot_ray_results.py