from hypergan.samplers.base_sampler import BaseSampler

import tensorflow as tf
import numpy as np

class BeganSampler(BaseSampler):
    def __init__(self, gan):
        BaseSampler.__init__(self, gan)
        self.x_v = None
        self.z_v = None
        self.created = False

    def sample(self, path, save_samples=False):
        gan = self.gan
        x_t = gan.inputs.x
        g_t = gan.generator.sample
        z_t = gan.encoder.sample
        
        rx_t = gan.discriminator.reconstruction

        sess = gan.session
        config = gan.config
        if(not self.created):
            self.x_v, self.z_v = sess.run([x_t, z_t])
            self.created=True

        rx_v, g_v = sess.run([rx_t, g_t], {x_t: self.x_v, z_t: self.z_v})
        stacks = []
        bs = gan.batch_size() // 2
        width = 8
        for i in range(1):
            stacks.append([self.x_v[i*width+width+j] for j in range(width)])
        for i in range(1):
            stacks.append([rx_v[i*width+width+j] for j in range(width)])
        for i in range(1):
            stacks.append([g_v[i*width+width+j] for j in range(width)])

        #[print(np.shape(s)) for s in stacks]
        images = np.vstack([np.hstack(s) for s in stacks])

        self.plot(images, path, save_samples)
        return [{'image': path, 'label': 'tiled x sample'}]


@@ -152,16 +152,33 @@
 elf, gan
+, samples_per_row
 ):%0A     
@@ -210,16 +210,33 @@
 elf, gan
+, samples_per_row
 )%0A      
@@ -307,24 +307,25 @@
  False%0A%0A
+%0A
     def 
 sample(s
@@ -316,16 +316,17 @@
     def 
+_
 sample(s
@@ -332,34 +332,8 @@
 self
-, path, save_samples=False
 ):%0A 
@@ -366,61 +366,27 @@
     
-x_t = gan.inputs.x%0A        g_t = gan.generator.sample
+config = gan.config
 %0A   
@@ -394,90 +394,53 @@
     
-z_t
+sess
  = gan.
-encoder.sample%0A        %0A        rx_t = gan.discriminator.reconstruction%0A
+session%0A        x_t = gan.inputs.x
 %0A   
@@ -448,54 +448,32 @@
     
-sess
+z_t
  = gan.
-session%0A        config = gan.config
+encoder.sample
 %0A   
@@ -583,16 +583,98 @@
 ed=True%0A
+        g_t = gan.generator.sample%0A        rx_t = gan.discriminator.reconstruction
 %0A       
@@ -815,41 +815,29 @@
 h = 
-8%0A        for i in range(1):%0A    
+self.samples_per_row%0A
     
@@ -858,17 +858,16 @@
 end(
-%5B
 self.x_v
 %5Bi*w
@@ -866,81 +866,10 @@
 .x_v
-%5Bi*width+width+j%5D for j in range(width)%5D)%0A        for i in range(1):%0A    
+)%0A
     
@@ -890,86 +890,14 @@
 end(
-%5B
 rx_v
-%5Bi*width+width+j%5D for j in range(width)%5D)%0A        for i in range(1):%0A    
+)%0A
     
@@ -918,52 +918,11 @@
 end(
-%5B
 g_v
-%5Bi*width+width+j%5D for j in range(width)%5D
 )%0A%0A 
@@ -997,39 +997,102 @@
 ack(
-%5Bnp.h
 stack
-(
 s)
- for s in stacks%5D)
+%0A        return %7B 'generator':images%7D%0A%0A%0A    def sample(self, path, save_samples=False):
 %0A%0A  

debug sampler includes extra sampling for began