import torch
import torch.nn as nn
import numpy as np

torch.manual_seed(42)
np.set_printoptions(3,linewidth=np.inf)

def weights_init(m):
    if isinstance(m, nn.Conv2d) or  isinstance(m, nn.ConvTranspose2d):
        torch.nn.init.uniform_(m.weight,-0.1,0.2)
        torch.nn.init.uniform_(m.bias,-0.1,0.2)

input_depth = 5
output_depth = 1
nb_kernels_intermediate = 8
if nb_kernels_intermediate % 4 != 0:
    raise ValueError('`nb_kernels_intermediate` is not divisible by 4.')
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv01 = nn.Conv2d(input_depth, nb_kernels_intermediate, kernel_size=(3, 3), padding='same', bias=True)
        self.conv02 = nn.Conv2d(nb_kernels_intermediate, nb_kernels_intermediate, kernel_size=(3, 3), padding='same', bias=True)
        self.conv03 = nn.Conv2d(nb_kernels_intermediate//4, output_depth, kernel_size=(3, 3), padding='same', bias=True)
        self.apply(weights_init)

    def forward(self, inputs):
        input, = inputs
        x = self.conv01(input)
        x = nn.MaxPool2d(2)(x)
        x = self.conv02(x)
        x = nn.PixelShuffle(2)(x)
        y = self.conv03(x)
        return y

# Here we assume the the NN was design for 8x8 blocks and we apply a process on a superblock 16x16.
patch_size = 16 # super block size (orgiginal block will be half the size)
shift = 1       # shift to apply on original block to demonstrate that the NN is NOT translation equivariant
receptive_field = 5 # property of the model
border = receptive_field # overlap should be == to the RF since supreblock is possible
size = patch_size + border*2 + shift
sublock_size=patch_size//2 # original block size

bb = "top-left"
if bb == "top-left":
 x0 = receptive_field
 y0 = receptive_field
elif bb == "top-right":
 x0 = receptive_field+sublock_size
 y0 = receptive_field
elif bb == "bottom-left":
 x0 = receptive_field
 y0 = receptive_field+sublock_size
elif bb == "bottom-right":
 x0 = receptive_field+sublock_size
 y0 = receptive_field+sublock_size
else:
    raise ValueError("`bb` is not valid.")
bleft = receptive_field 
bright = receptive_field 
btop = receptive_field
bbottom = receptive_field 

if receptive_field + patch_size > size:
    raise ValueError('The first crop will go out of the bounds of the output superblock.')
if y0 + sublock_size > size:
    raise ValueError('The second crop will go out of the bounds of the output superblock.')
if x0 + sublock_size > size:
    raise ValueError('The second crop will go out of the bounds of the output superblock.')
if y0 < btop:
    raise ValueError('The third crop will go out of the bounds of the input superblock.')
if y0 + sublock_size + bbottom > size:
    raise ValueError('The third crop will go out of the bounds of the input superblock.')
if x0 < bleft:
    raise ValueError('The third crop will go out of the bounds of the input superblock.')
if x0 + sublock_size + bright > size:
    raise ValueError('The third crop will go out of the bounds of the input superblock.')
if y0 + shift < btop:
    raise ValueError('The fourth crop will go out of the bounds of the input superblock.')
if y0 + sublock_size + bbottom + shift > size:
    raise ValueError('The fourth crop will go out of the bounds of the input superblock.')

# Create the model.
model = Model()

input_superblock = np.linspace(0.,1,size*size*input_depth).reshape((1,input_depth,size,size)).astype(np.float32) 
inputs_torch = [torch.from_numpy(input_superblock)]
inputs_torch[0].requires_grad = True
output_superblock = model(inputs_torch).detach().numpy()
output_superblock_cropped_full = output_superblock[0,:,receptive_field:receptive_field+patch_size,receptive_field:receptive_field+patch_size]
output_superblock_cropped = output_superblock[0,:,y0:y0+sublock_size,x0:x0+sublock_size]

# Compute for original block.
input_org = input_superblock[:,:,y0-btop:y0+sublock_size+bbottom,  x0-bleft:x0+sublock_size+bright]
inputs_torch = [ torch.from_numpy(input_org)]
inputs_torch[0].requires_grad=True
output_org= model(inputs_torch).detach().numpy()
output_org_cropped = output_org[0,:,btop:btop+sublock_size,bleft:bleft+sublock_size]

print("== Original block ==")
print("input {} (margin={})\n{}".format(input_org.shape, border, input_org))
print("output {} cropped to {}\n{}\n".format(output_org.shape,output_org_cropped.shape,output_org_cropped))

print("== Superblock ==")
print("input {} (margin={})\n{}".format(input_superblock.shape,border,input_superblock))
print("output {} cropped to {}\n{}\n".format( output_superblock.shape,output_superblock_cropped_full.shape,output_superblock_cropped_full))
print("output org block - {} part of superblock output\n{}\n".format(bb,output_org_cropped-output_superblock_cropped))

# Compute for a sublock translated.
input_org_shift = input_superblock[:,:,y0-btop+shift:y0+sublock_size+bbottom+shift,x0-bleft:x0+sublock_size+bright]
inputs_torch = [ torch.from_numpy(input_org_shift)]
inputs_torch[0].requires_grad=True
output_org_shift = model(inputs_torch).detach().numpy()
output_org_shift_cropped = output_org_shift[0,:,btop:btop+sublock_size,bleft:bleft+sublock_size]
output_org_shift_cropped = np.roll(output_org_shift_cropped,shift,axis=1)

print("== Original block shifted ==")
print("input {} (margin={}) shifted by 1\n{}".format(input_org_shift.shape, border, input_org_shift))
print("output {} shifted cropped to {}\n{}\n".format(output_org_shift.shape,output_org_shift_cropped.shape,output_org_shift_cropped))
print("output org block - {} shifted output (first {} lines not relevant due to the rolling)\n{}\n".format(bb,shift,output_org_cropped-output_org_shift_cropped))


