def create_connect_viewer(N, N1, imgs, count, W2):
    """
    Create the patch to show connections between layers.

    Parameters
    ----------
    N: int
        Number of rows.
    N1: int
        Number of elements in the first layer.
    imgs: ndarray
        Images of weights from the first layer.
    count: int
        Number of elements to show.
    W2: list
        Second hidden layer.
    """
    pv = PatchViewer((N, count), imgs.shape[1:3], is_color=imgs.shape[3] == 3)

    for i in xrange(N):
        w = W2[:, i]

        wneg = w[w < 0.]
        wpos = w[w > 0.]

        w /= np.abs(w).max()

        wa = np.abs(w)

        to_sort = zip(wa, range(N1), w)

        s = sorted(to_sort)

        for j in xrange(count):

            idx = s[N1-j-1][1]
            mag = s[N1-j-1][2]

            if mag > 0:
                act = (mag, 0)
            else:
                act = (0, -mag)

            pv.add_patch(imgs[idx, ...], rescale=True, activation=act)

    return pv

def show_samples(generator,Noise_Dim,data_obj,filename):
    if data_obj.pitch_scale:
        pitch_max = 1.0
    else:
        pitch_max = 108.0
    rows = 4
    sample_cols = 5
    input_noise = np.random.uniform(-1.0,1.0,(rows*sample_cols, Noise_Dim))
    samples = generator.predict(input_noise)
    topo_samples = samples.reshape(samples.shape[0],4,samples.shape[-1]/4)
    #get topological_view

    pv = PatchViewer(grid_shape=(rows,sample_cols + 1),patch_shape=(4,samples.shape[-1]/4), \
                     is_color=False)
    X = np.concatenate((data_obj.X_train,data_obj.X_val,data_obj.X_test),axis = 0)
    topo_X = X
    print('Shape of dataset is {}').format(X.shape)
    X = X.reshape(X.shape[0],X.shape[1]*X.shape[2])

    for i in xrange(topo_samples.shape[0]):
        topo_sample = patch_quantize_01(patch_thresholding(topo_samples[i,:]/pitch_max))
        pv.add_patch(topo_sample * 2. -1.,rescale=False)
        if(i + 1) % sample_cols ==0:
            sample = samples[i,:]
            dists = np.square(X - sample).sum(axis = 1)
            j = np.argmin(dists)
            match = patch_quantize_01(patch_thresholding(topo_X[j,:]/pitch_max))
            pv.add_patch(match*2-1,rescale=False,activation = 1)
    print "Saving %s ..."%filename
    pv.save(filename)

def create_patch_viewer(grid_shape, vis_chains, m):
    """
    Add the patches to show.

    Parameters
    ----------
    grid_shape: tuple
        The shape of the grid to show.
    vis_chains: numpy array
        Visibles chains.
    m: int
        Number of visible chains.
    """
    pv = PatchViewer(grid_shape, vis_chains.shape[1:3],
                     is_color=vis_chains.shape[-1] == 3)

    for i in xrange(m):
        pv.add_patch(vis_chains[i, :], rescale=False)

    return pv

    def __init__(self, generator, save_prefix, batch_size=20, grid_shape=(5, 4)):
        assert isinstance(generator, Generator)

        self.batch_sym = T.matrix('generate_batch')
        self.generate_f = theano.function([self.batch_sym],
                                          generator.dropout_fprop(self.batch_sym)[0])

        self.batch = generator.get_noise(batch_size).eval()
        self.save_prefix = save_prefix
        self.patch_viewer = PatchViewer(grid_shape=grid_shape, patch_shape=(32, 32),
                                        is_color=True)

class GenerateAndSave(TrainExtension):
    """
    Keeps track of what the generator in a (vanilla) GAN returns for a
    particular set of noise values.
    """

    def __init__(self, generator, save_prefix, batch_size=20, grid_shape=(5, 4)):
        assert isinstance(generator, Generator)

        self.batch_sym = T.matrix('generate_batch')
        self.generate_f = theano.function([self.batch_sym],
                                          generator.dropout_fprop(self.batch_sym)[0])

        self.batch = generator.get_noise(batch_size).eval()
        self.save_prefix = save_prefix
        self.patch_viewer = PatchViewer(grid_shape=grid_shape, patch_shape=(32, 32),
                                        is_color=True)

    def on_monitor(self, model, dataset, algorithm):
        samples = self.generate_f(self.batch).swapaxes(0, 3)

        self.patch_viewer.clear()
        for sample in samples:
            self.patch_viewer.add_patch(sample, rescale=True)

        fname = self.save_prefix + '.%05i.png' % model.monitor.get_epochs_seen()
        self.patch_viewer.save(fname)

    def on_monitor(self, *args, **kwargs):

        if not hasattr(self, 'record'):
            self.record = {}
            self.size = {}
            for dataset in self.datasets:
                assert tuple(dataset.view_converter.axes) == ('c', 0, 1, 'b')
                self.record[dataset] = dataset.get_topological_view().copy()
                self.size[dataset] = dataset.X.shape[0]
        else:
            for i, dataset in enumerate(self.datasets):
                size = self.size[dataset]
                assert dataset.X.shape[0] == size
                self.record[dataset] = np.concatenate((self.record[dataset], dataset.get_topological_view().copy()),
                        axis=-1)
                record_view = self.record[dataset].copy()
                record_view /= np.abs(record_view).max()
                pv = PatchViewer(grid_shape=(record_view.shape[3]/size, size),
                        patch_shape = record_view.shape[1:3], is_color = record_view.shape[0] == 3)
                for j in xrange(record_view.shape[3]):
                    pv.add_patch(np.transpose(record_view[:,:,:,j], (1, 2, 0)), rescale=False)
                print 'Dataset %d: ' % i
                pv.show()
                x = raw_input()

            chan_viewer.add_patch(block_viewer.image[:,:,viewer_dims] - 0.5)
            block_viewer.clear()

        main_viewer.add_patch(chan_viewer.image[:,:,viewer_dims] - 0.5)
        chan_viewer.clear()

    return copy.copy(main_viewer.image)


viewer_g = make_viewer(wvg, get_dims(model.n_g), (opts.height, opts.width), is_color=True)
viewer_h = make_viewer(wvh, get_dims(model.n_h), (opts.height, opts.width), is_color=True)
w_image = plot(wv)

viewer = PatchViewer((1, 3),
            (numpy.max((viewer_g.image.shape[0], viewer_h.image.shape[0], w_image.shape[0])), 
             numpy.max((viewer_g.image.shape[1], viewer_h.image.shape[1], w_image.shape[1]))),
            is_color = opts.color,
            pad=(0,10))

viewer_dims = slice(0, None) if opts.color else 0
viewer.add_patch(viewer_g.image[:,:, viewer_dims] - 0.5)
viewer.add_patch(viewer_h.image[:,:, viewer_dims] - 0.5)
viewer.add_patch(w_image[:,:, viewer_dims] - 0.5)

pl.axis('off')
pl.imshow(viewer.image, interpolation='nearest')
pl.savefig('filters_%s.png' % opts.path)
if not opts.noshow:
    pl.show()

from pylearn2.datasets.cifar10 import CIFAR10
from pylearn2.gui.patch_viewer import PatchViewer

dataset = CIFAR10(which_set = 'test')

pv = PatchViewer((10,1),(32,32),is_color=True)

T,y = dataset.get_batch_topo(10, include_labels = True)

for i in xrange(10):
    print dataset.label_names[y[i]]
    pv.add_patch(dataset.adjust_for_viewer(T[i,:,:,:]),rescale=False)

pv.show()

    if topo:
        X = Xt

    args = [X]

    Y_sequence = f(*args)

    for elem in Y_sequence:
        assert elem.shape[0] == m

    rows = m

    cols = 1+len(Y_sequence)


    pv = PatchViewer((rows, cols), (Xt.shape[1], Xt.shape[2]), is_color = True,
            pad = (8,8) )

    for i in xrange(m):

        #add original patch
        patch = Xt[i,:,:,:].copy()
        patch = dataset.adjust_for_viewer(patch)
        if patch.shape[-1] != 3:
            patch = np.concatenate( (patch,patch,patch), axis=2)
        pv.add_patch(patch, rescale = False, activation = (1,0,0))
        orig_patch = patch

        def label_to_vis(Y_elem):
            prod =  np.dot(Y_elem, templates)
            assert Y_elem.ndim == 1
            rval = np.zeros((1, prod.shape[0]))

    feat_dataset = DenseDesignMatrix(X = feat, view_converter = DefaultViewConverter([1, 1, feat.shape[1]] ) )

    print 'reassembling features'
    ns = 32 - size + 1
    depatchifier = ReassembleGridPatches( orig_shape  = (ns, ns), patch_shape=(1,1) )
    feat_dataset.apply_preprocessor(depatchifier)

    print 'making topological view'
    topo_feat = feat_dataset.get_topological_view()
    assert topo_feat.shape[0] == X.shape[0]

    print 'assembling visualizer'

    n = np.ceil(np.sqrt(model.nhid))

    pv3 = PatchViewer(grid_shape = (X.shape[0], num_filters), patch_shape=(ns,ns), is_color= False)
    pv4 = PatchViewer(grid_shape = (n,n), patch_shape = (size,size), is_color = True, pad = (7,7))
    pv5 = PatchViewer(grid_shape = (1,num_filters), patch_shape = (size,size), is_color = True, pad = (7,7))

    idx = sorted(range(model.nhid), key = lambda l : -topo_feat[:,:,:,l].std() )

    W = model.W.get_value()

    weights_view = dataset.get_weights_view( W.T )

    p_act = 1. / (1. + np.exp(- model.bias_hid.get_value()))
    p_act /= p_act.max()

    mu_act = np.abs(model.mu.get_value())
    mu_act /= mu_act.max()
    mu_act += 0.5

from pylearn2.gui.patch_viewer import PatchViewer
from pylearn2.utils import get_choice

print('Use test set?')
choices = {'y': 'test', 'n': 'train'}
which_set = choices[get_choice(choices)]

dataset = FoveatedNORB(which_set=which_set, center=True)

topo = dataset.get_topological_view()

b, r, c, ch = topo.shape

assert ch == 2

pv = PatchViewer((1, 2), (r, c), is_color=False)

i = 0
while True:
    patch = topo[i, :, :, :]
    patch = patch / np.abs(patch).max()

    pv.add_patch(patch[:,:,1], rescale=False)
    pv.add_patch(patch[:,:,0], rescale=False)

    pv.show()

    print(dataset.y[i])

    choices = {'g': 'goto image', 'q': 'quit'}

    normalizer = gdpdf_normalizer
elif t == SCORE:
    df = grad_func(- dataset.free_energy(X))
    mf = grad_func( - model.free_energy(X))

    normalizer = gpdf_normalizer
elif t == SCORED:
    df = grad_dir_func(- dataset.free_energy(X))
    mf = grad_dir_func( - model.free_energy(X))

    normalizer = gdpdf_normalizer
else:
    assert False


g.components.append(HeatMap(f = df, normalizer = normalizer))
dimg = g.render()

g.components.pop()
g.components.append(HeatMap(f = mf, normalizer = normalizer))

mimg = g.render()

from pylearn2.gui.patch_viewer import PatchViewer

pv = PatchViewer((2,1),dimg.shape[0:2], is_color = dimg.shape[2] == 3)
pv.add_patch(dimg,rescale = False)
pv.add_patch(mimg, rescale = False)
pv.show()

            block_viewer.clear()

        main_viewer.add_patch(chan_viewer.image[:,:,viewer_dims] - 0.5)
        chan_viewer.clear()

    return copy.copy(main_viewer.image)


w_image = plot(wv)
if opts.phi:
    phi_image = plot(phi)

nplots = 2 if opts.phi else 1
viewer = PatchViewer((1,nplots), 
            (w_image.shape[0], 
             w_image.shape[1]),
            is_color = opts.color,
            pad=(20,20))

viewer_dims = slice(0, None) if opts.color else 0
viewer.add_patch(w_image[:,:, viewer_dims] - 0.5)
if opts.phi:
    viewer_dims = slice(0, None) if opts.color else 1
    viewer.add_patch(phi_image[:,:,0] - 0.5)

pl.imshow(viewer.image, interpolation='nearest')
pl.savefig('filters_%s.png' % opts.path)
pl.close()

if not opts.noshow:
    viewer.show()

def get_dims(nf):
    num_rows = numpy.floor(numpy.sqrt(nf))
    return (int(num_rows), int(numpy.ceil(nf / num_rows)))

topo_shape = [opts.height, opts.width, opts.chans]
viewconv = DefaultViewConverter(topo_shape)
viewdims = slice(0, None) if opts.color else 0

# load model and retrieve parameters
model = serial.load(opts.path)
wv = model.Wv.get_value().T
if opts.mu:
    wv = wv * model.mu.get_value()[:, None]

wv_viewer = PatchViewer(get_dims(len(wv)), (opts.height, opts.width),
                        is_color = opts.color, pad=(2,2))
for i in xrange(len(wv)):
    topo_wvi = viewconv.design_mat_to_topo_view(wv[i:i+1])
    wv_viewer.add_patch(topo_wvi[0])
if opts.wv_only:
    wv_viewer.show()
    os.sys.exit()

wg = model.Wg.get_value()
wh = model.Wh.get_value()
wg_viewer2 = PatchViewer((opts.top, opts.top), (opts.height, opts.width),
                         is_color = opts.color, pad=(2,2))
wg_viewer1 = PatchViewer(get_dims(len(wg)/opts.top),
                         (wg_viewer2.image.shape[0], wg_viewer2.image.shape[1]),
                         is_color = opts.color, pad=(2,2))
for i in xrange(0, len(wg), opts.top):

    fp.close()
    print 'Applying inverse pipeline ...'
    preproc.inverse(dataset)
samples = dataset.X

# check for global scaling
if not opts.local:
    samples = samples / numpy.abs(samples).max()

##############
# PLOT FILTERS
##############

import pdb; pdb.set_trace()
viewer = PatchViewer(get_dims(model.batch_size),
                     (opts.height, opts.width),
                     is_color = opts.color,
                     pad=(2,2))

topo_shape = [opts.height, opts.width, opts.chans]
view_converter = DefaultViewConverter(topo_shape)
topo_view = view_converter.design_mat_to_topo_view(samples)

for chan_i in xrange(nplots):

    topo_chan = topo_view if opts.color else topo_view[..., chan_i:chan_i+1]

    for bi in xrange(model.batch_size):
        viewer.add_patch(topo_chan[bi])

    #pl.subplot(1, nplots, chan_i+1)
    #pl.imshow(viewer.image, interpolation=None)

        zeroed = dataset.get_topological_view(dataset.mapback_for_viewer(dataset.get_design_matrix(zeroed)))
        M = dataset.get_topological_view(dataset.mapback_for_viewer(design_X))
        print (M.min(), M.max())
        M_sequence = [ dataset.get_topological_view(dataset.mapback_for_viewer(mat)) for mat in design_X_sequence ]
    X = dataset.adjust_to_be_viewed_with(Xt,Xt,per_example=True)
    if X_sequence[0].ndim == 2:
        X_sequence = [ dataset.get_topological_view(mat) for mat in X_sequence ]
    X_sequence = [ dataset.adjust_to_be_viewed_with(mat,Xt,per_example=True) for mat in X_sequence ]

    if cost.supervised:
        rows += m

        Y_sequence = outputs[end_X_outputs:]


    pv = PatchViewer( (rows, cols), (X.shape[1], X.shape[2]), is_color = True,
            pad = (8,8) )

    for i in xrange(m):

        """
        #add original patch
        patch = X[i,:,:,:]
        if patch.shape[-1] != 3:
            patch = np.concatenate( (patch,patch,patch), axis=2)
        pv.add_patch(patch, rescale = False, activation = (1,0,0))
        orig_patch = patch

        #mark the masked areas as red
        mask_patch = drop_mask[i,:,:,0]
        if drop_mask.shape[-1] > 1 and mask_gen.sync_channels and mask_gen.n_channels > 1:
            assert np.all(mask_patch == drop_mask[i,:,:,1])

from pylearn2.config import yaml_parse
from pylearn2.gui.patch_viewer import PatchViewer
import numpy as np

dataset = yaml_parse.load(model.dataset_yaml_src)

grid_shape = None

rows = 4
sample_cols = 5

# For some reason format_as from VectorSpace is not working right
samples = model.generator.sample(rows * sample_cols).eval()
topo_samples = dataset.get_topological_view(samples)

pv = PatchViewer(grid_shape=(rows, sample_cols + 1), patch_shape=(32,32),
        is_color=True)
scale = np.abs(samples).max()

X = dataset.X
topo = dataset.get_topological_view()
index = 0
for i in xrange(samples.shape[0]):
    topo_sample = topo_samples[i, :, :, :]
    print topo_sample.min(), topo_sample.max()
    pv.add_patch(topo_sample / scale, rescale=False)

    if (i +1) % sample_cols == 0:
        sample = samples[i, :]
        dists = np.square(X - sample).sum(axis=1)
        j = np.argmin(dists)
        match = topo[j, :]

dataset = yaml_parse.load(model.dataset_yaml_src)

rows = 10
cols = 3
m = rows * cols

X = dataset.get_batch_design(100)

drop_mask, X_hat = f(X)

X, drop_mask, X_hat = [ dataset.get_topological_view(mat)
        for mat in [X, drop_mask, X_hat] ]
X = dataset.adjust_for_viewer(X)
X_hat = dataset.adjust_for_viewer(X_hat)

pv = PatchViewer( (rows, cols*3), (X.shape[1], X.shape[2]), is_color = True)

for i in xrange(m):
    #add original patch
    patch = X[i,:,:,:]
    if patch.shape[-1] != 3:
        patch = np.concatenate( (patch,patch,patch), axis=2)
    pv.add_patch(patch, rescale = False)

    #mark the masked areas as red
    mask_patch = drop_mask[i,:,:,0]
    red_channel = patch[:,:,0]
    green_channel = patch[:,:,1]
    blue_channel = patch[:,:,2]
    red_channel[mask_patch == 1] = 1.
    green_channel[mask_patch == 1] = -1.

rows = 10
cols = 5

X = dataset.get_batch_design(rows * cols )

noise = model.dnce.noise_conditional

assert noise.is_symmetric()

Y = function([],noise.random_design_matrix(sharedX(X)))()

prob = function([],T.nnet.sigmoid(model.free_energy(sharedX(Y))-model.free_energy(sharedX(X))))()
assert prob.ndim == 1

Xt = dataset.get_topological_view(X)
Yt = dataset.get_topological_view(Y)
Xt = dataset.adjust_for_viewer(Xt)
Yt = dataset.adjust_for_viewer(Yt)

pv = PatchViewer( (rows, cols * 2), Xt.shape[1:3], is_color = Xt.shape[-1] == 3)

for i in xrange(Xt.shape[0]):
    assert prob[i] >= 0.0
    assert prob[i] <= 1.0
    assert not np.isnan(prob[i])
    assert not np.isinf(prob[i])
    pv.add_patch(Xt[i,:,:,:], activation = prob[i])
    pv.add_patch(Yt[i,:,:,:], activation = 1. - prob[i])

pv.show()

        display_axes = ('b', 0, 1, 'c')

        X = Conv2DSpace.convert_numpy(X, model_axes, display_axes)
        Xt = Conv2DSpace.convert_numpy(Xt, model_axes, display_axes)
        drop_mask = Conv2DSpace.convert_numpy(drop_mask, model_axes, display_axes)
        X_sequence = [Conv2DSpace.convert_numpy(elem, model_axes, display_axes) for elem in X_sequence]
    X_sequence = [ dataset.adjust_to_be_viewed_with(mat,Xt,per_example=True) for mat in X_sequence ]

    if cost.supervised:
        rows += m

        drop_mask_Y = outputs[end_X_outputs]
        Y_sequence = outputs[end_X_outputs+1:]


    pv = PatchViewer( (rows, cols), (X.shape[1], X.shape[2]), is_color = True,
            pad = (8,8) )

    for i in xrange(m):

        #add original patch
        patch = X[i,:,:,:]
        if patch.shape[-1] != 3:
            assert patch.shape[-1] == 2, patch.shape
            patch = patch[:,:,0:1]
            patch = np.concatenate( (patch,patch,patch), axis=2)
        pv.add_patch(patch, rescale = False, activation = (1,0,0))
        orig_patch = patch

        #mark the masked areas as red
        mask_patch = drop_mask[i,:,:,0]
        if drop_mask.shape[-1] > 1 and mask_gen.sync_channels and mask_gen.n_channels > 1: