#!/usr/bin/env python3
from dataset import SIDDValData
from model import UNetD
import megengine.data as data
from utils import batch_PSNR
from tqdm import tqdm
import argparse
import pickle
import megengine


def test(args):
    valid_dataset = SIDDValData(args.data)
    valid_sampler = data.SequentialSampler(
        valid_dataset, batch_size=1, drop_last=False
    )
    valid_dataloader = data.DataLoader(
        valid_dataset,
        sampler=valid_sampler,
        num_workers=8,
    )
    model = UNetD(3)
    with open(args.checkpoint, "rb") as f:
        state = pickle.load(f)
    model.load_state_dict(state["state_dict"])
    model.eval()

    def valid_step(image, label):
        pred = model(image)
        pred = image - pred
        psnr_it = batch_PSNR(pred, label)
        return psnr_it

    def valid(func, data_queue):
        psnr_v = 0.
        for step, (image, label) in tqdm(enumerate(data_queue)):
            image = megengine.tensor(image)
            label = 
megengine.tensor(label)