import torch
import json
import numpy as np
from torch.autograd import Variable
import gzip
import yaml
from re import split
from matplotlib import pyplot


def showImg( im ):
    pyplot.imshow( im )
    pyplot.show()


def myOpen( fname, mode ):
    return open( fname, mode, encoding="utf-8" )

def readFile( fname ):
    opener, mode = ( gzip.open, 'rt' ) if fname[-3:] == '.gz' else ( open, 'r' )
    with opener( fname, mode ) as f:
        return f.read()

def readLines( fname ):
    return split('[\r\n]', readFile( fname ) )

def readJson( fname ):
    with myOpen( fname, 'r' ) as f:
        return json.load( f )

def writeFile( fname, contents ):
    with myOpen( fname, 'w' ) as f:
        f.write( contents )

def writeJson( fname, data ):
    with myOpen( fname, 'w') as outfile:
        json.dump(data, outfile)

def readYaml( fname ):
    with myOpen(fname, 'r') as fp:
        return yaml.load( fp )


config = readYaml('./config.yaml')


class averager(object):
    """Compute average for `torch.Variable` and `torch.Tensor`. """

    def __init__(self):
        self.reset()

    def add(self, v):
        if isinstance(v, Variable):
            count = v.data.numel()
            v = v.data.sum()
        elif isinstance(v, torch.Tensor):
            count = v.numel()
            v = v.sum()

        self.n_count += count
        self.sum += v

    def reset(self):
        self.n_count = 0
        self.sum = 0

    def val(self):
        res = 0
        if self.n_count != 0:
            res = self.sum / float(self.n_count)
        return res



def loadTrainedModel( model, opt ):
    """Load a pretrained model into given model"""

    print('loading pretrained model from %s' % opt.crnn)
    if( opt.cuda ):
        stateDict = torch.load(opt.crnn )
    else:
        stateDict = torch.load(opt.crnn, map_location={'cuda:0': 'cpu'} )

    #  Handle the case of some old torch version. It will save the data as module.<xyz> . Handle it
    if( list( stateDict.keys() )[0][:7] == 'module.' ):
        for key in list(stateDict.keys()): 
            stateDict[ key[ 7:] ] = stateDict[key]
            del stateDict[ key ]
    model.load_state_dict( stateDict )
    print('Completed loading pre trained model')

[((12, 4, 12, 23), 'matplotlib.pyplot.imshow', 'pyplot.imshow', ({(12, 19, 12, 21): 'im'}, {}), '(im)', False, 'from matplotlib import pyplot\n'), ((13, 4, 13, 17), 'matplotlib.pyplot.show', 'pyplot.show', ({}, {}), '()', False, 'from matplotlib import pyplot\n'), ((29, 15, 29, 29), 'json.load', 'json.load', ({(29, 26, 29, 27): 'f'}, {}), '(f)', False, 'import json\n'), ((37, 8, 37, 32), 'json.dump', 'json.dump', ({(37, 18, 37, 22): 'data', (37, 24, 37, 31): 'outfile'}, {}), '(data, outfile)', False, 'import json\n'), ((41, 15, 41, 30), 'yaml.load', 'yaml.load', ({(41, 26, 41, 28): 'fp'}, {}), '(fp)', False, 'import yaml\n'), ((81, 20, 81, 41), 'torch.load', 'torch.load', ({(81, 31, 81, 39): 'opt.crnn'}, {}), '(opt.crnn)', False, 'import torch\n'), ((83, 20, 83, 73), 'torch.load', 'torch.load', (), '', False, 'import torch\n')]