import argparse
import os
import pickle
import numpy as np
import matplotlib.pyplot as plt
plt.style.use('ggplot')

parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
parser.add_argument('--mnist', action='store_true', default=False,
                    help='open mnist result')
args = parser.parse_args()


def subplot(subplot, data_first, data_second, title):
    plt.subplot(subplot)
    if args.mnist:
        x = np.arange(0,100)
    else:
        x = np.arange(0,500)
    y_first = np.mean(data_first, axis=0)
    y_second = np.mean(data_second, axis=0)
    y_first_err = np.std(data_first, axis=0) / 2.
    y_second_err = np.std(data_second, axis=0) / 2. 
    
    plt.fill_between(x, y_first - y_first_err, y_first + y_first_err, color='m', alpha=0.3)
    plt.fill_between(x, y_second - y_second_err, y_second + y_second_err, color='c', alpha=0.3)
    plt.plot(x, y_first, color='r', label='Task A')
    plt.plot(x, y_second, color='g', label='Task B (transfer learning)')
    plt.legend(bbox_to_anchor=(0.8, 0.3), loc=2, ncol=1, fontsize=15)
    axes = plt.gca()

    if args.mnist:
        axes.set_xlim([0, 100])
        axes.set_ylim([0, 1.2])
    else:
        axes.set_xlim([0, 500])
        axes.set_ylim([0, 0.6])
    plt.title(title, fontsize=20, y = 0.9)
    plt.ylabel('Accuracy',fontsize=15)
    plt.xlabel('Generations',fontsize=15)
    plt.grid(True)


try: 
    if args.mnist:
        f = open(os.path.join('./result/result_mnist.pickle'))
        result = pickle.load(f)
        f.close()
        pathnet_first = []
        pathnet_second = []

        for res in result:
            pathnet_first.append(res[2])
            pathnet_second.append(res[3])

        subplot('111', pathnet_first, pathnet_second,'MNIST')
        plt.show()


    else:
        f = open(os.path.join('./result/result_cifar_svhn.pickle'))
        result = pickle.load(f)
        f.close()

        cifar_first = []
        cifar_second = []
        svhn_first = []
        svhn_second = []

        for res in result:
            if res[0] == 'pathnet_cifar_first':
                cifar_first.append(res[2])
                svhn_second.append(res[3])
            else:
                svhn_first.append(res[2])
                cifar_second.append(res[3])

    subplot('211', cifar_first, cifar_second,'CIFAR-10')
    subplot('212', svhn_first, svhn_second,'cSVHN')

    plt.show()

except IOError:
    print("Result file does not exist")

[((91, 114), 'matplotlib.pyplot.style.use', 'plt.style.use', (['"""ggplot"""'], {}), "('ggplot')\n", (104, 114), True, 'import matplotlib.pyplot as plt\n'), ((125, 185), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""PyTorch MNIST Example"""'}), "(description='PyTorch MNIST Example')\n", (148, 185), False, 'import argparse\n'), ((386, 406), 'matplotlib.pyplot.subplot', 'plt.subplot', (['subplot'], {}), '(subplot)\n', (397, 406), True, 'import matplotlib.pyplot as plt\n'), ((508, 535), 'numpy.mean', 'np.mean', (['data_first'], {'axis': '(0)'}), '(data_first, axis=0)\n', (515, 535), True, 'import numpy as np\n'), ((551, 579), 'numpy.mean', 'np.mean', (['data_second'], {'axis': '(0)'}), '(data_second, axis=0)\n', (558, 579), True, 'import numpy as np\n'), ((692, 783), 'matplotlib.pyplot.fill_between', 'plt.fill_between', (['x', '(y_first - y_first_err)', '(y_first + y_first_err)'], {'color': '"""m"""', 'alpha': '(0.3)'}), "(x, y_first - y_first_err, y_first + y_first_err, color='m',\n    alpha=0.3)\n", (708, 783), True, 'import matplotlib.pyplot as plt\n'), ((784, 880), 'matplotlib.pyplot.fill_between', 'plt.fill_between', (['x', '(y_second - y_second_err)', '(y_second + y_second_err)'], {'color': '"""c"""', 'alpha': '(0.3)'}), "(x, y_second - y_second_err, y_second + y_second_err, color\n    ='c', alpha=0.3)\n", (800, 880), True, 'import matplotlib.pyplot as plt\n'), ((880, 927), 'matplotlib.pyplot.plot', 'plt.plot', (['x', 'y_first'], {'color': '"""r"""', 'label': '"""Task A"""'}), "(x, y_first, color='r', label='Task A')\n", (888, 927), True, 'import matplotlib.pyplot as plt\n'), ((932, 1000), 'matplotlib.pyplot.plot', 'plt.plot', (['x', 'y_second'], {'color': '"""g"""', 'label': '"""Task B (transfer learning)"""'}), "(x, y_second, color='g', label='Task B (transfer learning)')\n", (940, 1000), True, 'import matplotlib.pyplot as plt\n'), ((1005, 1070), 'matplotlib.pyplot.legend', 'plt.legend', ([], {'bbox_to_anchor': '(0.8, 0.3)', 'loc': '(2)', 'ncol': '(1)', 'fontsize': '(15)'}), '(bbox_to_anchor=(0.8, 0.3), loc=2, ncol=1, fontsize=15)\n', (1015, 1070), True, 'import matplotlib.pyplot as plt\n'), ((1082, 1091), 'matplotlib.pyplot.gca', 'plt.gca', ([], {}), '()\n', (1089, 1091), True, 'import matplotlib.pyplot as plt\n'), ((1254, 1290), 'matplotlib.pyplot.title', 'plt.title', (['title'], {'fontsize': '(20)', 'y': '(0.9)'}), '(title, fontsize=20, y=0.9)\n', (1263, 1290), True, 'import matplotlib.pyplot as plt\n'), ((1297, 1332), 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""Accuracy"""'], {'fontsize': '(15)'}), "('Accuracy', fontsize=15)\n", (1307, 1332), True, 'import matplotlib.pyplot as plt\n'), ((1336, 1374), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""Generations"""'], {'fontsize': '(15)'}), "('Generations', fontsize=15)\n", (1346, 1374), True, 'import matplotlib.pyplot as plt\n'), ((1378, 1392), 'matplotlib.pyplot.grid', 'plt.grid', (['(True)'], {}), '(True)\n', (1386, 1392), True, 'import matplotlib.pyplot as plt\n'), ((2393, 2403), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (2401, 2403), True, 'import matplotlib.pyplot as plt\n'), ((438, 455), 'numpy.arange', 'np.arange', (['(0)', '(100)'], {}), '(0, 100)\n', (447, 455), True, 'import numpy as np\n'), ((477, 494), 'numpy.arange', 'np.arange', (['(0)', '(500)'], {}), '(0, 500)\n', (486, 494), True, 'import numpy as np\n'), ((598, 624), 'numpy.std', 'np.std', (['data_first'], {'axis': '(0)'}), '(data_first, axis=0)\n', (604, 624), True, 'import numpy as np\n'), ((649, 676), 'numpy.std', 'np.std', (['data_second'], {'axis': '(0)'}), '(data_second, axis=0)\n', (655, 676), True, 'import numpy as np\n'), ((1500, 1514), 'pickle.load', 'pickle.load', (['f'], {}), '(f)\n', (1511, 1514), False, 'import pickle\n'), ((1770, 1780), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (1778, 1780), True, 'import matplotlib.pyplot as plt\n'), ((1878, 1892), 'pickle.load', 'pickle.load', (['f'], {}), '(f)\n', (1889, 1892), False, 'import pickle\n'), ((1437, 1481), 'os.path.join', 'os.path.join', (['"""./result/result_mnist.pickle"""'], {}), "('./result/result_mnist.pickle')\n", (1449, 1481), False, 'import os\n'), ((1810, 1859), 'os.path.join', 'os.path.join', (['"""./result/result_cifar_svhn.pickle"""'], {}), "('./result/result_cifar_svhn.pickle')\n", (1822, 1859), False, 'import os\n')]