"""Generate charts from with .npy files containing custom callables through replay."""

import argparse
from datetime import datetime
import errno
import numpy as np
import matplotlib.pyplot as plt
import os
import pytz
import sys

def make_bar_plot(vals, title):
    print(len(vals))
    fig = plt.figure()
    plt.hist(vals, 10, facecolor='blue', alpha=0.5)
    plt.title(title)
    plt.xlim(1000,3000)
    return fig

def plot_trip_distribution(all_trip_energy_distribution):
    non_av_vals = []
    figures = []
    figure_names = []
    for key in all_trip_energy_distribution:
        if key != 'av':
            non_av_vals.extend(all_trip_energy_distribution[key])
        figures.append(make_bar_plot(all_trip_energy_distribution[key], key))
        figure_names.append(key)
    
    figure_names.append('All Non-AV')
    figures.append(make_bar_plot(non_av_vals, 'All Non-AV'))

    figure_names.append('All')
    figures.append(make_bar_plot(non_av_vals + all_trip_energy_distribution['av'], 'All'))

    return figure_names, figures
    
    

def parse_flags(args):
    """Parse training options user can specify in command line.

    Returns
    -------
    argparse.Namespace
        the output parser object
    """
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description="Parse argument used when running a Flow simulation.",
        epilog="python train.py EXP_CONFIG")
    parser.add_argument("target_folder", type=str,
                        help='Folder containing results')
    parser.add_argument("--output_folder", type=str, required=False, default=None,
                        help='Folder to save charts to.')
    parser.add_argument("--show_images", action='store_true',
                        help='Whether to display charts.')
    parser.add_argument("--heatmap", type=str, required=False,
                        help='Make a heatmap of the supplied variable.')
    return parser.parse_args(args)


if __name__ == "__main__":
    flags = parse_flags(sys.argv[1:])

    date = datetime.now(tz=pytz.utc)
    date = date.astimezone(pytz.timezone('US/Pacific')).strftime("%m-%d-%Y")

    if flags.output_folder:
        if not os.path.exists(flags.output_folder):
            try:
                os.makedirs(flags.output_folder)
            except OSError as exc:
                if exc.errno != errno.EEXIST:
                    raise

    info_dicts = []
    custom_callable_names = set()
    exp_names = []
    for (dirpath, dir_names, file_names) in os.walk(flags.target_folder):
        for file_name in file_names:
            if file_name[-8:] == "info.npy":
                exp_name = os.path.basename(dirpath)
                info_dict = np.load(os.path.join(dirpath, file_name), allow_pickle=True).item()

                info_dicts.append(info_dict)
                print(info_dict.keys())
                exp_names.append(exp_name)
                custom_callable_names.update(info_dict.keys())

    idxs = np.argsort(exp_names)
    exp_names = [exp_names[i] for i in idxs]
    info_dicts = [info_dicts[i] for i in idxs]

    if flags.heatmap is not None:
        heatmap = np.zeros((4, 6))
        pr_spacing = np.around(np.linspace(0, 0.3, 4), decimals=2)
        apr_spacing = np.around(np.linspace(0, 0.5, 6), decimals=2)
        for exp_name, info_dict in zip(exp_names, info_dicts):
            apr_bucket = int(np.around(float(exp_name.split('_')[1][3:]) / 0.1))
            pr_bucket = int(np.around(float(exp_name.split('_')[0][2:]) / 0.1))

            if flags.heatmap not in info_dict:
                print(exp_name)
                continue
            else:
                val = np.mean(info_dict[flags.heatmap])
                print(exp_name, pr_bucket, pr_spacing[pr_bucket], apr_bucket, apr_spacing[apr_bucket], val)
                heatmap[pr_bucket, apr_bucket] = val

        fig = plt.figure()
        plt.imshow(heatmap, interpolation='nearest', cmap='seismic', aspect='equal', vmin=1500, vmax=3000)
        plt.title(flags.heatmap)
        plt.yticks(ticks=np.arange(len(pr_spacing)), labels=pr_spacing)
        plt.ylabel("AV Penetration")
        plt.xticks(ticks=np.arange(len(apr_spacing)), labels=apr_spacing)
        plt.xlabel("Aggressive Driver Penetration")
        plt.colorbar()
        plt.show()
        plt.close(fig)

    else:
        for name in custom_callable_names:
            y_vals = [np.mean(info_dict[name]) for info_dict in info_dicts]
            y_stds = [np.std(info_dict[name]) for info_dict in info_dicts]
            x_pos = np.arange(len(exp_names))

            plt.bar(x_pos, y_vals, align='center', alpha=0.5)
            plt.xticks(x_pos, [exp_name for exp_name in exp_names], rotation=60)
            plt.xlabel('Experiment')
            plt.title('I210 Replay Result: {}'.format(name))
            plt.tight_layout()
            if flags.output_folder:
                plt.savefig(os.path.join(flags.output_folder, '{}-plot.png'.format(name)))

            plt.show()

[((14, 10, 14, 22), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((15, 4, 15, 51), 'matplotlib.pyplot.hist', 'plt.hist', (), '', True, 'import matplotlib.pyplot as plt\n'), ((16, 4, 16, 20), 'matplotlib.pyplot.title', 'plt.title', ({(16, 14, 16, 19): 'title'}, {}), '(title)', True, 'import matplotlib.pyplot as plt\n'), ((17, 4, 17, 23), 'matplotlib.pyplot.xlim', 'plt.xlim', ({(17, 13, 17, 17): '(1000)', (17, 18, 17, 22): '(3000)'}, {}), '(1000, 3000)', True, 'import matplotlib.pyplot as plt\n'), ((48, 13, 51, 44), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((66, 11, 66, 36), 'datetime.datetime.now', 'datetime.now', (), '', False, 'from datetime import datetime\n'), ((80, 44, 80, 72), 'os.walk', 'os.walk', ({(80, 52, 80, 71): 'flags.target_folder'}, {}), '(flags.target_folder)', False, 'import os\n'), ((91, 11, 91, 32), 'numpy.argsort', 'np.argsort', ({(91, 22, 91, 31): 'exp_names'}, {}), '(exp_names)', True, 'import numpy as np\n'), ((96, 18, 96, 34), 'numpy.zeros', 'np.zeros', ({(96, 27, 96, 33): '(4, 6)'}, {}), '((4, 6))', True, 'import numpy as np\n'), ((111, 14, 111, 26), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((112, 8, 112, 106), 'matplotlib.pyplot.imshow', 'plt.imshow', (), '', True, 'import matplotlib.pyplot as plt\n'), ((113, 8, 113, 32), 'matplotlib.pyplot.title', 'plt.title', ({(113, 18, 113, 31): 'flags.heatmap'}, {}), '(flags.heatmap)', True, 'import matplotlib.pyplot as plt\n'), ((115, 8, 115, 36), 'matplotlib.pyplot.ylabel', 'plt.ylabel', ({(115, 19, 115, 35): '"""AV Penetration"""'}, {}), "('AV Penetration')", True, 'import matplotlib.pyplot as plt\n'), ((117, 8, 117, 51), 'matplotlib.pyplot.xlabel', 'plt.xlabel', ({(117, 19, 117, 50): '"""Aggressive Driver Penetration"""'}, {}), "('Aggressive Driver Penetration')", True, 'import matplotlib.pyplot as plt\n'), ((118, 8, 118, 22), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((119, 8, 119, 18), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((120, 8, 120, 22), 'matplotlib.pyplot.close', 'plt.close', ({(120, 18, 120, 21): 'fig'}, {}), '(fig)', True, 'import matplotlib.pyplot as plt\n'), ((70, 15, 70, 50), 'os.path.exists', 'os.path.exists', ({(70, 30, 70, 49): 'flags.output_folder'}, {}), '(flags.output_folder)', False, 'import os\n'), ((97, 31, 97, 53), 'numpy.linspace', 'np.linspace', ({(97, 43, 97, 44): '0', (97, 46, 97, 49): '0.3', (97, 51, 97, 52): '4'}, {}), '(0, 0.3, 4)', True, 'import numpy as np\n'), ((98, 32, 98, 54), 'numpy.linspace', 'np.linspace', ({(98, 44, 98, 45): '0', (98, 47, 98, 50): '0.5', (98, 52, 98, 53): '6'}, {}), '(0, 0.5, 6)', True, 'import numpy as np\n'), ((128, 12, 128, 61), 'matplotlib.pyplot.bar', 'plt.bar', (), '', True, 'import matplotlib.pyplot as plt\n'), ((129, 12, 129, 80), 'matplotlib.pyplot.xticks', 'plt.xticks', (), '', True, 'import matplotlib.pyplot as plt\n'), ((130, 12, 130, 36), 'matplotlib.pyplot.xlabel', 'plt.xlabel', ({(130, 23, 130, 35): '"""Experiment"""'}, {}), "('Experiment')", True, 'import matplotlib.pyplot as plt\n'), ((132, 12, 132, 30), 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((136, 12, 136, 22), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((67, 27, 67, 54), 'pytz.timezone', 'pytz.timezone', ({(67, 41, 67, 53): '"""US/Pacific"""'}, {}), "('US/Pacific')", False, 'import pytz\n'), ((72, 16, 72, 48), 'os.makedirs', 'os.makedirs', ({(72, 28, 72, 47): 'flags.output_folder'}, {}), '(flags.output_folder)', False, 'import os\n'), ((83, 27, 83, 52), 'os.path.basename', 'os.path.basename', ({(83, 44, 83, 51): 'dirpath'}, {}), '(dirpath)', False, 'import os\n'), ((107, 22, 107, 55), 'numpy.mean', 'np.mean', ({(107, 30, 107, 54): 'info_dict[flags.heatmap]'}, {}), '(info_dict[flags.heatmap])', True, 'import numpy as np\n'), ((124, 22, 124, 46), 'numpy.mean', 'np.mean', ({(124, 30, 124, 45): 'info_dict[name]'}, {}), '(info_dict[name])', True, 'import numpy as np\n'), ((125, 22, 125, 45), 'numpy.std', 'np.std', ({(125, 29, 125, 44): 'info_dict[name]'}, {}), '(info_dict[name])', True, 'import numpy as np\n'), ((84, 36, 84, 68), 'os.path.join', 'os.path.join', ({(84, 49, 84, 56): 'dirpath', (84, 58, 84, 67): 'file_name'}, {}), '(dirpath, file_name)', False, 'import os\n')]