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', True, 'import matplotlib.pyplot as plt\n'), (8, 'argparse.ArgumentParser', 'argparse.ArgumentParser', False, 'import argparse\n'), (15, 'matplotlib.pyplot.subplot', 'plt.subplot', True, 'import matplotlib.pyplot as plt\n'), (20, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (21, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (25, 'matplotlib.pyplot.fill_between', 'plt.fill_between', True, 'import matplotlib.pyplot as plt\n'), (26, 'matplotlib.pyplot.fill_between', 'plt.fill_between', True, 'import matplotlib.pyplot as plt\n'), (27, 'matplotlib.pyplot.plot', 'plt.plot', True, 'import matplotlib.pyplot as plt\n'), (28, 'matplotlib.pyplot.plot', 'plt.plot', True, 'import matplotlib.pyplot as plt\n'), (29, 'matplotlib.pyplot.legend', 'plt.legend', 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', True, 'import matplotlib.pyplot as plt\n'), (39, 'matplotlib.pyplot.ylabel', 'plt.ylabel', True, 'import matplotlib.pyplot as plt\n'), (40, 'matplotlib.pyplot.xlabel', 'plt.xlabel', True, 'import matplotlib.pyplot as plt\n'), (41, 'matplotlib.pyplot.grid', 'plt.grid', 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', True, 'import numpy as np\n'), (19, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (22, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (23, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (47, 'pickle.load', 'pickle.load', False, 'import pickle\n'), (57, 'matplotlib.pyplot.show', 'plt.show', True, 'import matplotlib.pyplot as plt\n'), (62, 'pickle.load', 'pickle.load', False, 'import pickle\n'), (46, 'os.path.join', 'os.path.join', False, 'import os\n'), (61, 'os.path.join', 'os.path.join', False, 'import os\n')]