# Path: selfsup/multi/methods/base.py
# class Method:
#     def __init__(self, name, batch_size):
#         self.name = name
#         self.batch_size = batch_size
# 
#     def input_size(self):
#         return 0
# 
#     def batch_loader(self):
#         raise NotImplemented("Abstract base class")
# 
#     def build_network(network, y, phase_test):
#         raise NotImplemented("Abstract base class")
# 
#     @property
#     def basenet_settings():
#         return {}
# 
#     def make_path(self, name, ext, iteration):
#         return f'img/{self.name}_{name}_{iteration}.png'


import selfsup
import tensorflow as tf
import os
import deepdish as dd
import numpy as np
import itertools
import selfsup.jigsaw
from .base import Method
from collections import OrderedDict
    patches = []

    for i, j in dd.multi_range(size, size):
        #tf.slice(x, [

        M = WINDOW_SIZE - patch_size + 1
        assert M > 0, f'Jigsaw: Window size ({WINDOW_SIZE}) and patch size ({patch_size}) not compatible'
        limit = np.array([1, M, M, 1])
        offset = np.array([0, i * WINDOW_SIZE, j * WINDOW_SIZE, 0]) + tf.random_uniform(
                [4], dtype=tf.int32,
                maxval=M,
                ) % limit
        patch = tf.slice(x, offset, [N, patch_size, patch_size, C])
        patches.append(patch)

    patches1 = tf.stack(patches, axis=1)
    xyz = np.arange(batch_size)[:, np.newaxis] * size**2 + (perm - 1)
    #import ipdb
    #ipdb.set_trace()
    perm0 = tf.reshape(xyz, [-1])

    patches_flat = tf.reshape(patches1, [-1] + patches1.get_shape().as_list()[2:])
    #import ipdb
    ##ipdb.set_trace()
    patches2 = tf.gather(patches_flat, perm0)

    #return tf.reshape(patches2, [-1, PATCH_SIZE, PATCH_SIZE, C])
    return patches2



# Path: selfsup/multi/methods/base.py
# class Method:
#     def __init__(self, name, batch_size):
#         self.name = name
#         self.batch_size = batch_size
# 
#     def input_size(self):
#         return 0
# 
#     def batch_loader(self):
#         raise NotImplemented("Abstract base class")
# 
#     def build_network(network, y, phase_test):
#         raise NotImplemented("Abstract base class")
# 
#     @property
#     def basenet_settings():
#         return {}
# 
#     def make_path(self, name, ext, iteration):
#         return f'img/{self.name}_{name}_{iteration}.png'


# Path: selfsup/multi/methods/jigsaw.py
import selfsup
import tensorflow as tf
import os
import deepdish as dd
import numpy as np
import itertools
import selfsup.jigsaw
from .base import Method
from collections import OrderedDict

    patches = []

    for i, j in dd.multi_range(size, size):
        #tf.slice(x, [

        M = WINDOW_SIZE - patch_size + 1
        assert M > 0, f'Jigsaw: Window size ({WINDOW_SIZE}) and patch size ({patch_size}) not compatible'
        limit = np.array([1, M, M, 1])
        offset = np.array([0, i * WINDOW_SIZE, j * WINDOW_SIZE, 0]) + tf.random_uniform(
                [4], dtype=tf.int32,
                maxval=M,
                ) % limit
        patch = tf.slice(x, offset, [N, patch_size, patch_size, C])
        patches.append(patch)

    patches1 = tf.stack(patches, axis=1)
    xyz = np.arange(batch_size)[:, np.newaxis] * size**2 + (perm - 1)
    #import ipdb
    #ipdb.set_trace()
    perm0 = tf.reshape(xyz, [-1])

    patches_flat = tf.reshape(patches1, [-1] + patches1.get_shape().as_list()[2:])
    #import ipdb
    ##ipdb.set_trace()
    patches2 = tf.gather(patches_flat, perm0)

    #return tf.reshape(patches2, [-1, PATCH_SIZE, PATCH_SIZE, C])
    return patches2



class Jigsaw(Method):