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

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

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
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

patch_size = 16  # original block size
border     = 2   # overlap value
size = patch_size + 2*border
# property of the network
receptive_field=5 
dT=dict()
dB=dict()
dL=dict()
dR=dict()

dT["top-left"]=0
dB["top-left"]=0
dL["top-left"]=0
dR["top-left"]=0

dT["top-right"]=0
dB["top-right"]=0
dL["top-right"]=1
dR["top-right"]=0

dT["bottom-left"]=1
dB["bottom-left"]=0
dL["bottom-left"]=0
dR["bottom-left"]=0

dT["bottom-right"]=1
dB["bottom-right"]=0
dL["bottom-right"]=1
dR["bottom-right"]=0

sublock_size = patch_size//2 # subblock size
bb = "bottom-right" # parts to compare

# Create the model.
model = Model()

input_org = np.linspace(0.,1,size*size*input_depth).reshape((1,input_depth,size,size)).astype(np.float32)
inputs_torch = [ torch.from_numpy(input_org)]
inputs_torch[0].requires_grad=True
output_org = model(inputs_torch).detach().numpy()
print("== Original block ==")
print("input {} (margin={})\n{}".format(input_org.shape, border, input_org))
print("output {} \n{}\n".format(output_org.shape,output_org))

sublock_size=patch_size//2

if bb == "top-left":
 x0 = border
 y0 = border
 bleft = border + dL[bb]
 bright = receptive_field + dR[bb]
 btop = border + dT[bb]
 bbottom = receptive_field + dB[bb]
elif bb == "top-right":
 x0 = border+sublock_size
 y0 = border
 bleft = receptive_field + dL[bb]
 bright = border + dR[bb]
 btop = border + dT[bb]
 bbottom = receptive_field+ dB[bb]
elif bb == "bottom-left":
 x0 = border
 y0 = border+sublock_size
 bleft = border + dL[bb]
 bright = receptive_field + dR[bb] 
 btop = receptive_field + dT[bb]
 bbottom = border + dB[bb]
elif bb == "bottom-right":
 x0 = border+sublock_size
 y0 = border+sublock_size
 bleft = receptive_field + dL[bb]
 bright = border + dR[bb] + dR[bb]
 btop = receptive_field+ dT[bb]
 bbottom = border + dB[bb]

output_org_cropped = output_org[0,:,y0:y0+sublock_size,x0:x0+sublock_size]

input_subblock = input_org[:,:,y0-btop:y0+sublock_size+bbottom,  x0-bleft:x0+sublock_size+bright]
inputs_torch = [ torch.from_numpy(input_subblock)]
inputs_torch[0].requires_grad=True
output_sublock = model(inputs_torch).detach().numpy()
output_subblock_cropped = output_sublock[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 {} \n{}\n".format(output_org.shape,output_org))
print("output org {} cropped to {}\n{}\n".format(bb,output_org_cropped.shape,output_org_cropped))

print("== subblock {} ==".format(bb))
print("input {}\n{}".format(input_subblock.shape,input_subblock))
print("output {} cropped to {}\n{}\n".format( output_sublock.shape,output_subblock_cropped.shape,output_subblock_cropped))
print("diff output subblock - output cropped\n{}\n".format(output_subblock_cropped-output_org_cropped))



