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)

    dummy_input = torch.randn(1, 6, 100, 100)
    torch.onnx.export(model, dummy_input, MODEL_PATH + ".onnx", verbose=False, opset_version=11)


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 = ()
        i = 0
        piclen = len(b['input_arrays'])
        for input_array, known_array in zip(b['input_arrays'], b['known_arrays']):
            input_array = DataLoader.preprocess(input_array, precision=np.float32)

            input_array = np.expand_dims(input_array, 0)
            known_array = np.expand_dims(known_array, 0)
            input_tensor = torch.cat((torch.from_numpy(input_array), torch.from_numpy(known_array)), 1)
            out = model(input_tensor)
            out = DataLoader.postprocess(out.cpu().detach().numpy())

            rest = out * (1 - known_array)
            rest = rest[1 - known_array > 0]

            outarr = (*outarr, rest)

            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)