import os
import sys
import numpy as np
import matplotlib.pyplot as plt
import math
import pylab
import sys
import argparse
import re
from pylab import figure, show, legend, ylabel

from mpl_toolkits.axes_grid1 import host_subplot

if __name__ == "__main__":
    plt.ion()
    host = host_subplot(111)
    host.set_xlabel("Iterations")
    host.set_ylabel("Loss")
    plt.subplots_adjust(right=0.75)


    while True:
        parser = argparse.ArgumentParser(description='makes a plot from Caffe output')
        parser.add_argument('output_file', help='file of captured stdout and stderr')
        args = parser.parse_args()

        f = open(args.output_file, 'r')

        training_iterations = []
        training_loss = []

        test_iterations = []
        test_accuracy = []
        test_loss = []

        check_test = False
        check_test2 = False
        for line in f:

            # if check_test:
            #     #test_accuracy.append(float(line.strip().split(' = ')[-1]))
            #     check_test = False
            #     check_test2 = True
            # elif check_test2:
            if 'Test net output' in line and 'loss = ' in line:
                # print line
                #print line.strip().split(' ')
                test_loss.append(float(line.strip().split(' ')[-2]))
                check_test2 = False
            # else:
            #     test_loss.append(0)
            #     check_test2 = False

            if '] Iteration ' in line and 'loss = ' in line:
                arr = re.findall(r'ion \b\d+\b,', line)
                training_iterations.append(int(arr[0].strip(',')[4:]))
                training_loss.append(float(line.strip().split(' = ')[-1]))

            if '] Iteration ' in line and 'Testing net' in line:
                arr = re.findall(r'ion \b\d+\b,', line)
                test_iterations.append(int(arr[0].strip(',')[4:]))
                check_test = True

        print 'train iterations len: ', len(training_iterations)
        print 'train loss len: ', len(training_loss)
        print 'test loss len: ', len(test_loss)
        print 'test iterations len: ', len(test_iterations)
        #print 'test accuracy len: ', len(test_accuracy)

        # if len(test_iterations) != len(test_accuracy):  # awaiting test...
        #     print 'mis-match'
        #     print len(test_iterations[0:-1])
        #     test_iterations = test_iterations[0:-1]

        f.close()
        #  plt.plot(training_iterations, training_loss, '-', linewidth=2)
        #  plt.plot(test_iterations, test_accuracy, '-', linewidth=2)
        #  plt.show()

        # host = host_subplot(111)  # , axes_class=AA.Axes)
        # plt.subplots_adjust(right=0.75)

        #par1 = host.twinx()

        # host.set_xlabel("iterations")
        # host.set_ylabel("log loss")
        #par1.set_ylabel("validation accuracy")

        host.clear()
        host.clear()
        host.set_xlabel("Iterations")
        host.set_ylabel("Loss")
        #p1, = host.plot(training_iterations, training_loss, label="training loss")
        if len(training_iterations) == len(training_loss):
            p1, = host.plot(training_iterations, training_loss, label="training loss")
        if len(test_iterations) == len(test_loss):
            p3, = host.plot(test_iterations, test_loss, label="valdation loss")
        #p2, = par1.plot(test_iterations, test_accuracy, label="validation accuracy")

        host.legend(loc=2)

        #host.axis["left"].label.set_color(p1.get_color())
        #par1.axis["right"].label.set_color(p2.get_color())
        #fig = plt.figure()
        #fig.patch.set_facecolor('white')

        #axes = plt.gca()
        #ymin, ymax = min(training_loss), max(training_loss)
        #axes.set_xlim([xmin, xmax])
        #axes.set_ylim([0, ymax])
        #plt.yticks([0, 0.2, 0.4, 0.6, 0.8,1.0, 1.2, 1.4, 1.6])
        plt.grid()
        plt.draw()
        plt.show()
        plt.pause(5)