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")

[(6, 'matplotlib.pyplot.style.use', 'plt.style.use', (['"""ggplot"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (8, 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""PyTorch MNIST Example"""'}), False, 'import argparse\n'), (15, 'matplotlib.pyplot.subplot', 'plt.subplot', (['subplot'], {}), True, 'import matplotlib.pyplot as plt\n'), (20, 'numpy.mean', 'np.mean', (['data_first'], {'axis': '(0)'}), True, 'import numpy as np\n'), (21, 'numpy.mean', 'np.mean', (['data_second'], {'axis': '(0)'}), True, 'import numpy as np\n'), (25, 'matplotlib.pyplot.fill_between', 'plt.fill_between', (['x', '(y_first - y_first_err)', '(y_first + y_first_err)'], {'color': '"""m"""', 'alpha': '(0.3)'}), True, 'import matplotlib.pyplot as plt\n'), (26, 'matplotlib.pyplot.fill_between', 'plt.fill_between', (['x', '(y_second - y_second_err)', '(y_second + y_second_err)'], {'color': '"""c"""', 'alpha': '(0.3)'}), True, 'import matplotlib.pyplot as plt\n'), (27, 'matplotlib.pyplot.plot', 'plt.plot', (['x', 'y_first'], {'color': '"""r"""', 'label': '"""Task A"""'}), True, 'import matplotlib.pyplot as plt\n'), (28, 'matplotlib.pyplot.plot', 'plt.plot', (['x', 'y_second'], {'color': '"""g"""', 'label': '"""Task B (transfer learning)"""'}), True, 'import matplotlib.pyplot as plt\n'), (29, 'matplotlib.pyplot.legend', 'plt.legend', ([], {'bbox_to_anchor': '(0.8, 0.3)', 'loc': '(2)', 'ncol': '(1)', 'fontsize': '(15)'}), True, 'import matplotlib.pyplot as plt\n'), (30, 'matplotlib.pyplot.gca', 'plt.gca', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (38, 'matplotlib.pyplot.title', 'plt.title', (['title'], {'fontsize': '(20)', 'y': '(0.9)'}), True, 'import matplotlib.pyplot as plt\n'), (39, 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""Accuracy"""'], {'fontsize': '(15)'}), True, 'import matplotlib.pyplot as plt\n'), (40, 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""Generations"""'], {'fontsize': '(15)'}), True, 'import matplotlib.pyplot as plt\n'), (41, 'matplotlib.pyplot.grid', 'plt.grid', (['(True)'], {}), True, 'import matplotlib.pyplot as plt\n'), (81, 'matplotlib.pyplot.show', 'plt.show', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (17, 'numpy.arange', 'np.arange', (['(0)', '(100)'], {}), True, 'import numpy as np\n'), (19, 'numpy.arange', 'np.arange', (['(0)', '(500)'], {}), True, 'import numpy as np\n'), (22, 'numpy.std', 'np.std', (['data_first'], {'axis': '(0)'}), True, 'import numpy as np\n'), (23, 'numpy.std', 'np.std', (['data_second'], {'axis': '(0)'}), True, 'import numpy as np\n'), (47, 'pickle.load', 'pickle.load', (['f'], {}), False, 'import pickle\n'), (57, 'matplotlib.pyplot.show', 'plt.show', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (62, 'pickle.load', 'pickle.load', (['f'], {}), False, 'import pickle\n'), (46, 'os.path.join', 'os.path.join', (['"""./result/result_mnist.pickle"""'], {}), False, 'import os\n'), (61, 'os.path.join', 'os.path.join', (['"""./result/result_cifar_svhn.pickle"""'], {}), False, 'import os\n')]