import pickle

import numpy as np
import torch
from PIL import Image

import Compress
import DataLoader
import ex4
from ImageImpaint import get_train_device
from netio import load_model, eval_evalset, write_to_pickle


def apply_model(filepath: str):
    device = get_train_device()

    img = Image.open(filepath)
    model = load_model()
    model.to(device)

    pic = DataLoader.crop_image(img)
    pic = DataLoader.preprocess(pic, precision=np.float32)
    pic = ex4.ex4(pic, (5, 5), (4, 4))[0]
    Image.fromarray((np.transpose(DataLoader.postprocess(pic), (1, 2, 0)).astype(np.uint8))).save("filename_grid.jpg")
    out = model(torch.from_numpy(pic).to(device))
    out = out.cpu().detach().numpy()
    out = DataLoader.postprocess(out)
    out = np.transpose(out, (1, 2, 0))
    im = Image.fromarray(out)
    im.save("filename.jpg", format="jpeg")

def test():
    # read the provided testing pickle file
    print("Generating pickle file with privided test data")
    PICKEL_PATH = "test"

    model = load_model()
    model.eval()

    loader,_ = DataLoader.get_image_loader("training/", np.float32)
    outarr = np.zeros(dtype=np.uint8, shape=(8663, 3, 100, 100))
    targetarr = np.zeros(dtype=np.uint8, shape=(8663, 3, 100, 100))

    i = 0
    for input, target in loader:
        out = model(input)
        out = DataLoader.postprocess(out.cpu().detach().numpy())
        outarr[i] = out
        targetarr[i] = DataLoader.postprocess(target.cpu().detach().numpy())
        print(f'\rApplying model [{i}/{len(loader)}]', end='')
        i += 1
        if i==8663:
            break
    write_to_pickle(PICKEL_PATH + "_pred.pkl", list(outarr))
    # compress the generated pickle arr
    Compress.compress(PICKEL_PATH + "_pred.pkl")

    write_to_pickle(PICKEL_PATH + "_target.pkl", list(targetarr))
    # compress the generated pickle arr
    Compress.compress(PICKEL_PATH + "_target.pkl")

if __name__ == '__main__':
    # apply_model("training/000/000017.jpg")
    eval_evalset()
    # test()