breakhis_v2.py

# -*- coding: utf-8 -*-
"""Breakhis v2

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/17dqHsfx5kp81I9rVdmts3T1HUOjr9oZ2
"""

from google.colab import drive
drive.mount('/content/drive')

# Commented out IPython magic to ensure Python compatibility.
# %tensorflow_version 2.x
import tensorflow as tf
import tensorflow_addons as tfa
import numpy as np
import os
import matplotlib.pyplot as plt
from pickle import load, dump
import sys
import functools

"""# Load Data"""

# patient ids/class e.g. ids[0] : patients with Adenosis 
path = 'drive/My Drive/TFRecord_ds/'
ids = load(open(path+'ids.pkl','rb') )
zoom_ids = load(open(path+'zoom_ids.pkl','rb'))

raw_image_dataset = tf.data.TFRecordDataset(tf.data.Dataset.list_files("drive/My\ Drive/TFRecord_ds/img_*.tfrecord",shuffle=False))

image_feature_description = {
    'id' : tf.io.FixedLenFeature([], tf.string),
    'zoom' :  tf.io.FixedLenFeature([], tf.int64), 
    'label': tf.io.FixedLenFeature([], tf.int64),
    'image_raw': tf.io.FixedLenFeature([], tf.string),
}

def _parse_image_function(example_proto):
  # Parse the input tf.Example proto using the dictionary above.
  return tf.io.parse_single_example(example_proto, image_feature_description)

parsed_image_dataset = raw_image_dataset.map(_parse_image_function)

"""# Split Dataset"""

from random import shuffle, seed
from collections import defaultdict
seed(10)

other_paper = True
set_indices = []
paper_indices  = [(),(),(),(),(21,6,11),(),(2,1,2),()]

def filter_patient_fn(img_dict,id_set,i):
  return tf.math.logical_and(
                       tf.math.equal(img_dict['label'], tf.constant(i, dtype='int64')), 
                       tf.reduce_any(tf.convert_to_tensor(list(map(lambda id: tf.math.equal(img_dict['id'], id),id_set)))))

def split_dataset(ds, train_ratio, val_ratio, ids):
  splitted_ids = {k:[] for k in ['Val', 'Train', 'Test']}
  num_of_patients = defaultdict(int)
  # split every class patientwise
  for i,id_set in enumerate(ids):
    id_set = list(sorted(id_set))
    shuffle(id_set)
    ids_size = len(id_set)
    if (i!=4 and i!=6) or other_paper:
      test_idx = round(train_ratio*ids_size)
      # keep at least 1 patient for testing
      test_idx = test_idx-1 if test_idx==ids_size else test_idx
      val_idx  = round(val_ratio*test_idx)
      # keep at least 1 patient for validation
      val_idx = 1 if not val_idx else val_idx
    else :
      test_idx = ids_size - paper_indices[i][2]
      val_idx = paper_indices[i][1]

    
    splitted_ids['Val'].append(id_set[:val_idx])
    splitted_ids['Train'].append(id_set[val_idx:test_idx])
    splitted_ids['Test'].append(id_set[test_idx:])

    num_of_patients['Val']+=val_idx
    num_of_patients['Train']+=test_idx-val_idx
    num_of_patients['Test']+=ids_size-test_idx
    set_indices.append(( test_idx-val_idx,val_idx,ids_size-test_idx))

  dss = {k : tuple(ds.filter(lambda im: filter_patient_fn(im, set_ids, i))
               for i,set_ids in enumerate(ids_lists)) for k,ids_lists in splitted_ids.items() }
 
  
  ret = {k: functools.reduce(lambda a, b: a.concatenate(b), d).map(lambda img_dict : (img_dict['image_raw'], img_dict['label'])) for k, d in dss.items()}
  return ret,num_of_patients

"""## Split per zoom"""

train_zoom, val_zoom, test_zoom = {}, {}, {}
train_zoom_sz, val_zoom_sz, test_zoom_sz = {}, {}, {}

for k in zoom_ids.keys():
  ds_at_zoom = parsed_image_dataset.filter(lambda img : tf.math.equal(img['zoom'], tf.constant(k, dtype='int64')))
  splitted = split_dataset(ds_at_zoom, 0.9, 1/9, ids)
  dss = splitted[0]
  num_of_patients = splitted[1]
  train_zoom[k], val_zoom[k], test_zoom[k] = dss['Train'], dss['Val'], dss['Test']
  train_zoom_sz[k],val_zoom_sz[k],test_zoom_sz[k] =  num_of_patients['Train'], num_of_patients['Val'], num_of_patients['Test']

zoom = 40
train_parsed, val_parsed , test_parsed =  train_zoom[zoom], val_zoom[zoom], test_zoom[zoom]
train_patients, val_patients, test_patients = train_zoom_sz[zoom], val_zoom_sz[zoom], test_zoom_sz[zoom]

print(train_patients, val_patients, test_patients)

"""## Whole Dataset Split"""

splitted = split_dataset(parsed_image_dataset, 0.8, 0.1, ids)
train_parsed, val_parsed , test_parsed = splitted[0]['Train'], splitted[0]['Val'], splitted[0]['Test']
train_patients, val_patients , test_patients = splitted[1]['Train'], splitted[1]['Val'], splitted[1]['Test']

"""# Count"""

# Save time ^^
distribution = defaultdict(list)

#0.9, 1/9, zoom 40
distribution['Train'] =  [70, 204, 58, 121, 674, 161, 83, 91]
distribution['Val'] =  [15, 23, 13, 16, 92, 29, 51, 19]
distribution['Test'] =  [29, 26, 38, 12, 98, 15, 22, 35]

# 0.8, .15 whole ds
# distribution['Train'] = [262, 665, 235, 452, 2491, 554, 351, 353]
# distribution['Val'] = [61, 81, 60, 64, 234, 119, 201, 65]
# distribution['Test'] = [121, 268, 158, 53, 726, 119, 74, 142]

# attention by selection distr
# distribution['Train'] = [70, 123, 58, 105, 459, 73, 63, 61]
# distribution['Val'] = [15, 37, 13, 16, 167, 75, 51, 19]
# distribution['Test'] = [29, 93, 38, 28, 238, 57, 42, 65] 

# my distrib 
# distribution['Train'] = [70, 123, 58, 105, 436, 73, 83, 61]
# distribution['Val'] = [15, 37, 13, 16, 190, 75, 51, 19]
# distribution['Test'] = [29, 93, 38, 28, 238, 57, 22, 65]

distribution = defaultdict(list)
for distr, parsed in zip(['Train', 'Val', 'Test'], [train_parsed, val_parsed, test_parsed]):
  distribution[distr] = [0 for i in range(8)]

  for _,l in parsed.as_numpy_iterator():
    distribution[distr][l]+=1

  print(f"{distr}: {distribution[distr]}")

train_size = sum(distribution['Train'])
val_size = sum(distribution['Val'])              
test_size = sum(distribution['Test'])              

# size of train, validation and test set respectively
train_size, val_size, test_size

max_num_of_instances = max(distribution['Train'])
number_of_transformations = [ max_num_of_instances//v - 1 for v in distribution['Train']]

# number of transformations needed per class to improve balance
number_of_transformations

"""# Preprocess"""

# run for no augmentation
number_of_transformations = [2 for i in range(8)]

AUTOTUNE = tf.data.experimental.AUTOTUNE # https://www.tensorflow.org/guide/data_performance
batch_size = 32
image_size = 224

from random import randint, random, choice

# def rotate_image(image, deg):
#   image  = tfa.image.rotate(image,tf.constant(deg*1.))
#   return image
  
# def random_rotate(x):  
#   image  =  tfa.image.rotate(x,tf.constant(randint(0, 359)*1.))
#   return image
def apply_rnd_transformation(x, transformations):
  f = choice(transformations)
  return f[0](*((x,)+f[1:])) if random()<.5 else x

def augment_data(ds,transformations):
  augmented = [ds.map(lambda x,l: ( f[0](*((x,)+f[1:])), l) ) for f in transformations  ]
  ret = functools.reduce(lambda a, b: a.concatenate(b), [ds]+augmented)
  return ret
# for a set of transformation functions f1,f2,..,fn return fn(...(f2(f1(x))))
def fold_transformations(ds, transformations):
  return ds.map(lambda x,l: ((functools.reduce(lambda i,f: f[0](*((i,)+f[1:])), transformations, x)), l))

def apply_rnd_transformations(ds, transformations):
  return ds.map(lambda x,l: ( apply_rnd_transformation(x,transformations), l))

def mean_stdev(ds, size):
  means = np.zeros((3))
  sum_tmp = np.zeros((3))
  for im,l in ds.as_numpy_iterator():
    for i in range(3):
      means[i]+=tf.reduce_mean(im[:,:,i])
  means = (batch_size * means)/(size)

  for im,l in ds.as_numpy_iterator():
    for i in range(3):
      sum_tmp[i]+=(tf.reduce_sum((im[:,:,i] - means[i])**2))

  stdevs = np.sqrt(sum_tmp / (size * 224 *224 ))
  means = np.repeat([np.repeat([means], 224, axis=0)], 224,axis=0)
  stdevs = np.repeat([np.repeat([stdevs], 224, axis=0)], 224,axis=0)
  return means, stdevs

from operator import truediv

def _input_fn(ds): 
  return ds.repeat().batch(batch_size).prefetch(buffer_size=AUTOTUNE)
  
preprocess = [(tf.io.decode_jpeg, 3 ), (tf.image.resize, [image_size, image_size]),
              # (tf.image.resize_with_crop_or_pad, 460, 700), # not all images have exactly the same dimensions
              ( tf.cast, tf.float32),
              (tf.image.per_image_standardization, )
              #  (truediv, 255.)
              ]
val ,test = [ fold_transformations(ds, preprocess).cache()
                      for ds in [val_parsed , test_parsed ]
                    ]
train_per_class = [ train_parsed.filter(lambda _,l : tf.math.equal(l, tf.constant(i, 'int64'))) for i in range(8)]
train_per_class = [  fold_transformations(ds, preprocess).cache() for ds in train_per_class] 

aug_size=train_size
if any(number_of_transformations):
  transformations = [
                    (tf.image.flip_up_down,  ), (tf.image.flip_left_right,) , (tf.image.rot90, 1) ,  (tf.image.rot90, 3), 
                      (tf.image.rot90, 2) ]#,(tf.image.random_brightness, 0.3),  (tf.roll, [round(image_size*.2) for _ in range(3)], [0,1,2]) , 
                      # (tf.image.random_contrast,.2, .5), (tf.image.adjust_saturation, .5) ]

  rnd_transformations = [(tf.image.random_flip_up_down, 42 ), (tf.image.random_flip_left_right, 42)  ] #,(tf.image.random_brightness, 0.3, 42),  (tf.image.random_contrast,.2, .5, 42)]

  ds_size =  min( d*(min(t, len(transformations))+1) for d,t in zip(distribution['Train'], number_of_transformations))

  aug_train_per_class = [ augment_data(ds, transformations[:min(n, len(transformations))])  for ds,n in zip(train_per_class, number_of_transformations)]

  # aug_train_per_class[4] = apply_rnd_transformations(aug_train_per_class[4], transformations) # apply random transformations to majority class
  aug_distribution = [d*(min(t, len(transformations))+1) for d,t in zip(distribution['Train'], number_of_transformations)]
  aug_size = sum(aug_distribution)

# train set too large to fit in shuffle buffer, use sampling  to randomly shuffle 
train_concatenated = tf.data.experimental.sample_from_datasets(train_per_class, [.125]*8, seed=10) 

# train_concatenated = tf.data.experimental.sample_from_datasets(train_per_class, \
                                                              #  [(train_size - d)/(train_size*(train_size-1)) for d in distribution['Train']],
                                                              #  seed=10)
# GCN
# means_stdevs = [ mean_stdev(functools.reduce(lambda a, b: a.concatenate(b), [train_concatenated, val, test]), sum([aug_size, val_size, test_size]))]*3
# sets  = [ st.map(lambda im,l: ((im-mu)/sg,l)) for st,(mu,sg) in zip([train_concatenated, val, test], means_stdevs)]
# train_concatenated, val, test = sets

train_ds = _input_fn(train_concatenated.shuffle(aug_size, seed=10))
val_ds = _input_fn(val.shuffle(val_size, seed=10))
test_ds = _input_fn(test.shuffle(test_size, seed=10))

steps = aug_size//batch_size
# class_weight = {k : (1/distr)*(aug_size/8.) for k,distr in enumerate(aug_distribution) } # argument to pass as class weight in fit()
# class_weight = {k : (1/ (distr*7))*(aug_size)/8. for k,(distr, tr) in enumerate(zip(distribution['Train'], number_of_transformations)) }
print(aug_size)
print(steps)

val_ds = _input_fn(val.concatenate(test).shuffle(val_size+test_size, seed=10))

"""# Model

## Resnet18
"""

!pip install torch torchvision
!pip install onnx onnxruntime
!pip install fire
!git clone https://github.com/AxisCommunications/onnx-to-keras.git

import torchvision.models as models
import torch

resnet18 = models.resnet18(pretrained=True)
dummy_input = torch.randn(32, 3, 224, 224)
torch.onnx.export(resnet18, dummy_input, "resnet18.onnx")

!python3 onnx-to-keras/onnx2keras.py resnet18.onnx resnet18.h5

import re
from keras.models import Model

def insert_layer_nonseq(model, layer_regex, insert_layer_factory,
                        insert_layer_name=None, position='after'):

    # Auxiliary dictionary to describe the network graph
    network_dict = {'input_layers_of': {}, 'new_output_tensor_of': {}}

    # Set the input layers of each layer
    for layer in model.layers:
        for node in layer._outbound_nodes:
            layer_name = node.outbound_layer.name
            if layer_name not in network_dict['input_layers_of']:
                network_dict['input_layers_of'].update(
                        {layer_name: [layer.name]})
            else:
                # network_dict['input_layers_of'][layer_name].append(layer.name)

                if (layer.name not in network_dict['input_layers_of'][layer_name]):
                  network_dict['input_layers_of'][layer_name].append(layer.name)

    # Set the output tensor of the input layer
    network_dict['new_output_tensor_of'].update(
            {model.layers[0]._name: model.input})

    # Iterate over all layers after the input
    model_outputs = []
    for layer in model.layers[1:]:

        # Determine input tensors
        layer_input = [network_dict['new_output_tensor_of'][layer_aux] 
                for layer_aux in network_dict['input_layers_of'][layer.name]]
        if len(layer_input) == 1:
            layer_input = layer_input[0]

        # Insert layer if name matches the regular expression
        if re.match(layer_regex, layer.name):
            if position == 'replace':
                x = layer_input
            elif position == 'after':
                x = layer(layer_input)
            elif position == 'before':
                pass
            else:
                raise ValueError('position must be: before, after or replace')

            new_layer = insert_layer_factory()
            if insert_layer_name:
                new_layer.name = insert_layer_name
            else:
                new_layer._name = '{}_{}'.format(layer.name, new_layer.name)
            x = new_layer(x)
            print('New layer: {} Old layer: {} Type: {}'.format(new_layer.name,
                                                            layer.name, position))
            if position == 'before':
                x = layer(x)
        else:
            x = layer(layer_input)

        # Set new output tensor (the original one, or the one of the inserted
        # layer)
        network_dict['new_output_tensor_of'].update({layer._name: x})

        # Save tensor in output list if it is output in initial model
        if layer_name in model.output_names:
            model_outputs.append(x)

    return Model(inputs=model.inputs, outputs=model_outputs)

from tensorflow.keras.layers import *
tf.keras.backend.clear_session()
mod = tf.keras.models.load_model('resnet18.h5')

t_model = tf.keras.Model(inputs=mod.input, outputs=mod.get_layer(index=-4).output)
t_model.trainable=False
for layer in t_model.layers[56:]:
  layer.trainable = True
# model.summary()

x = t_model.output    
# x = Dropout(.5)(x)
x = GlobalAveragePooling2D()(x)
x = Dropout(.5)(x)
x = Dense(128)(x)
# x = Dropout(.5)(x)
x = Dense(8)(x)
out = Softmax()(x)
hope_model = tf.keras.Model(inputs=t_model.input, outputs=out)
conf = defaultdict()
for layer in hope_model.layers[56:]:
  conf = layer.get_config()
  if 'kernel_regularizer' in conf:
    layer.kernel_regularizer = tf.keras.regularizers.l2(0.001)
  
# hope_model.summary()

def dropout_layer_factory():
    return Dropout(rate=0.5)
drop = insert_layer_nonseq(hope_model, '.*conv2d.*', dropout_layer_factory)

tmp = Model(inputs=drop.input, outputs=drop.output[-1])

tmp.compile(optimizer=tf.keras.optimizers.SGD(learning_rate=0.001, momentum=0.9), loss='sparse_categorical_crossentropy', metrics='acc')

history = tmp.fit(train_ds, epochs=5, shuffle=False, steps_per_epoch=steps,validation_data=val_ds, validation_steps=(val_size+test_size)//batch_size )

history = hope_model.fit(train_ds, epochs=20, initial_epoch=5, steps_per_epoch=steps,validation_data=val_ds, validation_steps=(val_size+test_size)//batch_size )

history = hope_model.fit(train_ds, epochs=40, initial_epoch=20,steps_per_epoch=steps,validation_data=val_ds, validation_steps=(val_size+test_size)//batch_size )

"""## Any pretrained"""

def model(pretrained, size, no_fine_tune=True, freeze=1 ):
      PRETRAINED_MODEL=pretrained(input_shape=(size, size,3), include_top=False, weights='imagenet')

      PRETRAINED_MODEL.trainable = True
      if no_fine_tune:
        for layer in PRETRAINED_MODEL.layers[:round(len(PRETRAINED_MODEL.layers)*freeze)]:
          layer.trainable =  False
      dropout_layer1 = tf.keras.layers.Dropout(0.5)

      dense_layer2 = tf.keras.layers.Dense(512,activation='relu')
      dropout_layer3 = tf.keras.layers.Dropout(0.5)

      dense_layer1 = tf.keras.layers.Dense(256,kernel_regularizer=tf.keras.regularizers.l1(0.0001),activation='relu')
      dropout_layer2 = tf.keras.layers.Dropout(0.5)

      dense_layer3 = tf.keras.layers.Dense(128,activation='relu')
      dropout_layer4 = tf.keras.layers.Dropout(0.5)

      dense_layer4 = tf.keras.layers.Dense(16,kernel_regularizer=tf.keras.regularizers.l2(),activation='relu')

      global_max_layer = tf.keras.layers.GlobalMaxPooling2D()
      global_avg_layer = tf.keras.layers.GlobalAveragePooling2D()
      
      batch_norm = tf.keras.layers.BatchNormalization()
      batch_norm1 = tf.keras.layers.BatchNormalization()

      flatten = tf.keras.layers.Flatten()
      
      # add top layer  classification
      prediction_layer = tf.keras.layers.Dense(8,activation='softmax')
      # prediction_layer = tf.keras.layers.Dense(4, activation='linear', kernel_regularizer=tf.keras.regularizers.l2())
      dropout_layer2 = tf.keras.layers.Dropout(0.5)
      model = tf.keras.Sequential([PRETRAINED_MODEL,global_avg_layer, dropout_layer1, batch_norm, dense_layer1, batch_norm1,dropout_layer2,dense_layer3 ,prediction_layer ])
      # model = tf.keras.Sequential([PRETRAINED_MODEL,global_avg_layer, dropout_layer1,batch_norm,dense_layer1,dropout_layer2,batch_norm1,dense_layer3,dropout_layer4,prediction_layer ])

      lr = tf.keras.optimizers.schedules.ExponentialDecay(initial_learning_rate=0.0002, decay_rate=0.7, decay_steps=2)
      model.compile(optimizer=tf.optimizers.RMSprop(learning_rate=0.00001, decay=.000001), loss=tf.keras.losses.sparse_categorical_crossentropy , metrics=["accuracy"])
      
      model.summary()
      return model

the_model = model(tf.keras.applications.VGG16,image_size, freeze=.5) 
learn_control = tf.keras.callbacks.ReduceLROnPlateau(monitor='val_accuracy', patience=5,
                                  verbose=1,factor=0.2, min_lr=1e-7)

history = the_model.fit(train_ds, epochs=10,steps_per_epoch=steps, validation_data=val_ds, validation_steps=val_size// batch_size,  class_weight=class_weight)

history = the_model.fit(train_ds, epochs=20,initial_epoch=10,steps_per_epoch=steps, validation_data=val_ds, validation_steps=val_size// batch_size,  class_weight=class_weight)

"""# Attention Model"""

from tensorflow.keras import Model
from tensorflow.keras.layers import *

import tensorflow as tf
import numpy as np

"""## Metrics"""

class ImageLevelAccuracy(tf.keras.metrics.Metric):

  def __init__(self, name='image_level_accuracy', **kwargs):
    super(ImageLevelAccuracy, self).__init__(name=name, **kwargs)
    self.n_rec= self.add_weight(name='n_rec', initializer='zeros')
    self.cnt = self.add_weight(name='count', initializer='zeros')
    self.majority = 3*tf.ones(4, dtype='int64')
    

  def update_state(self, y_true, y_pred, sample_weight=None):
    y_pred = tf.reshape(tf.argmax(y_pred, axis=1), [4,5])
    y_true = tf.reshape(y_true,[4,5])
    v = tf.equal(y_true, y_pred)
    # count correctly classified patches
    v = tf.math.reduce_sum(tf.cast(v, dtype='int64'), axis=1)  
    # count correctly classified images TODO: if no majority in training should we select more patches?
    v = tf.reduce_sum(tf.cast(tf.greater_equal(v,self.majority), 'int64'))
    self.n_rec.assign_add(tf.cast(v,'float32'))
    self.cnt.assign_add(tf.cast(4, dtype='float32'))

  def result(self):
    return self.n_rec/self.cnt
 
  def reset_states(self):
    self.n_rec.assign(0)
    self.cnt.assign(0)

my_acc = ImageLevelAccuracy()

"""## Soft Attention"""

def image_cross_entropy_loss(y_true, y_pred):
  class_prob = tf.reduce_sum(tf.one_hot(tf.squeeze(y_true), depth=8)*y_pred, axis=1)
  y_image = tf.reduce_mean(tf.reshape(class_prob, (4,5)), axis=1)
  return -tf.reduce_mean(tf.math.log(y_image))

def residual_unit(inp, patch_size=112, is_first=False):
  # A
  a = Conv2D(input_shape=(patch_size, patch_size, 3),filters=64, kernel_size=1)(inp) if is_first else Conv2D(filters=64, kernel_size=1)(inp)
  a = BatchNormalization( momentum=.1, epsilon=1e-5)(a) # PyTorch different default values

  # B
  # b = Conv2D(input_shape=(patch_size+2, patch_size+2, 3),filters=64, kernel_size=1)(inp) if is_first else Conv2D(filters=64, kernel_size=1)(inp) # when zero padding at beginning
  b =  Conv2D(filters=64, kernel_size=1)(inp)

  b = BatchNormalization(momentum=.1, epsilon=1e-5)(b) # PyTorch different default values
  b = PReLU(shared_axes=[1,2], alpha_initializer=tf.keras.initializers.constant(.25))(b)
  b = Conv2D(filters=64, kernel_size=3)(b)
  b = BatchNormalization(momentum=.1, epsilon=1e-5)(b) # PyTorch different default values
  b = PReLU(shared_axes=[1,2], alpha_initializer=tf.keras.initializers.constant(.25))(b)
  b = Conv2D(filters=64, kernel_size=1)(b)
  b = BatchNormalization(momentum=.1, epsilon=1e-5)(b)

  # a = patch_size X patch_size, b = (patch_size-2) X (patch_size-2) 
  b = ZeroPadding2D()(b)
  

  output = Add()([a,b])
  
  return output

patch = Input(shape=( 112,112,3),name='input' )

# SaNet
# let patch be the patch
ru = residual_unit(patch, True)

# Trunk 
t = residual_unit(ru)
t = residual_unit(t)

# Mask
m = MaxPool2D(pool_size=3, strides=2)(ru)
m = residual_unit(m)
m = residual_unit(m)

m = MaxPool2D(pool_size=3, strides=2)(m)
m = residual_unit(m)
m = residual_unit(m)

m = UpSampling2D(interpolation='bilinear')(m)
m = BatchNormalization(momentum=.1, epsilon=1e-5)(m)
m = PReLU(shared_axes=[1,2], alpha_initializer=tf.keras.initializers.constant(.25))(m)
m = residual_unit(m)

m = ZeroPadding2D()(m)
m = UpSampling2D(interpolation='bilinear')(m)
m = BatchNormalization(momentum=.1, epsilon=1e-5)(m)
m = PReLU(shared_axes=[1,2], alpha_initializer=tf.keras.initializers.constant(.25))(m)
m = residual_unit(m)

m = Conv2D(filters=64, kernel_size=1, activation='sigmoid')(m)

x = Multiply()([m+1, t])
# x = AveragePooling2D(pool_size=5, strides=1)(x)
x = GlobalAveragePooling2D()(x)
# x = Flatten()(x)
deep = Dense(128, activation='relu', name='feature_layer')(x)
output = Dense(8)(deep)
output = Softmax()(output)

lr = tf.keras.optimizers.schedules.ExponentialDecay(initial_learning_rate=.01, decay_steps=steps*200, decay_rate=.95)

SA_Net = Model(inputs=patch, outputs=output)
SA_Net.compile(optimizer=tf.keras.optimizers.Adam(), loss=image_cross_entropy_loss, metrics=[my_acc])
# SA_Net.summary()

SA_Net_features = Model(inputs=SA_Net.input, outputs=SA_Net.get_layer('feature_layer').output)

np.random.seed(12)

train_iter = train_ds.as_numpy_iterator()
val_iter = val_ds.as_numpy_iterator()

"""### Testing area"""

batch = val_iter.next()
# positions = np.hstack([(np.random.randint(66,394, size=(4,5))/460.) , (np.random.randint(66,634, size=(4,5))/700.)]).reshape(4,5,2)
# shallow = np.zeros((4,5,7))
# shallows=[]
# deeps=[]
# avg_loss = 0.
# avg_acc = 0.

# guard=0
# cropped = [0]*4
# l = []
# b = []

# show_batch(batch, positions)
# exctract 5 patches for each image in mini batch
glmps = [tf.image.extract_glimpse(batch[0], size=(112,112), offsets=positions[:, i, :], centered=True, normalized=True) for i in range(5)] 
glmps = tf.concat(glmps, 0)
labels = np.vstack([batch[1]]*5 ).T
# 6 loc_x, 7 loc_y
# shallow[:,:,5] = positions[:,:,0]
# shallow[:,:,6] = positions[:,:,1]
# # 3 pred, 4 true pred
# shallow[:,:,3] = np.argmax(SA_Net.predict_on_batch(glmps), axis=1).reshape(4,5)
# shallow[:,:,4] = labels


# shallow_t = tf.convert_to_tensor(shallow, dtype=tf.float64)
# deep = tf.reshape(SA_Net_features(glmps), (4,5,128))

# if not val:
#   deeps.append(deep)
#   shallows.append(shallow_t)

# out = [o.numpy() for o in De_Net([deep,shallow_t ]) ] 

# patches = glmps.numpy().reshape(4,5, 112,112,3)
# # print(positions[0])
# # iterate over DeNet actions, select patches and update centers for next iter
# for i,im in enumerate(out[1]):
#   for j,o in enumerate(im):
#     if (out[0][i][j][0]>=out[0][i][j][1] and cropped[i] <5) or (guard>15 and cropped[i]<5):
#       b.append(patches[i][j])
#       l.append(labels[i][j])
#       cropped[i]+=1
#     # avoid croping the edges
#     y = np.clip(o[0]+np.random.normal(scale=20/460), 0,1) *164/230
#     x = np.clip(o[1]+np.random.normal(scale=20/700), 0, 1)*294/350

#     positions[i,j,:] = [y , x ]
# # out = SA_Net.train_on_batch(glmps.numpy, labels)

step = 0
historical_loss = 0
len(l)

loss, acc =  SA_Net.train_on_batch(np.array(b[:20]), np.array(l[:20]))
      
avg_loss = ( avg_loss*step + loss) / (step + 1) # running average
avg_acc = (avg_acc*step + acc) / (step + 1)


# (historical?) training loss, state feauture 0
historical_loss =  ( historical_loss*(step + shallow[0][0][2] ) + loss) / (step + shallow[0][0][2] + 1)
shallow[:,:,0] = np.array([historical_loss]*20).reshape(4,5)
shallow[:,:,2] = np.array([shallow[0][0][2]+1]*20).reshape(4,5) # num of iterations (epochs?)

l , acc = SA_Net.train_on_batch(glmps.numpy(), labels.reshape(20), )

l, acc

features = SA_Net.evaluate(glmps.numpy(), labels.reshape(20), batch_size=20, steps=1)
features

tmp = tf.reduce_sum(tf.one_hot(tf.convert_to_tensor(labels.reshape(20), dtype='int64'),8)*pred, axis=-1) #, tf.math.reduce_max(pred, 1)
tf.reshape(tmp, [4,5])
# -tf.reduce_mean(tf.math.log(tmp))

outputs[0]

SA_Net.get_layer(index=4).get_weights()

tf.reduce_max(SA_Net_features(glmps), axis=-1)

my_acc.update_state(tf.convert_to_tensor(labels.reshape(20), dtype='int64'), tf.convert_to_tensor(pred, 'float64'))
my_acc.result()

pred = SA_Net.predict(glmps.numpy() )
tf.argmax(pred, axis=1).numpy().reshape(4,5), labels

tf.reduce_mean(tf.keras.backend.sparse_categorical_crossentropy(labels.reshape(20),pred))

positions

def show_batch(batch, positions):
  bboxes = np.apply_along_axis( lambda y : [(y[0]+1)/2-66/460, (y[1]+1)/2-66/700, (y[0]+1)/2+66/460, (y[1]+1)/2+66/700], axis=-1, arr=positions )
  # bboxes = np.apply_along_axis( lambda y : [y[0]-66/460, y[1]-66/700, y[0]+66/460, y[1]+66/700], axis=-1, arr=positions )
  colors = np.array([[0.2, 0.0, 1.0]])
  to_show = tf.image.draw_bounding_boxes(batch[0], bboxes, colors)
  f, axarr = plt.subplots(2,2, figsize=(20,10))
  for i, im in enumerate(tf.unstack(to_show)):
    axarr[i//2,i % 2].imshow(im)

# show_batch(batch, positions)

"""## Hard Attention"""

# DeNet
deep = Input(shape=(5,128), batch_size=4, name='deep_input')
other_state = Input(shape=(5, 7), batch_size=4,name='other_input' )

norm_other_state = Lambda(lambda x:tf.keras.backend.l2_normalize(x,axis=2))(other_state)
norm_deep = Lambda(lambda x: tf.keras.backend.l2_normalize(x, axis=2))(deep)

fc_relu = Dense(32, activation='relu',kernel_initializer=tf.keras.initializers.RandomUniform(minval=-0.01, maxval=0.01))(deep)
fc_tanh = Dense(10, activation='tanh',kernel_initializer=tf.keras.initializers.RandomUniform(minval=-0.01, maxval=0.01))(other_state)
fc = Concatenate()([fc_tanh, fc_relu])
lstm_input = Dense(24,activation='relu',kernel_initializer=tf.keras.initializers.RandomUniform(minval=-0.01, maxval=0.01))(fc)
lstm_output = LSTM(24, return_sequences=True)(lstm_input)
action = Dense(2, activation='sigmoid',kernel_initializer=tf.keras.initializers.RandomUniform(minval=-0.01, maxval=0.01), \
               bias_initializer=tf.keras.initializers.Constant(2))(lstm_output)
loc_y = Dense(2, activation='sigmoid', kernel_initializer=tf.keras.initializers.RandomUniform(minval=-0.01, maxval=0.01))(lstm_output)
loc_x = Dense(2, activation='sigmoid', kernel_initializer=tf.keras.initializers.RandomUniform(minval=-0.01, maxval=0.01))(lstm_output)

# crop = LSTM(2)(lstm_input)
# loc_y = LSTM(236)(lstm_input)
# loc_x = LSTM(476)(lstm_input)

De_Net = Model(inputs=[deep, other_state], outputs = [action, loc_y, loc_x] )
# De_Net = Model(inputs=[deep, other_state], outputs = [action, loc] )
# De_Net = Model(inputs=[deep, other_state], outputs = [lstm_output] )


De_Net.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001) )
# De_Net.summary()

"""## Experiments

### B2B 1 val 1 train
"""

def step_fn(batch, step, avg_loss=0., avg_acc=0., val=False):
  global positions, shallow

  
  cropped = 0
  guard = 0 # prevent infinite loop while cropping

  b = np.empty((4,5,112,112,3)) # batch for SaNet
  l = np.zeros((4,5), dtype='int')-1 # labels of batch for SaNet
  
  ds_per_ep = []
  shs_per_ep = []
  actions_per_ep = []
  actions_ts = np.zeros((3,4,5))
  
  while cropped<20: 
    guard+=1
    # exctract 5 patches for each image in mini batch
    glmps = [tf.image.extract_glimpse(batch[0], size=(112,112), offsets=positions[:, i, :], centered=True, normalized=True) for i in range(5)] 
    glmps = tf.concat(glmps, 0)
    labels = np.vstack([batch[1] for _ in range(5)]).T

    # 6 loc_x, 7 loc_y
    shallow[:,:,5] = positions[:,:,0]
    shallow[:,:,6] = positions[:,:,1]
    # 3 pred, 4 true pred
    shallow[:,:,3] = np.argmax(SA_Net.predict_on_batch(glmps), axis=1).reshape(4,5)
    shallow[:,:,4] = labels


    shallow_t = tf.convert_to_tensor(shallow, dtype=tf.float64)
    deep = tf.reshape(SA_Net_features(glmps), (4,5,128))

    if not val:
        ds_per_ep.append(deep)
        shs_per_ep.append(shallow_t)

    out = [o.numpy() for o in De_Net([deep,shallow_t ]) ] 

    patches = glmps.numpy().reshape(4,5, 112,112,3)
    # iterate over DeNet actions, select patches and update centers for next iter
    for i,im in enumerate(out[1]):
      for j,o in enumerate(im):
        actions_ts[0][i][j] = 0
        if l[i][j]==-1 and (out[0][i][j][1]-out[0][i][j][0]>=0. or guard>10 ):
          b[i][j] = patches[i][j]
          l[i][j] = labels[i][j]
          cropped+=1
          actions_ts[0][i][j] = 1
        # avoid croping the edges
        y = np.clip(tf.random.normal([1], mean=o[0], stddev=o[1], dtype=tf.float32), 66/230 -1, 394/230-1) #*164/230
        x = np.clip(tf.random.normal([1], mean=out[2][i][j][0], stddev=out[2][i][j][1], dtype=tf.float32), 66/350-1,634/350-1) #*294/350

        positions[i,j,:] = [y , x ]
        actions_ts[1][i][j] = y
        actions_ts[2][i][j] = x
    if not val:
      actions_per_ep.append(actions_ts)
  
  loss, acc = SA_Net.train_on_batch(b.reshape((20,112,112,3)), l.reshape((20,1)) ) if not val else SA_Net.test_on_batch(b.reshape((20,112,112,3)), l.reshape((20,1)) ) 
    
  avg_loss = ( avg_loss*step + loss) / (step + 1) # running average
  avg_acc = (avg_acc*step + acc) / (step + 1)
      
  if not val:
    return avg_loss, avg_acc, actions_per_ep, ds_per_ep, shs_per_ep, loss
  else:
    return avg_loss, avg_acc, guard


def train_epoch( historical_loss=0.):
  global positions, shallow
  avg_loss = 0.
  avg_acc = 0.
  
  deeps = []
  shallows = []
  actions = []
  rews = []
  for step in range(train_size//4):
    batch = train_iter.next()
    avg_loss, avg_acc, actions_per_ep, ds_per_ep, shs_per_ep, loss = step_fn(batch, step, avg_loss, avg_acc)
    
    sys.stdout.write(f"\r {len(actions_per_ep)} Step {step+1}: Loss {avg_loss} , Accuracy {avg_acc}")
    sys.stdout.flush()

    deeps.append(ds_per_ep)
    shallows.append(shs_per_ep)
    actions.append(actions_per_ep)

    # (historical?) training loss, state feauture 0
    historical_loss =  ( historical_loss*(step + shallow[0][0][2] ) + loss) / (step + shallow[0][0][2] + 1)
    shallow[:,:,0] = np.array([historical_loss]*20).reshape(4,5)
    shallow[:,:,2] = np.array([shallow[0][0][2]+1]*20).reshape(4,5) # num of iterations (epochs?)

    rp=0.
    val_loss = 0.
    for val_step in range(5):
      val_batch = val_iter.next()
      val_loss,rp, _ = step_fn(val_batch,val_step,val_loss, rp, val=True)

    rc  = -np.log(step/253 + 1e-15) if val_loss <= .25 else 0. 
    if rc : print('Found it!')   
    rews.append(rp+rc)
  
  print()
  val_acc=0.
  val_loss = 0.
  for step in range(val_size//4):
    val_batch = val_iter.next()
    val_loss,val_acc, guard = step_fn(val_batch,step, val_loss, val_acc,True)  
    
    sys.stdout.write(f"\r {guard} Step {step+1}:                          Val Loss {val_loss} , Val Accuracy {val_acc}")
    sys.stdout.flush()
 # update at the end of validation or during?
  shallow[:,:,1] = np.array([max(val_acc, shallow[0][0][1])]*20).reshape(4,5) # best val acc
    
  return deeps, shallows,actions, rews,historical_loss

# Commented out IPython magic to ensure Python compatibility.
# %%time
# positions = np.hstack([(np.random.randint(66,394, size=(4,5)))/230. - 1, np.random.randint(66,634, size=(4,5))/350. -1]).reshape(4,5,2)
# # positions = np.hstack([(np.random.randint(66,394, size=(4,5)))/460. , (np.random.randint(66,634, size=(4,5)))/700. ]).reshape(4,5,2)
# shallow = np.zeros((4,5,7)) # all other but deep feautures of DeNet input
# hist_rew = 0.
# 
# for epoch in range(5):
#   print(f"Epoch {epoch+1}")
#   deeps,shallows, actions, rews , hist_loss= train_epoch(historical_loss=0. if not epoch else hist_loss)
#   
#   print()
#   with tf.GradientTape() as policy_tape:
#     sums = []
#     for ep, (ds, os) in enumerate(zip(deeps,shallows)):
#       outputs = [De_Net([d,o]) for d,o in zip(ds,os)]
#       for ts,o in enumerate(outputs):
#         select = tf.reduce_sum(tf.one_hot(actions[ep][ts][0],depth=2)*o[0], axis=-1)
#         pdf_y = tf.exp(-0.5 *((actions[ep][ts][1] - o[1][:,:,0]) / (o[1][:,:,1]))**2) * 1/(o[1][:,:,1]*tf.sqrt(2 *np.pi))
#         pdf_x = tf.exp(-0.5 *((actions[ep][ts][2] - o[2][:,:,0]) / (o[2][:,:,1]))**2) * 1/(o[2][:,:,1]*tf.sqrt(2 *np.pi))
#         outputs[ts] = tf.stack([select, pdf_y,pdf_x], axis=-1  )
#       
#       log_probs = tf.keras.backend.log(outputs)
#       sums.append(tf.reduce_sum(log_probs, axis=[0]))
#     sums = tf.stack(sums, axis=0)   
#       
#     policy_loss = -tf.reduce_mean(sums*(np.array([[[[r]*3]*5]*4 for r in rews])-hist_rew), axis=0)
#     policy_variables = De_Net.trainable_variables
#     grads = policy_tape.gradient(policy_loss, policy_variables)
#     De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
#   
#   hist_rew = (hist_rew*epoch+np.mean(rews))/(epoch+1)
# 
# show_batch(batch, positions)

for epoch in range(5,10):
  print(f"Epoch {epoch+1}")
  deeps,shallows, actions, rews , hist_loss= train_epoch(historical_loss=0. if not epoch else hist_loss)
  
  print()
  with tf.GradientTape() as policy_tape:
    sums = []
    for ep, (ds, os) in enumerate(zip(deeps,shallows)):
      outputs = [De_Net([d,o]) for d,o in zip(ds,os)]
      for ts,o in enumerate(outputs):
        select = tf.reduce_sum(tf.one_hot(actions[ep][ts][0],depth=2)*o[0], axis=-1)
        pdf_y = tf.exp(-0.5 *((actions[ep][ts][1] - o[1][:,:,0]) / (o[1][:,:,1]))**2) * 1/(o[1][:,:,1]*tf.sqrt(2 *np.pi))
        pdf_x = tf.exp(-0.5 *((actions[ep][ts][2] - o[2][:,:,0]) / (o[2][:,:,1]))**2) * 1/(o[2][:,:,1]*tf.sqrt(2 *np.pi))
        outputs[ts] = tf.stack([select, pdf_y,pdf_x], axis=-1  )
      
      log_probs = tf.keras.backend.log(outputs)
      sums.append(tf.reduce_sum(log_probs, axis=[0,1,2]))
    sums = tf.stack(sums, axis=0)   
      
    policy_loss = -tf.reduce_mean(sums*(np.array([[r] for r in rews])-hist_rew), axis=0)
    policy_variables = De_Net.trainable_variables
    grads = policy_tape.gradient(policy_loss, policy_variables)
    De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  
  hist_rew = (hist_rew*epoch+np.mean(rews))/(epoch+1)

show_batch(batch, positions)

for epoch in range(10,20):
  print(f"Epoch {epoch+1}")
  deeps,shallows, actions, rews , hist_loss= train_epoch(historical_loss=0. if not epoch else hist_loss)
  
  print()
  with tf.GradientTape() as policy_tape:
    sums = []
    for ep, (ds, os) in enumerate(zip(deeps,shallows)):
      outputs = [De_Net([d,o]) for d,o in zip(ds,os)]
      for ts,o in enumerate(outputs):
        select = tf.reduce_sum(tf.one_hot(actions[ep][ts][0],depth=2)*o[0], axis=-1)
        pdf_y = tf.exp(-0.5 *((actions[ep][ts][1] - o[1][:,:,0]) / (o[1][:,:,1]))**2) * 1/(o[1][:,:,1]*tf.sqrt(2 *np.pi))
        pdf_x = tf.exp(-0.5 *((actions[ep][ts][2] - o[2][:,:,0]) / (o[2][:,:,1]))**2) * 1/(o[2][:,:,1]*tf.sqrt(2 *np.pi))
        outputs[ts] = tf.stack([select, pdf_y,pdf_x], axis=-1  )
      
      log_probs = tf.keras.backend.log(outputs)
      sums.append(tf.reduce_sum(log_probs, axis=[0,1,2]))
    sums = tf.stack(sums, axis=0)   
      
    policy_loss = -tf.reduce_mean(sums*(np.array([[r] for r in rews])-hist_rew), axis=0)
    policy_variables = De_Net.trainable_variables
    grads = policy_tape.gradient(policy_loss, policy_variables)
    De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  
  hist_rew = (hist_rew*epoch+np.mean(rews))/(epoch+1)

show_batch(batch, positions)

"""### Experiment 6/10 
* 1 train 5 val
* Zeorpadding before multiply in SaNet
* normalized, .125
* Default adam
* loc= continuous actions
"""

# Commented out IPython magic to ensure Python compatibility.
# %%time
# positions = np.hstack([(np.random.randint(66,394, size=(4,5)))/230. - 1, np.random.randint(66,634, size=(4,5))/350. -1]).reshape(4,5,2)
# # positions = np.hstack([(np.random.randint(66,394, size=(4,5)))/460. , (np.random.randint(66,634, size=(4,5)))/700. ]).reshape(4,5,2)
# shallow = np.zeros((4,5,7)) # all other but deep feautures of DeNet input
# hist_rew = 0.
# 
# for epoch in range(5):
#   print(f"Epoch {epoch+1}")
#   deeps,shallows, actions, rews , hist_loss= train_epoch(historical_loss=0. if not epoch else hist_loss)
#   
#   print()
#   with tf.GradientTape() as policy_tape:
#     sums = []
#     for ep, (ds, os) in enumerate(zip(deeps,shallows)):
#       outputs = [De_Net([d,o]) for d,o in zip(ds,os)]
#       for ts,o in enumerate(outputs):
#         select = tf.reduce_sum(tf.one_hot(actions[ep][ts][0],depth=2)*o[0], axis=-1)
#         pdf_y = tf.exp(-0.5 *((actions[ep][ts][1] - o[1][:,:,0]) / (o[1][:,:,1]))**2) * 1/(o[1][:,:,1]*tf.sqrt(2 *np.pi))
#         pdf_x = tf.exp(-0.5 *((actions[ep][ts][2] - o[2][:,:,0]) / (o[2][:,:,1]))**2) * 1/(o[2][:,:,1]*tf.sqrt(2 *np.pi))
#         outputs[ts] = tf.stack([select, pdf_y,pdf_x], axis=-1  )
#       
#       log_probs = tf.keras.backend.log(outputs)
#       sums.append(tf.reduce_sum(log_probs, axis=[0,1,2]))
#     sums = tf.stack(sums, axis=0)   
#       
#     policy_loss = -tf.reduce_mean(sums*(np.array([[r] for r in rews])-hist_rew), axis=0)
#     policy_variables = De_Net.trainable_variables
#     grads = policy_tape.gradient(policy_loss, policy_variables)
#     De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
#   
#   hist_rew = (hist_rew*epoch+np.mean(rews))/(epoch+1)
# 
# show_batch(batch, positions)

# Commented out IPython magic to ensure Python compatibility.
# %%time
# for epoch in range(5,15):
#   print(f"Epoch {epoch+1}")
#   deeps,shallows, actions, rews , hist_loss= train_epoch(historical_loss=0. if not epoch else hist_loss)
#   
#   print()
#   with tf.GradientTape() as policy_tape:
#     sums = []
#     for ep, (ds, os) in enumerate(zip(deeps,shallows)):
#       outputs = [De_Net([d,o]) for d,o in zip(ds,os)]
#       for ts,o in enumerate(outputs):
#         select = tf.reduce_sum(tf.one_hot(actions[ep][ts][0],depth=2)*o[0], axis=-1)
#         pdf_y = tf.exp(-0.5 *((actions[ep][ts][1] - o[1][:,:,0]) / (o[1][:,:,1]))**2) * 1/(o[1][:,:,1]*tf.sqrt(2 *np.pi))
#         pdf_x = tf.exp(-0.5 *((actions[ep][ts][2] - o[2][:,:,0]) / (o[2][:,:,1]))**2) * 1/(o[2][:,:,1]*tf.sqrt(2 *np.pi))
#         outputs[ts] = tf.stack([select, pdf_y,pdf_x], axis=-1  )
#       
#       log_probs = tf.keras.backend.log(outputs)
#       sums.append(tf.reduce_sum(log_probs, axis=[0,1,2]))
#     sums = tf.stack(sums, axis=0)   
#       
#     policy_loss = -tf.reduce_mean(sums*(np.array([[r] for r in rews])-hist_rew), axis=0)
#     policy_variables = De_Net.trainable_variables
#     grads = policy_tape.gradient(policy_loss, policy_variables)
#     De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
#   
#   hist_rew = (hist_rew*epoch+np.mean(rews))/(epoch+1)
# 
# show_batch(batch, positions)

# Commented out IPython magic to ensure Python compatibility.
# %%time
# for epoch in range(15,25):
#   print(f"Epoch {epoch+1}")
#   deeps,shallows, actions, rews , hist_loss= train_epoch(historical_loss=0. if not epoch else hist_loss)
#   
#   print()
#   with tf.GradientTape() as policy_tape:
#     sums = []
#     for ep, (ds, os) in enumerate(zip(deeps,shallows)):
#       outputs = [De_Net([d,o]) for d,o in zip(ds,os)]
#       for ts,o in enumerate(outputs):
#         select = tf.reduce_sum(tf.one_hot(actions[ep][ts][0],depth=2)*o[0], axis=-1)
#         pdf_y = tf.exp(-0.5 *((actions[ep][ts][1] - o[1][:,:,0]) / (o[1][:,:,1]))**2) * 1/(o[1][:,:,1]*tf.sqrt(2 *np.pi))
#         pdf_x = tf.exp(-0.5 *((actions[ep][ts][2] - o[2][:,:,0]) / (o[2][:,:,1]))**2) * 1/(o[2][:,:,1]*tf.sqrt(2 *np.pi))
#         outputs[ts] = tf.stack([select, pdf_y,pdf_x], axis=-1  )
#       
#       log_probs = tf.keras.backend.log(outputs)
#       sums.append(tf.reduce_sum(log_probs, axis=[0,1,2]))
#     sums = tf.stack(sums, axis=0)   
#       
#     policy_loss = -tf.reduce_mean(sums*(np.array([[r] for r in rews])-hist_rew), axis=0)
#     policy_variables = De_Net.trainable_variables
#     grads = policy_tape.gradient(policy_loss, policy_variables)
#     De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
#   
#   hist_rew = (hist_rew*epoch+np.mean(rews))/(epoch+1)
# 
# show_batch(batch, positions)

"""### Experiment 5/10
 * B2B
 * normalized, .125
 * Default adam
 * loc= continuous actions
"""

# Commented out IPython magic to ensure Python compatibility.
# %%time
# positions = np.hstack([(np.random.randint(66,394, size=(4,5)))/230. - 1, np.random.randint(66,634, size=(4,5))/350. -1]).reshape(4,5,2)
# # positions = np.hstack([(np.random.randint(66,394, size=(4,5)))/460. , (np.random.randint(66,634, size=(4,5)))/700. ]).reshape(4,5,2)
# shallow = np.zeros((4,5,7)) # all other but deep feautures of DeNet input
# hist_rew = 0.
# 
# for epoch in range(5):
#   print(f"Epoch {epoch+1}")
#   deeps,shallows, actions, rews , hist_loss= train_epoch(historical_loss=0. if not epoch else hist_loss)
#   
#   print()
#   with tf.GradientTape() as policy_tape:
#     sums = []
#     for ep, (ds, os) in enumerate(zip(deeps,shallows)):
#       outputs = [De_Net([d,o]) for d,o in zip(ds,os)]
#       for ts,o in enumerate(outputs):
#         select = tf.reduce_sum(tf.one_hot(actions[ep][ts][0],depth=2)*o[0], axis=-1)
#         pdf_y = tf.exp(-0.5 *((actions[ep][ts][1] - o[1][:,:,0]) / (o[1][:,:,1]))**2) * 1/(o[1][:,:,1]*tf.sqrt(2 *np.pi))
#         pdf_x = tf.exp(-0.5 *((actions[ep][ts][2] - o[2][:,:,0]) / (o[2][:,:,1]))**2) * 1/(o[2][:,:,1]*tf.sqrt(2 *np.pi))
#         outputs[ts] = tf.stack([select, pdf_y,pdf_x], axis=-1  )
#       
#       log_probs = tf.keras.backend.log(outputs)
#       sums.append(tf.reduce_sum(log_probs, axis=[0,1,2]))
#     sums = tf.stack(sums, axis=0)   
#       
#     policy_loss = -tf.reduce_mean(sums*(np.array([[r] for r in rews])-hist_rew), axis=0)
#     policy_variables = De_Net.trainable_variables
#     grads = policy_tape.gradient(policy_loss, policy_variables)
#     De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
#   
#   hist_rew = (hist_rew*epoch+np.mean(rews))/(epoch+1)
# 
# show_batch(batch, positions)

for epoch in range(5,10):
  print(f"Epoch {epoch+1}")
  deeps,shallows, actions, rews , hist_loss= train_epoch(historical_loss=0. if not epoch else hist_loss)
  
  print()
  with tf.GradientTape() as policy_tape:
    sums = []
    for ep, (ds, os) in enumerate(zip(deeps,shallows)):
      outputs = [De_Net([d,o]) for d,o in zip(ds,os)]
      for ts,o in enumerate(outputs):
        select = tf.reduce_sum(tf.one_hot(actions[ep][ts][0],depth=2)*o[0], axis=-1)
        pdf_y = tf.exp(-0.5 *((actions[ep][ts][1] - o[1][:,:,0]) / (o[1][:,:,1]))**2) * 1/(o[1][:,:,1]*tf.sqrt(2 *np.pi))
        pdf_x = tf.exp(-0.5 *((actions[ep][ts][2] - o[2][:,:,0]) / (o[2][:,:,1]))**2) * 1/(o[2][:,:,1]*tf.sqrt(2 *np.pi))
        outputs[ts] = tf.stack([select, pdf_y,pdf_x], axis=-1  )
      
      log_probs = tf.keras.backend.log(outputs)
      sums.append(tf.reduce_sum(log_probs, axis=[0,1,2]))
    sums = tf.stack(sums, axis=0)   
      
    policy_loss = -tf.reduce_mean(sums*(np.array([[r] for r in rews])-hist_rew), axis=0)
    policy_variables = De_Net.trainable_variables
    grads = policy_tape.gradient(policy_loss, policy_variables)
    De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  
  hist_rew = (hist_rew*epoch+np.mean(rews))/(epoch+1)

show_batch(batch, positions)

for epoch in range(10,12):
  print(f"Epoch {epoch+1}")
  deeps,shallows, actions, rews , hist_loss= train_epoch(historical_loss=0. if not epoch else hist_loss)
  
  print()
  with tf.GradientTape() as policy_tape:
    sums = []
    for ep, (ds, os) in enumerate(zip(deeps,shallows)):
      outputs = [De_Net([d,o]) for d,o in zip(ds,os)]
      for ts,o in enumerate(outputs):
        select = tf.reduce_sum(tf.one_hot(actions[ep][ts][0],depth=2)*o[0], axis=-1)
        pdf_y = tf.exp(-0.5 *((actions[ep][ts][1] - o[1][:,:,0]) / (o[1][:,:,1]))**2) * 1/(o[1][:,:,1]*tf.sqrt(2 *np.pi))
        pdf_x = tf.exp(-0.5 *((actions[ep][ts][2] - o[2][:,:,0]) / (o[2][:,:,1]))**2) * 1/(o[2][:,:,1]*tf.sqrt(2 *np.pi))
        outputs[ts] = tf.stack([select, pdf_y,pdf_x], axis=-1  )
      
      log_probs = tf.keras.backend.log(outputs)
      sums.append(tf.reduce_sum(log_probs, axis=[0,1,2]))
    sums = tf.stack(sums, axis=0)   
      
    policy_loss = -tf.reduce_mean(sums*(np.array([[r] for r in rews])-hist_rew), axis=0)
    policy_variables = De_Net.trainable_variables
    grads = policy_tape.gradient(policy_loss, policy_variables)
    De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  
  hist_rew = (hist_rew*epoch+np.mean(rews))/(epoch+1)

show_batch(batch, positions)

"""## Validation at epoch end"""

def the_loop(iter,  val, size, start=False, historical_loss=0.):
  '''
  training or validation loop
  Arguments
    iter: data iterator 
    val: boolean, it is true if validation
    size: data size
    start: True when training for first epochs to select all cropped patches (may be unnecessary?)
  Returns
    deeps,shallows : all LSTM inputs to update gradients 
    avg_acc : average training/validation accuracy
  '''
  global positions, shallow
  avg_loss = 0.
  avg_acc = 0.
  
  deeps = []
  shallows = []
  actions_per_ep = []
  actions_ts = np.zeros((3,4,5))

  for step in range(size//4):
    batch = iter.next()
    cropped = 0
    guard = 0 # prevent infinite loop while cropping

    b = np.empty((4,5,112,112,3)) # batch for SaNet
    l = np.zeros((4,5), dtype='int')-1 # labels of batch for SaNet
    while cropped<20: 
      guard+=1
      # exctract 5 patches for each image in mini batch
      glmps = [tf.image.extract_glimpse(batch[0], size=(112,112), offsets=positions[:, i, :], centered=True, normalized=True) for i in range(5)] 
      glmps = tf.concat(glmps, 0)
      labels = np.vstack([batch[1] for _ in range(5)]).T

      # 6 loc_x, 7 loc_y
      shallow[:,:,5] = positions[:,:,0]
      shallow[:,:,6] = positions[:,:,1]
      # 3 pred, 4 true pred
      shallow[:,:,3] = np.argmax(SA_Net.predict_on_batch(glmps), axis=1).reshape(4,5)
      shallow[:,:,4] = labels


      shallow_t = tf.convert_to_tensor(shallow, dtype=tf.float64)
      deep = tf.reshape(SA_Net_features(glmps), (4,5,128))

      if not val:
        deeps.append(deep)
        shallows.append(shallow_t)

      out = [o.numpy() for o in De_Net([deep,shallow_t ]) ] 

      patches = glmps.numpy().reshape(4,5, 112,112,3)
      # print(positions[0])
      # iterate over DeNet actions, select patches and update centers for next iter
      for i,im in enumerate(out[1]):
        for j,o in enumerate(im):
          actions_ts[0][i][j] = 0
          if l[i][j]==-1 and (out[0][i][j][1]-out[0][i][j][0]>=0. or guard>10 ):
            b[i][j] = patches[i][j]
            l[i][j] = labels[i][j]
            cropped+=1
            actions_ts[0][i][j] = 1
          # avoid croping the edges
          y = np.clip(np.random.normal(loc=o[0], scale=o[1]) , 66/230 -1, 394/230-1) #*164/230
          x = np.clip(np.random.normal(loc= out[2][i][j][0], scale=out[2][i][j][1]), 66/350-1,634/350-1) #*294/350

          positions[i,j,:] = [y , x ]
          actions_ts[1][i][j] = y
          actions_ts[2][i][j] = x
      if not val:
        actions_per_ep.append(actions_ts)
   
    loss, acc = SA_Net.test_on_batch(b.reshape((20,112,112,3)), l.reshape((20,1))) if val else SA_Net.train_on_batch(b.reshape((20,112,112,3)), l.reshape((20,1)))
      
    avg_loss = ( avg_loss*step + loss) / (step + 1) # running average
    avg_acc = (avg_acc*step + acc) / (step + 1)

    if not val :
      # (historical?) training loss, state feauture 0
      historical_loss =  ( historical_loss*(step + shallow[0][0][2] ) + loss) / (step + shallow[0][0][2] + 1)
      shallow[:,:,0] = np.array([historical_loss]*20).reshape(4,5)
      shallow[:,:,2] = np.array([shallow[0][0][2]+1]*20).reshape(4,5) # num of iterations (epochs?)
    sys.stdout.write(f"\r {guard if guard>0 else ''} Step {step+1}: {'                         Val ' if val else ''}Loss {avg_loss} , {'Val ' if val else ''}Accuracy {avg_acc}")
    sys.stdout.flush()
  # update at the end of validation or during?
    if val:  
      shallow[:,:,1] = np.array([max(avg_acc, shallow[0][0][1])]*20).reshape(4,5) # best val acc
  return deeps, shallows,actions_per_ep,avg_acc,historical_loss

# Commented out IPython magic to ensure Python compatibility.
# %%time
# # #random initial patch center positions
# positions = np.hstack([(2*np.random.randint(66,394, size=(4,5)))/460. - 1, (2 *np.random.randint(66,634, size=(4,5)))/700. -1]).reshape(4,5,2)
# # positions = np.hstack([(np.random.randint(66,394, size=(4,5)))/460. , (np.random.randint(66,634, size=(4,5)))/700. ]).reshape(4,5,2)
# shallow = np.zeros((4,5,7)) # all other but deep feautures of DeNet input
# hist_rew = 0.
# rews = []
# for epoch in range(5):
#   print(f"Epoch {epoch+1}")
#   deeps,shallows,actions,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
#   
#   print()
# 
#   _,__,_,reward,_ = the_loop(val_iter,  True, val_size)
# # update DeNet for every input
# 
#   with tf.GradientTape() as policy_tape:
# 
#     outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
#     for ts,o in enumerate(outputs):
#       select = tf.reduce_sum(tf.one_hot(actions[ts][0],depth=2)*o[0], axis=-1)
#       pdf_y = tf.exp(-0.5 *((actions[ts][1] - o[1][:,:,0]) / (o[1][:,:,1]))**2) * 1/(o[1][:,:,1]*tf.sqrt(2 *np.pi))
#       pdf_x = tf.exp(-0.5 *((actions[ts][2] - o[2][:,:,0]) / (o[2][:,:,1]))**2) * 1/(o[2][:,:,1]*tf.sqrt(2 *np.pi))
#       outputs[ts] = tf.stack([select, pdf_y,pdf_x], axis=-1  )
#       
#     log_probs = tf.keras.backend.log(outputs)
#     policy_loss = -tf.math.reduce_sum(log_probs, axis=0)*(reward-hist_rew)
# 
# 
# 
#   policy_variables = De_Net.trainable_variables
#   grads = policy_tape.gradient(policy_loss, policy_variables)
#   # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
#   De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
#   print()
#   hist_rew = (hist_rew*epoch+reward)/(epoch+1)
# show_batch(batch, positions)

# Commented out IPython magic to ensure Python compatibility.
# %%time
# for epoch in range(15,20):
#   print(f"Epoch {epoch+1}")
#   deeps,shallows,actions,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
#   
#   print()
# 
#   _,__,_,reward,_ = the_loop(val_iter,  True, val_size)
# # update DeNet for every input
# 
#   with tf.GradientTape() as policy_tape:
# 
#     outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
#     for ts,o in enumerate(outputs):
#       select = tf.reduce_sum(tf.one_hot(actions[ts][0],depth=2)*o[0], axis=-1)
#       pdf_y = tf.exp(-0.5 *((actions[ts][1] - o[1][:,:,0]) / (o[1][:,:,1]))**2) * 1/(o[1][:,:,1]*tf.sqrt(2 *np.pi))
#       pdf_x = tf.exp(-0.5 *((actions[ts][2] - o[2][:,:,0]) / (o[2][:,:,1]))**2) * 1/(o[2][:,:,1]*tf.sqrt(2 *np.pi))
#       outputs[ts] = tf.stack([select, pdf_y,pdf_x], axis=-1  )
#       
#     log_probs = tf.keras.backend.log(outputs)
#     policy_loss = -tf.math.reduce_sum(log_probs, axis=0)*(reward-hist_rew)
# 
# 
# 
#   policy_variables = De_Net.trainable_variables
#   grads = policy_tape.gradient(policy_loss, policy_variables)
#   # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
#   De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
#   print()
#   hist_rew = (hist_rew*epoch+reward)/(epoch+1)
# show_batch(batch, positions)

for epoch in range(10, 15):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input

  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    for i,o in enumerate(outputs):
      tmp = [tf.reshape(t, [4,5,1]) for t in tf.unstack(o[1], axis=-1)]
      outputs[i] = tf.squeeze(tf.stack([o[0], tmp[0], tmp[1]], axis=-1))
    ## code for Dense(2) x 2  
    # actions, locs = tf.unstack(outputs, axis=1)
    # log_probs = tf.keras.backend.log(tf.cast(tf.concat([actions, locs], axis=-1), 'float64'))
    log_probs = tf.keras.backend.log(outputs)
    policy_loss = -tf.math.reduce_sum(log_probs, axis=[1, 0])*(reward-hist_rew)



  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)
show_batch(batch, positions)

test_iter = test_ds.as_numpy_iterator()

_,_,fr,_ = the_loop(test_iter, True, test_size)

"""### Experiment 26/9

Saved models 120 epochs, SaNet did train 33-->47%, not so good generalization. 

```
Epoch 1
 1 Step 252: Loss 1.4167722812727868 , Accuracy 0.3353174603174603
 1 Step 104:                          Val Loss 9.90010475324764 , Val Accuracy 0.4567307692307692
Epoch 2
 1 Step 252: Loss 1.529041473679837 , Accuracy 0.3521825396825396
 1 Step 104:                          Val Loss 4.304744035170259 , Val Accuracy 0.43028846153846156
Epoch 3
 1 Step 252: Loss 1.520365798659253 , Accuracy 0.3541666666666667
 1 Step 104:                          Val Loss 4.343825320963965 , Val Accuracy 0.45913461538461536
Epoch 4
 1 Step 252: Loss 1.6632714321511601 , Accuracy 0.3382936507936508
 1 Step 104:                          Val Loss 3.4249503813110858 , Val Accuracy 0.17548076923076916
Epoch 5
 1 Step 252: Loss 1.5349681016283183 , Accuracy 0.3680555555555557
 1 Step 104:                          Val Loss 3.468753189421617 , Val Accuracy 0.12980769230769232
 ```
"""

show_batch(batch, positions)

for epoch in range(5, 15):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input

  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    for i,o in enumerate(outputs):
      tmp = [tf.reshape(t, [4,5,1]) for t in tf.unstack(o[1], axis=-1)]
      outputs[i] = tf.squeeze(tf.stack([o[0], tmp[0], tmp[1]], axis=-1))
    ## code for Dense(2) x 2  
    # actions, locs = tf.unstack(outputs, axis=1)
    # log_probs = tf.keras.backend.log(tf.cast(tf.concat([actions, locs], axis=-1), 'float64'))
    log_probs = tf.keras.backend.log(outputs)
    policy_loss = -tf.math.reduce_sum(tf.reduce_sum(log_probs, axis=1), axis=0)*(reward-hist_rew)



  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)
show_batch(batch, positions)

for epoch in range(15, 30):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input

  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    for i,o in enumerate(outputs):
      tmp = [tf.reshape(t, [4,5,1]) for t in tf.unstack(o[1], axis=-1)]
      outputs[i] = tf.squeeze(tf.stack([o[0], tmp[0], tmp[1]], axis=-1))
    ## code for Dense(2) x 2  
    # actions, locs = tf.unstack(outputs, axis=1)
    # log_probs = tf.keras.backend.log(tf.cast(tf.concat([actions, locs], axis=-1), 'float64'))
    log_probs = tf.keras.backend.log(outputs)
    policy_loss = -tf.math.reduce_sum(tf.reduce_sum(log_probs, axis=1), axis=0)*(reward-hist_rew)



  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)
show_batch(batch, positions)

for epoch in range(30, 50):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input

  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    for i,o in enumerate(outputs):
      tmp = [tf.reshape(t, [4,5,1]) for t in tf.unstack(o[1], axis=-1)]
      outputs[i] = tf.squeeze(tf.stack([o[0], tmp[0], tmp[1]], axis=-1))
    ## code for Dense(2) x 2  
    # actions, locs = tf.unstack(outputs, axis=1)
    # log_probs = tf.keras.backend.log(tf.cast(tf.concat([actions, locs], axis=-1), 'float64'))
    log_probs = tf.keras.backend.log(outputs)
    policy_loss = -tf.math.reduce_sum(tf.reduce_sum(log_probs, axis=1), axis=0)*(reward-hist_rew)



  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)
show_batch(batch, positions)

for epoch in range(50, 70):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input

  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    for i,o in enumerate(outputs):
      tmp = [tf.reshape(t, [4,5,1]) for t in tf.unstack(o[1], axis=-1)]
      outputs[i] = tf.squeeze(tf.stack([o[0], tmp[0], tmp[1]], axis=-1))
    ## code for Dense(2) x 2  
    # actions, locs = tf.unstack(outputs, axis=1)
    # log_probs = tf.keras.backend.log(tf.cast(tf.concat([actions, locs], axis=-1), 'float64'))
    log_probs = tf.keras.backend.log(outputs)
    policy_loss = -tf.math.reduce_sum(tf.reduce_sum(log_probs, axis=1), axis=0)*(reward-hist_rew)



  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)
show_batch(batch, positions)

for epoch in range(70, 100):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input

  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    for i,o in enumerate(outputs):
      tmp = [tf.reshape(t, [4,5,1]) for t in tf.unstack(o[1], axis=-1)]
      outputs[i] = tf.squeeze(tf.stack([o[0], tmp[0], tmp[1]], axis=-1))
    ## code for Dense(2) x 2  
    # actions, locs = tf.unstack(outputs, axis=1)
    # log_probs = tf.keras.backend.log(tf.cast(tf.concat([actions, locs], axis=-1), 'float64'))
    log_probs = tf.keras.backend.log(outputs)
    policy_loss = -tf.math.reduce_sum(tf.reduce_sum(log_probs, axis=1), axis=0)*(reward-hist_rew)



  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)
show_batch(batch, positions)

for epoch in range(100, 120):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input

  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    for i,o in enumerate(outputs):
      tmp = [tf.reshape(t, [4,5,1]) for t in tf.unstack(o[1], axis=-1)]
      outputs[i] = tf.squeeze(tf.stack([o[0], tmp[0], tmp[1]], axis=-1))
    ## code for Dense(2) x 2  
    # actions, locs = tf.unstack(outputs, axis=1)
    # log_probs = tf.keras.backend.log(tf.cast(tf.concat([actions, locs], axis=-1), 'float64'))
    log_probs = tf.keras.backend.log(outputs)
    policy_loss = -tf.math.reduce_sum(tf.reduce_sum(log_probs, axis=1), axis=0)*(reward-hist_rew)



  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)
show_batch(batch, positions)

"""### Experiment 25/9 
Same config as 24/9 but noise=50/dim

```
Epoch 1
 1 Step 252: Loss 1.3684685696928123 , Accuracy 0.35912698412698413
 1 Step 104:                          Val Loss 10.374517702811046 , Val Accuracy 0.4567307692307692
Epoch 2
 1 Step 252: Loss 1.5369671746563618 , Accuracy 0.3541666666666667
 1 Step 104:                          Val Loss 5.845840767083607 , Val Accuracy 0.4567307692307692
Epoch 3
 1 Step 252: Loss 1.5951837076110735 , Accuracy 0.3412698412698414
 1 Step 104:                          Val Loss 3.824998480645438 , Val Accuracy 0.44471153846153844
Epoch 4
 1 Step 252: Loss 1.5572875102848867 , Accuracy 0.35714285714285715
 1 Step 104:                          Val Loss 3.1672398932278156 , Val Accuracy 0.3389423076923076
Epoch 5
 1 Step 252: Loss 1.5700362925491647 , Accuracy 0.36904761904761907
 1 Step 104:                          Val Loss 2.811896381183312 , Val Accuracy 0.40865384615384615
 ```
"""

for epoch in range(5,10):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input

  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    for i,o in enumerate(outputs):
      tmp = [tf.reshape(t, [4,5,1]) for t in tf.unstack(o[1], axis=-1)]
      outputs[i] = tf.squeeze(tf.stack([o[0], tmp[0], tmp[1]], axis=-1))
    # actions, locs = tf.unstack(outputs, axis=1)
    # log_probs = tf.keras.backend.log(tf.cast(tf.concat([actions, locs], axis=-1), 'float64'))
    log_probs = tf.keras.backend.log(outputs)
    policy_loss = -tf.math.reduce_sum(tf.reduce_mean(log_probs, axis=1), axis=0)*(reward-hist_rew)



  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)
show_batch(batch, positions)

for epoch in range(10,20):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input

  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    for i,o in enumerate(outputs):
      tmp = [tf.reshape(t, [4,5,1]) for t in tf.unstack(o[1], axis=-1)]
      outputs[i] = tf.squeeze(tf.stack([o[0], tmp[0], tmp[1]], axis=-1))
    # actions, locs = tf.unstack(outputs, axis=1)
    # log_probs = tf.keras.backend.log(tf.cast(tf.concat([actions, locs], axis=-1), 'float64'))
    log_probs = tf.keras.backend.log(outputs)
    policy_loss = -tf.math.reduce_sum(tf.reduce_mean(log_probs, axis=1), axis=0)*(reward-hist_rew)



  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)
show_batch(batch, positions)

for epoch in range(20,30):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input

  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    for i,o in enumerate(outputs):
      tmp = [tf.reshape(t, [4,5,1]) for t in tf.unstack(o[1], axis=-1)]
      outputs[i] = tf.squeeze(tf.stack([o[0], tmp[0], tmp[1]], axis=-1))
    # actions, locs = tf.unstack(outputs, axis=1)
    # log_probs = tf.keras.backend.log(tf.cast(tf.concat([actions, locs], axis=-1), 'float64'))
    log_probs = tf.keras.backend.log(outputs)
    policy_loss = -tf.math.reduce_sum(tf.reduce_mean(log_probs, axis=1), axis=0)*(reward-hist_rew)



  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)
show_batch(batch, positions)

"""### Experiment 24/9
Very Promising ! \\


```
Epoch 1
 1 Step 252: Loss 1.3671510603205204 , Accuracy 0.37896825396825407
 1 Step 104:                          Val Loss 8.442637980027998 , Val Accuracy 0.4567307692307692
 ```
"""

show_batch(batch, positions)

for epoch in range(7,20):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input

  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    for i,o in enumerate(outputs):
      tmp = [tf.reshape(t, [4,5,1]) for t in tf.unstack(o[1], axis=-1)]
      outputs[i] = tf.squeeze(tf.stack([o[0], tmp[0], tmp[1]], axis=-1))
    # actions, locs = tf.unstack(outputs, axis=1)
    # log_probs = tf.keras.backend.log(tf.cast(tf.concat([actions, locs], axis=-1), 'float64'))
    log_probs = tf.keras.backend.log(outputs)
    policy_loss = -tf.math.reduce_sum(tf.reduce_mean(log_probs, axis=1), axis=0)*(reward-hist_rew)



  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)
show_batch(batch, positions)

for epoch in range(20,40):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input

  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    for i,o in enumerate(outputs):
      tmp = [tf.reshape(t, [4,5,1]) for t in tf.unstack(o[1], axis=-1)]
      outputs[i] = tf.squeeze(tf.stack([o[0], tmp[0], tmp[1]], axis=-1))
    # actions, locs = tf.unstack(outputs, axis=1)
    # log_probs = tf.keras.backend.log(tf.cast(tf.concat([actions, locs], axis=-1), 'float64'))
    log_probs = tf.keras.backend.log(outputs)
    policy_loss = -tf.math.reduce_sum(tf.reduce_mean(log_probs, axis=1), axis=0)*(reward-hist_rew)



  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)
show_batch(batch, positions)

for epoch in range(40,50):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input

  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    for i,o in enumerate(outputs):
      tmp = [tf.reshape(t, [4,5,1]) for t in tf.unstack(o[1], axis=-1)]
      outputs[i] = tf.squeeze(tf.stack([o[0], tmp[0], tmp[1]], axis=-1))
    # actions, locs = tf.unstack(outputs, axis=1)
    # log_probs = tf.keras.backend.log(tf.cast(tf.concat([actions, locs], axis=-1), 'float64'))
    log_probs = tf.keras.backend.log(outputs)
    policy_loss = -tf.math.reduce_sum(tf.reduce_mean(log_probs, axis=1), axis=0)*(reward-hist_rew)



  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)
show_batch(batch, positions)

for epoch in range(50,60):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input

  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    for i,o in enumerate(outputs):
      tmp = [tf.reshape(t, [4,5,1]) for t in tf.unstack(o[1], axis=-1)]
      outputs[i] = tf.squeeze(tf.stack([o[0], tmp[0], tmp[1]], axis=-1))
    # actions, locs = tf.unstack(outputs, axis=1)
    # log_probs = tf.keras.backend.log(tf.cast(tf.concat([actions, locs], axis=-1), 'float64'))
    log_probs = tf.keras.backend.log(outputs)
    policy_loss = -tf.math.reduce_sum(tf.reduce_mean(log_probs, axis=1), axis=0)*(reward-hist_rew)



  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)
show_batch(batch, positions)

for epoch in range(60,80):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input

  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    for i,o in enumerate(outputs):
      tmp = [tf.reshape(t, [4,5,1]) for t in tf.unstack(o[1], axis=-1)]
      outputs[i] = tf.squeeze(tf.stack([o[0], tmp[0], tmp[1]], axis=-1))
    # actions, locs = tf.unstack(outputs, axis=1)
    # log_probs = tf.keras.backend.log(tf.cast(tf.concat([actions, locs], axis=-1), 'float64'))
    log_probs = tf.keras.backend.log(outputs)
    policy_loss = -tf.math.reduce_sum(tf.reduce_mean(log_probs, axis=1), axis=0)*(reward-hist_rew)



  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)
show_batch(batch, positions)

# !mkdir models
tf.keras.models.save_model(De_Net, 'DeNet/') #120 epochs
tf.keras.models.save_model(SA_Net, 'SaNet/') #120 epochs
# !du -h models
# !rm -r models

cp -r SaNet drive/My\ Drive/

"""### Experiment 23/9"""

for epoch in range(5,10):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input
  
  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    actions, locs = tf.unstack(outputs, axis=1)
    log_probs = tf.keras.backend.log(tf.cast(tf.concat([tf.reduce_max(actions, axis=-1, keepdims=True), locs], axis=-1), 'float64'))
    policy_loss = tf.math.reduce_sum(tf.reduce_mean(log_probs, axis=1), axis=0)*(reward-hist_rew)

  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)

for epoch in range(40,50):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ , hist_loss= the_loop(train_iter, False, train_size, historical_loss=0. if not epoch else hist_loss)
  
  print()

  _,__,reward,_ = the_loop(val_iter,  True, val_size)
# update DeNet for every input
  
  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    actions, locs = tf.unstack(outputs, axis=1)
    log_probs = tf.keras.backend.log(tf.cast(tf.concat([tf.reduce_max(actions, axis=-1, keepdims=True), locs], axis=-1), 'float64'))
    policy_loss = tf.math.reduce_sum(tf.reduce_mean(log_probs, axis=1), axis=0)*(reward-hist_rew)

  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  # grads = tuple(tf.clip_by_norm(g, 5.) for g in grads)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()
  hist_rew = (hist_rew*epoch+reward)/(epoch+1)

"""### Old Experiment"""

# Older version for zoom 100 in which the first 5 patches where selected during validation (action was ignored)
for epoch in range(20, 41 ):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ = the_loop(train_iter, False, train_size)
  
  print()

  _,__,reward = the_loop(val_iter,  True, val_size)


  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    log_probs = tf.transpose(tf.reshape(tf.nn.log_softmax(outputs), (len(outputs)*20, 3)))
    policy_loss = -tf.math.reduce_mean(log_probs*reward, axis=1)

  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()

for epoch in range(40, 81 ):
  print(f"Epoch {epoch+1}")
  deeps,shallows,_ = the_loop(train_iter, False, train_size)
  
  print()

  _,__,reward = the_loop(val_iter,  True, val_size)


  with tf.GradientTape() as policy_tape:

    outputs = [De_Net([d,o]) for d,o in zip(deeps,shallows)]
    log_probs = tf.transpose(tf.reshape(tf.nn.log_softmax(outputs), (len(outputs)*20, 3)))
    policy_loss = -tf.math.reduce_mean(log_probs*reward, axis=1)

  policy_variables = De_Net.trainable_variables
  grads = policy_tape.gradient(policy_loss, policy_variables)
  De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
  print()

"""## Implementation with tf.environment 

DeNet--> Agent, SaNet-->Environment = Big overhead, too slow
"""

!sudo pip install tf-agents
# from tf_agents.environments import py_environment
# from tf_agents.environments import tf_environment
# from tf_agents.environments import tf_py_environment
# from tf_agents.environments import utils
# from tf_agents import specs 
# from tf_agents.environments import wrappers
# from tf_agents.environments import suite_gym
# from tf_agents.trajectories import time_step as ts
# from tf_agents.utils import common

train_iter = train_ds.unbatch().batch(1).as_numpy_iterator()

val_iter = val_ds.unbatch().batch(4).as_numpy_iterator()

def evaluation(init_state):
  def make_val_batch_glimpses(x,y,state):
    for i in range(5):
      locations = get_actions(state).tolist() 
      glimpse =  tf.image.extract_glimpse(x, size=(112,112), offsets=locations, centered=False, normalized=False)
      tmp_batch = (glimpse, y) if not i else (tf.concat([tmp_batch[0],glimpse], 0), np.append(tmp_batch[1], y))
      
      preds = SA_Net(glimpse)
      features = SA_Net_features(glimpse)
      # predictions to labels
      preds = list(np.argmax(preds, axis=1)) 
      
      # update state
      for i, pred in enumerate(preds):
        state[i]['pred_label'] = pred
        state[i]['true_label'] = y[i] # array index = 0 first item in batch
        state[i]['loc_y'] = locations[i][0]
        state[i]['loc_x'] = locations[i][1]       
        state[i]['deep_features'] = features[i].numpy()
    return tmp_batch

  val_acc = []
  val_loss = []

  for i in range(5):
    x0,y0 = val_iter.next()
    x,y = make_val_batch_glimpses(x0,y0,the_env.val_batch)
    vl, va = SA_Net.test_on_batch(x,y)
    val_acc.append(va)
    val_loss.append(vl)
    # sys.stdout.write(f"\r{i}/{val_size//4}")
    # sys.stdout.flush()

  val_loss = np.array(val_loss)
  igamma = np.argwhere( val_loss < .25)
  acc = np.mean(val_acc)
  los = np.mean(val_loss)
  print(f'Val Loss: {los} -- Val Acc: {acc}')
  return igamma[0][0] if igamma.size else 0., acc

def get_action(observation, network=De_Net, first=False):
  deep_inp = np.array(observation['deep_features']).reshape(1, 1,128)
  stats = np.array([v for k,v in observation.items() if k!='deep_features']).reshape(1, 1,7)
  deep_inp_tensor = tf.convert_to_tensor(deep_inp)
  stats_tensor = tf.convert_to_tensor(stats)

  net_output = network([deep_inp_tensor, stats_tensor])
  output = np.vstack([tf.math.argmax(logits, axis=-1)for logits in  net_output]).T[0]
  loc =  tf.round(tf.random.normal([2], np.mean((output[1:]+112).tolist()),np.sqrt(112)))

  action = [1 if first else output[0] , loc.numpy()]
  return action, net_output

def get_actions(observations, network=De_Net ):
  deep_inp =  np.array([observation['deep_features'] for observation in observations]).reshape(4, 1,128)
  stats =  np.array([v for observation in observations for k,v in observation.items()  if k!='deep_features']).reshape(4,1,7)
  deep_inp_tensor = tf.convert_to_tensor(deep_inp)
  stats_tensor = tf.convert_to_tensor(stats)
  output = network([deep_inp_tensor, stats_tensor])
  output = np.vstack([tf.math.argmax(logits, axis=-1).numpy() for logits in  output]).T
  # locs =  tf.round(tf.random.normal([2], np.mean((output[1:]+112).tolist()),np.sqrt(112)))

  return output[:, 1:]+112

def get_input(observations):
  # n = len(observations)
  deep_inp =  np.array([observation['deep_features'] for observation in observations]).reshape(4, 1,128)
  stats =  np.array([v for observation in observations for k,v in observation.items()  if k!='deep_features']).reshape(4,1,7)
  deep_inp_tensor = tf.convert_to_tensor(deep_inp)
  stats_tensor = tf.convert_to_tensor(stats)
  return [ deep_inp_tensor, stats_tensor]

def validation(init_state):
  def make_val_batch_glimpses(x,y,state):
    for i in range(5):
      locations = get_actions(state).tolist() 
      glimpse =  tf.image.extract_glimpse(x, size=(112,112), offsets=locations, centered=False, normalized=False)
      tmp_batch = (glimpse, y) if not i else (tf.concat([tmp_batch[0],glimpse], 0), np.append(tmp_batch[1], y))
      
      preds = SA_Net(glimpse)
      features = SA_Net_features(glimpse)
      # predictions to labels
      preds = list(np.argmax(preds, axis=1)) 
      
      # update state
      for i, pred in enumerate(preds):
        state[i]['pred_label'] = pred
        state[i]['true_label'] = y[i] # array index = 0 first item in batch
        state[i]['loc_y'] = locations[i][0]
        state[i]['loc_x'] = locations[i][1]       
        state[i]['deep_features'] = features[i].numpy()
    return tmp_batch

  val_acc = []
  val_loss = []

  for i in range(val_size//4):
    x0,y0 = val_iter.next()
    x,y = make_val_batch_glimpses(x0,y0,the_env.val_batch)
    vl, va = SA_Net.test_on_batch(x,y)
    val_acc.append(va)
    val_loss.append(vl)
    sys.stdout.write(f"\r{i+1}/{val_size//4}")
    sys.stdout.flush()

  val_loss = np.array(val_loss)
  igamma = np.argwhere( val_loss < .25)
  acc = np.mean(val_acc)
  los = np.mean(val_loss)
  print(f'Val Loss: {los} -- Val Acc: {acc}')
  return igamma[0][0] if igamma.size else 0., acc

from collections import deque
FIRST = ts.StepType.FIRST
MID = ts.StepType.MID
LAST = ts.StepType.LAST

class BreakHisEnv(py_environment.PyEnvironment):

  def __init__(self,iter=train_iter, network=SA_Net, features_exctractor=SA_Net_features):
    

    
     
    action_spec = specs.ArraySpec((3, ), dtype=np.uint32, name='action')
    observation_spec = specs.ArraySpec(shape=(135,), dtype=np.float32, name='observation')
    time_step_spec = ts.time_step_spec(observation_spec)

    self.iter = iter
    self._network = network
    self._feature_extractor = features_exctractor
    keys = ['avg_loss', 'max_acc', 'num_iters', 'pred_label', 'true_label', 'loc_y', 'loc_x', 'deep_features']
    vals = [0. for _ in range(7)] 
    vals.append([0. for _ in range(128)])
    self._state = {k:v for k,v in zip(keys, vals)}
    self.cropped = 0
    self.steps = 0
    self.accs = []
    self._cur_train_batch = self.iter.next()
    self.val_batch = deque()
    self.val_cnt =  0
    self.tr_steps = 0

  def action_spec(self):
    return self._action_spec

  def observation_spec(self):
    return self._observation_spec

  def curr_ts(self):
    step_type = MID if self.steps else FIRST
    reward = 0.
    if self.cropped == 20 or self.steps == 20:
      # episode ended must train SA_Net

      loss, acc = self._network.train_on_batch(self._cur_batch[0], tf.convert_to_tensor(self._cur_batch[1], dtype=tf.uint8))
      # validate
      self.tr_steps +=1

      print(f'Step {self.tr_steps} Loss: {loss}, Accuracy: {acc}')
      self.accs.append(acc)
      # val_acc = get_val_acc(get_actions(list(self.val_batch)), self.val_batch)
      igamma, val_acc = validation(the_env.val_batch)

      # update learning status representation
      self._state['avg_loss'] = (self._state['avg_loss']*self._state['num_iters'] + loss) / (self._state['num_iters'] + 1)
      self._state['max_acc'] = max(self._state['max_acc'], val_acc)
      self._state['num_iters'] +=1

      # set reward
      # preds = self._network.predict(val_set)
      # val_losses = tf.keras.loss.tf.keras.losses.SparseCategoricalCrossentropy()(val_set, preds)
      # igamma = np.argwhere(val_losses < .25)
      # rc = -np.log(igamma/200)
      reward = val_acc #+ rc    
      step_type = LAST
    
    return ts.TimeStep(step_type, reward,  0.,self._state)

  def _reset(self):
    self.cropped = 0 
    self.steps = 0
    return self.curr_ts()

  def _step(self, action):
    self.steps+=1
    

    if not (self.cropped%5):
      self._cur_train_batch = self.iter.next()

    location = action[1] 
    # fill batch for training when episode is ended
    curr_glimpse = (tf.image.extract_glimpse(self._cur_train_batch[0], size=(112,112), offsets=[location], centered=False, normalized=False), self._cur_train_batch[1])
    if action[0] :
      self._cur_batch = curr_glimpse if not (self.cropped%20) else (tf.concat([self._cur_batch[0],curr_glimpse[0]], 0), np.append(self._cur_batch[1], curr_glimpse[1]))
      self.cropped+=1
    

    # update incoming data statictics

    # print('predict ')
    pred = self._network.predict(curr_glimpse[0])
    # predictions to labels
    pred = list(np.argmax(pred, axis=1))      
    self._state['pred_label'] = pred[0]
    self._state['true_label'] = curr_glimpse[1][0] # array index = 0 first item in batch
    self._state['loc_y'] = location[0]
    self._state['loc_x'] = location[1]


    # print('features')
    features = self._feature_extractor.predict(curr_glimpse[0], steps=1)
    self._state['deep_features'] = features.reshape(1,128)
  
    self.val_batch.append(self._state)

    if len(self.val_batch)==5:
      self.val_batch.popleft()


    if action[0]>1:
      raise ValueError('`action` should be 0 or 1.')
    
    return self.curr_ts()

class ValEnvironment(BreakHisEnv):
  def __init__(self, initial_state):
    super().__init__(iter=val_iter)
    self._state = initial_state
    self.acc = 0.0
    # self.batched_state = self._state
  
  def curr_ts(self):
    step_type = MID if self.steps else FIRST
    reward = 0.
    if self.steps== (5*(val_size // 4)) :
      

      _, val_acc = self._network.evaluate(self._cur_batch[0].numpy() , self._cur_batch[1], steps=val_size//4, batch_size=5)
      # self.acc = (self.acc*self._state['num_iters'] + val_acc) / (self._state['num_iters'] + 1)
      self.acc = val_acc
      step_type = LAST
      # self.cropped = 0
    
    return ts.TimeStep(step_type,self.acc,  0.,self._state) 
  
  def _step(self, action):
  #  action is always 1
    self.steps+=1
   

    if not (self.cropped%5):
      self._cur_train_batch = self.iter.next()
    # batched action
    locations = action.tolist() 
    # fill batch for training when episode is ended
    curr_glimpse = tf.image.extract_glimpse(self._cur_train_batch[0], size=(112,112), offsets=locations, centered=False, normalized=False)
    self._cur_batch = (curr_glimpse, self._cur_train_batch[1]) if  self.steps == 1 else (tf.concat([self._cur_batch[0],curr_glimpse], 0), np.append(self._cur_batch[1],self._cur_train_batch[1]))
    self.cropped+=1

    # update incoming data statictics

    # print('predict ')
    preds = self._network(curr_glimpse)
    features = self._feature_extractor(curr_glimpse)
    # predictions to labels
    preds = list(np.argmax(preds, axis=1)) 
    # self.batched_state = []     
    for i, pred in enumerate(preds):
      self._state[i]['pred_label'] = pred
      self._state[i]['true_label'] = self._cur_train_batch[1][i] # array index = 0 first item in batch
      self._state[i]['loc_y'] = locations[i][0]
      self._state[i]['loc_x'] = locations[i][1]


      # print('features')
      self._state[i]['deep_features'] = features[i].numpy()
      # self.batched_state.append(self._state)
    
    return self.curr_ts()

import sys 
def get_val_acc(actions, init_obs):
  val_env = ValEnvironment(init_obs)
  val_env.reset()
  for i in range(val_size//4):
    for _ in range(5):
      timestep = val_env.step(actions)
      actions =  get_actions(timestep.observation)
    sys.stdout.write(f"\r{i}/{val_size//4}")
    sys.stdout.flush()
  return timestep.reward

the_env._cur_batch[1]

# Commented out IPython magic to ensure Python compatibility.
# %%time 
# get_val_acc(actions, obs)

# Commented out IPython magic to ensure Python compatibility.
# %%time 
# first = True
# the_env = BreakHisEnv()
# for i in range(1):
#   the_env.tr_steps = 0
#   print(f'Epoch {i}')
#   for j in range(1):
#     tst = the_env.reset()
#     actions = []
#     states = []
#     logits_list = []
# 
#     while tst.step_type!= LAST:
#       action, _ = get_action(tst.observation, first=first)
#       actions.append(action)
#       # logits_list.append(logits)
#       states.append(tst.observation)
#       tst = the_env.step(action)
#     reward = tst.reward
#     first = False
#     actions = np.array(actions)
#     select = actions[:, 0]
#     lx = np.array([ l[0] for l in actions[:,1]])
#     ly = np.array([ l[1] for l in actions[:,1]])
#     acts = [select, lx, ly]
#     actions_space = [2, 236, 476]
# 
# 
# # with tf.GradientTape() as policy_tape:
# #   outputs = [ De_Net(get_input(states[i:i+4])) for i in range(0, len(states), 4) ]
# #   t = [ np.array([tf.unstack(a)  for a in output]).T for output in outputs]
# #   logits = np.array([ [tf.nn.log_softmax(a) for a in o] for out in t for o in out ])
# 
# #   # logits = [[ tf.nn.log_softmax(b) for b in a]  for a in logits_list]
# #   one_hot = [tf.one_hot(a, n) for a,n in zip(acts, actions_space)]
# 
# #   log_probs = [ tf.math.reduce_sum(tf.stack(l)* o, axis=1) for l, o in zip(logits.T, one_hot)]
# 
# #   policy_loss = -tf.math.reduce_mean(log_probs* tf.stop_gradient(reward), axis = 1)
# 
# # policy_variables = De_Net.trainable_variables
# # grads = policy_tape.gradient(policy_loss, policy_variables)
# # De_Net.optimizer.apply_gradients(zip(grads, policy_variables))
# # print(f'Training Acc: {np.mean(the_env.accs)}')
# # validation(the_env.val_batch)

validation(the_env.val_batch)

def show(im):
  plt.figure()
  plt.imshow(im)
  plt.axis('off')

for im in tf.unstack(the_env._cur_batch[0]):
  show(tf.clip_by_value(im, 0.,1.))

tmp = [0 for _ in range(19)]
tmp.append( tst.reward)
reward =tst.reward

actions = np.array(actions)
select = actions[:, 0]
lx = np.array([ l[0] for l in actions[:,1]])
ly = np.array([ l[1] for l in actions[:,1]])
acts = [select, lx, ly]
actions_space = [2, 460, 700]


with tf.GradientTape() as policy_tape:
  outputs = [ De_Net(get_input(states[i:i+4])) for i in range(0, len(states), 4) ]
  t = [ np.array([tf.unstack(a)  for a in output]).T for output in outputs]
  logits = np.array([ [tf.nn.log_softmax(a) for a in o] for out in t for o in out ])

  # logits = [[ tf.nn.log_softmax(b) for b in a]  for a in logits_list]
  one_hot = [tf.one_hot(a, n) for a,n in zip(acts, actions_space)]

  log_probs = [ tf.math.reduce_sum(tf.stack(l)* o, axis=1) for l, o in zip(logits.T, one_hot)]

  policy_loss = -tf.math.reduce_mean(log_probs* tf.stop_gradient(reward), axis = 1)

policy_variables = De_Net.trainable_variables
grads = policy_tape.gradient(policy_loss, policy_variables)
De_Net.optimizer.apply_gradients(zip(grads, policy_variables))

actions = [ [1, [randint(0,459), randint(0,699)]] for _ in range(4)]
rews = []
actions_h = []
# while True:
obs = []

for action in actions:
  timestep = environment.step(action)
  obs.append(timestep.observation)
  rews.append(timestep.reward)
  actions_h.append(action)
#     if timestep.reward >0 :
#       break
  
#   if timestep.reward >0 :
#     break    
#   actions = get_action(De_Net, np.array(obs))

observations = np.array(obs)
deep_inp = observations[:,-128:].reshape(4,1,128)
stats = observations[:,:7].reshape(4,1,7)
deep_inp_tensor = tf.convert_to_tensor(deep_inp)
stats_tensor = tf.convert_to_tensor(stats)

l, a = De_Net.fit([deep_inp_tensor, stats_tensor], steps_per_epoch=1, epochs=1)

class ActionNet(network.Network):

  def __init__(self, input_tensor_spec, output_tensor_spec, model=De_Net):
    super(ActionNet, self).__init__(
        input_tensor_spec=input_tensor_spec,
        state_spec=(),
        name='ActionNet')
    self._output_tensor_spec = output_tensor_spec   
    self._model=model
    self.location = [random(),random()]

  def call(self, observations, step_type, network_state):
    del step_type
    # TODO observations as tensors
    deep_inp = tf.convert_to_tensor(observations[-128:].reshape(1,1,128))
    stats = tf.convert_to_tensor(np.append(observations[:5],self.location).reshape(1,1,7))
    
    output = self._model([deep_inp, stats]).numpy()[0]
    to_crop = 1 if output[0]>0 else 0
    self.location = output[1:]
    patches = 
    actions = [to_crop, patches]

    # Scale and shift actions to the correct range if necessary.
    return actions, network_state

"""# Convolutional Custom Model"""

from tensorflow.keras import Model
from tensorflow.keras.layers import Input, Dense, Dropout, Conv2D, MaxPool2D,GlobalMaxPool2D,Flatten
from tensorflow.keras.layers.experimental.preprocessing import Normalization

inputs = Input((400,400,3), name='input')

x = Conv2D(32, 3, strides=2)(inputs)
x = Conv2D(32, 3)(x)
x = Conv2D(32, 3)(x)
x = MaxPool2D(3, 1)(x)
x = Dropout(0.25)(x)

x = Conv2D(64, 5)(x)
x = Conv2D(64, 5)(x)
x = MaxPool2D(5, 1)(x)
x = Dropout(0.25)(x)

x = Conv2D(128, 7)(x)
x = MaxPool2D(7, 1)(x)
x = Conv2D(256, 7)(x)
x = GlobalMaxPool2D()(x)

x = Dropout(0.25)(x)
x = Dense(256)(x)
x = Dropout(0.25)(x)

output = Dense(8, activation='softmax', name='output')(x)

model = Model(inputs=[inputs], outputs=[output])

model.summary()
model.compile(optimizer=tf.optimizers.Adam(), loss=tf.keras.losses.sparse_categorical_crossentropy , metrics=["accuracy"])

history = model.fit(train_ds, epochs=5, steps_per_epoch=steps, validation_data=val_ds, validation_steps=val_size//batch_size)

history = model.fit(train_ds, epochs=15,initial_epoch=10, steps_per_epoch=steps, validation_data=val_ds, validation_steps=val_size// batch_size)

"""# Evaluate"""

# plot diagnostic learning curves
def summarize_diagnostics(history):
	plt.figure(figsize=(16, 8))
	plt.suptitle('Training Curves')
	# plot loss
	plt.subplot(121)
	plt.title('Cross Entropy Loss')
	plt.plot(history.history['loss'], color='blue', label='train')
	plt.plot(history.history['val_loss'], color='orange', label='val')
	plt.legend(loc='upper right')
	# plot accuracy
	plt.subplot(122)
	plt.title('Classification Accuracy')
	plt.plot(history.history['sparse_categorical_accuracy'], color='blue', label='train')
	plt.plot(history.history['val_sparse_categorical_accuracy'], color='orange', label='val')
	plt.legend(loc='lower right')
	return plt
 
# print test set evaluation metrics
def model_evaluation(test_ds, model, evaluation_steps):
	print('\nTest set evaluation metrics')
	loss0,accuracy0 = model.evaluate(test_ds, steps = evaluation_steps)
	print("loss: {:.2f}".format(loss0))
	print("accuracy: {:.2f}".format(accuracy0))
	return loss0,accuracy0

# the addition we made here is that the plots may be saved to a given directory
def model_report(test_ds, model, history, evaluation_steps = 10, direc = None):
	plt = summarize_diagnostics(history)
	if direc != None:
		plt.savefig(direc)
	plt.show()
	return model_evaluation(test_ds, model, evaluation_steps)

model_report(test_ds, hope_model, history, evaluation_steps=test_size//batch_size)