import os
import pickle

import numpy as np
import torch
import Compress
import DataLoader
from Net import ImageNN

MODEL_PATH = 'impaintmodel.pt'
PICKEL_PATH = 'impaintmodel.pkl'


def save_model(model: torch.nn.Module):
    print(f"Saved raw model to {MODEL_PATH}")
    torch.save(model, MODEL_PATH)


def eval_evalset():
    # read the provided testing pickle file
    print("Generating pickle file with privided test data")
    try:
        os.unlink(PICKEL_PATH)
    except:
        pass

    model = load_model()
    model.eval()
    with open('testing/inputs.pkl', 'rb') as handle:
        b: dict = pickle.load(handle)
        outarr = np.zeros(dtype=np.uint8, shape=(len(b['input_arrays']), 3, 100, 100))
        i = 0
        piclen = len(b['input_arrays'])
        for pic in b['input_arrays']:
            pic = DataLoader.preprocess(pic, precision=np.float32)
            out = model(torch.from_numpy(pic))
            out = DataLoader.postprocess(out.cpu().detach().numpy())
            outarr[i] = out

            print(f'\rApplying model [{i}/{piclen}]', end='')
            i += 1

        write_to_pickle(PICKEL_PATH, list(outarr))
    # compress the generated pickle arr
    Compress.compress(PICKEL_PATH)


def write_to_pickle(filename: str, data):
    with open(filename, 'wb') as handle:
        pickle.dump(data, handle, protocol=pickle.HIGHEST_PROTOCOL)


def load_model():
    return torch.load(MODEL_PATH)