import numpy as np
from PIL import Image
from ISR.models import RDN

img_name = './data/meerkat.png'
img = Image.open(img_name)

print(img.size)

scale = 2
process_box = 120

out_img = Image.new('RGB', (img.size[0] * scale, img.size[1] * scale))

rdn = RDN(arch_params={'C':6, 'D':20, 'G':64, 'G0':64, 'x':2})
rdn.model.load_weights('weights/sample_weights/rdn-C6-D20-G64-G064-x2/ArtefactCancelling/rdn-C6-D20-G64-G064-x2_ArtefactCancelling_epoch219.hdf5')

xNum = img.size[0] // process_box
yNum = img.size[1] // process_box
xDiv = img.size[0] % process_box
yDiv = img.size[1] % process_box

for i in range(0, xNum):
    for j in range(0, yNum):
        left = i * process_box
        top = j * process_box
        right = process_box + left
        bottom = process_box + top
        print(top, left, "->", bottom, right)
        cropped_example = img.crop((left, top, right, bottom))

        lr_img = np.array(cropped_example)

        sr_img = rdn.predict(lr_img)
        new_img = Image.fromarray(sr_img)

        out_img.paste(new_img, (i * process_box * scale, j * process_box * scale))

top = yNum * process_box
bottom = img.size[1]
for i in range(0, xNum):
    left = i * process_box
    right = process_box + left
    print(top, left, "->", bottom, right)
    cropped_example = img.crop((left, top, right, bottom))
    lr_img = np.array(cropped_example)

    sr_img = rdn.predict(lr_img)
    new_img = Image.fromarray(sr_img)

    out_img.paste(new_img, (left * scale, top * scale))

left = xNum * process_box
right = img.size[0]
for j in range(0, yNum):
    top = j * process_box
    bottom = process_box + top

    print(top, left, "->", bottom, right)
    cropped_example = img.crop((left, top, right, bottom))
    lr_img = np.array(cropped_example)

    sr_img = rdn.predict(lr_img)
    new_img = Image.fromarray(sr_img)

    out_img.paste(new_img, (left * scale, top * scale))

left = xNum * process_box
right = img.size[0]
top = yNum * process_box
bottom = img.size[1]

print(top, left, "->", bottom, right)
cropped_example = img.crop((left, top, right, bottom))
lr_img = np.array(cropped_example)

sr_img = rdn.predict(lr_img)
new_img = Image.fromarray(sr_img)

out_img.paste(new_img, (left * scale, top * scale))

out_img.save('outimg.png')