import time
import sys
import os
import shutil
import matplotlib
import valid_tool as vt
import tester
from func import *
matplotlib.use('Agg')
import matplotlib.pyplot as plt
plt.ion()

class ProgressWriter:
    def __init__(self, total, width=80):
        self.total = total
        self.counter = 0
        self.progress = 0
        self.width = width

    def increment(self, amount=1):
        self.counter += amount
        while self.counter / float(self.total) > \
              self.progress / float(self.width):
            sys.stdout.write('.')
            sys.stdout.flush()
            self.progress += 1
        if self.counter == self.total and \
           self.progress < self.width:
            print

class Logger:
    def __init__(self, trainer, csv=True, filename=None):
        self.trainer = trainer
        self.startTime = time.time()
        self.saveCsv = csv
        if filename is None:
            self.outFilename = os.path.join(
                trainer.outputFolder, trainer.name + '.csv')
        else:
            self.outFilename = filename

    def logMsg(self, msg):
        print msg

    def logTrainStats(self):
        timeElapsed = time.time() - self.startTime
        stats = 'N: %3d T: %5d  TE: %8.4f  TR: %8.4f  VE: %8.4f  VR: %8.4f' % \
                (self.trainer.epoch,
                 timeElapsed,
                 self.trainer.loss[self.trainer.totalStep],
                 self.trainer.rate[self.trainer.totalStep],
                 self.trainer.validLoss[self.trainer.totalStep],
                 self.trainer.validRate[self.trainer.totalStep])
        print stats

        if self.saveCsv:
            statsCsv = '%d,%.4f,%.4f,%.4f,%.4f' % \
                    (self.trainer.epoch,
                     self.trainer.loss[self.trainer.totalStep],
                     self.trainer.rate[self.trainer.totalStep],
                     self.trainer.validLoss[self.trainer.totalStep],
                     self.trainer.validRate[self.trainer.totalStep])
            with open(self.outFilename, 'a+') as f:
                f.write('%s\n' % statsCsv)
    pass

class Plotter:
    def __init__(self, trainer):
        self.trainer = trainer
        self.startTime = time.time()
        self.trainer.epoch = 0
        self.lossFigFilename = \
            os.path.join(trainer.outputFolder, trainer.name + '_loss.png')
        self.errFigFilename = \
            os.path.join(trainer.outputFolder, trainer.name + '_err.png')
        self.trainer.epoch = 0

    def plot(self):
        plt.figure(1)
        plt.clf()
        plt.plot(np.arange(self.trainer.totalStep + 1),
                 self.trainer.loss[0 : self.trainer.totalStep + 1], 
                 'b-x')
        if self.trainer.trainOpt['needValid']:
            plt.plot(np.arange(self.trainer.totalStep + 1),
                     self.trainer.validLoss[0 : self.trainer.totalStep + 1], 
                     'g-o')
            plt.legend(['Train', 'Valid'])