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 = 1
output_depth = 1
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv01 = nn.Conv2d(input_depth, 4, kernel_size=(3, 3), padding='same', bias=True)
        self.conv02 = nn.Conv2d(4, 4, kernel_size=(3, 3), padding='same', bias=True)
        self.conv03 = nn.Conv2d(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 = self.conv02(x)
        y = self.conv03(x)
        return y

# Here we assume the the NN was design for the whole frame (here 16x16).
patch_size = 16      # `patch_size`, defining the frame size, can be changed.
receptive_field = 3  # property of the model
border = 0           # no overlap for frame base by default
size = patch_size + border*2

# processing per block
bleft = receptive_field
bright = receptive_field
btop = receptive_field
bbottom = receptive_field

sublock_size = 4 # `subblock_size`, defining the block size, can be changed.
x0 = 5           # `x0` and `y0`, defining the location, can be changed.
y0 = 4

# sanity check: in principle, block on border is OK but must be filled with 0
if y0 + border + sublock_size > size:
    raise ValueError('The sub-block goes out of the bottom bound of the frame.')
if x0 + border + sublock_size > size:
    raise ValueError('The sub-block goes out of the right bound of the frame.')
if y0 < btop:
    raise ValueError('`y0` is strictly smaller than `btop`.')
if x0 < bleft:
    raise ValueError('`x0` is strictly smaller than `bleft`.')
if bleft < receptive_field or bright < receptive_field or btop < receptive_field or bbottom < receptive_field:
    raise ValueError('The receptive filed must be considered during the subblock-based inference.')

# Create the model.
model = Model()

input_frame = np.linspace(0.,1,size*size*input_depth).reshape((1,input_depth,size,size)).astype(np.float32) 
inputs_torch = [ torch.from_numpy(input_frame)]
inputs_torch[0].requires_grad=True
output = model(inputs_torch).detach().numpy()

# Extract result of the subblock.
output_frame_cropped = output[0,:,y0+border:y0+border+sublock_size,x0+border:x0+border+sublock_size]

# Compute for subblock.
input_subblock = input_frame[:,:,y0-btop+border:y0+sublock_size+bbottom+border,  x0-bleft+border:x0+sublock_size+bright+border]
inputs_torch = [ torch.from_numpy(input_subblock)]
inputs_torch[0].requires_grad=True
output_subblock = model(inputs_torch).detach().numpy()
output_subblock_cropped = output_subblock[0,:,btop:btop+sublock_size,bleft:bleft+sublock_size]

print("== Original frame ==")
print("input {} (margin={})\n{}".format(input_frame.shape, border, input_frame))
print("output {}\n{}\n".format(output.shape,output))

print("== subblock at {},{} of size {}x{} ==".format(x0,y0,sublock_size,sublock_size))
print("input {} (margin={})\n{}".format(input_subblock.shape,receptive_field,input_subblock))
print("output {} cropped to {}\n{}\n".format( output.shape,output_subblock_cropped.shape,output_subblock_cropped))
print("diff output subblock - part of frame\n{}\n".format(output_subblock_cropped-output_frame_cropped))



