#!/bin/env python
"""Evaluate a demosaicking model."""
import argparse
import logging

import torch as th
from torch.utils.data import DataLoader

import demosaicnet


log = logging.getLogger(__name__)

class PSNR(th.nn.Module):
    def __init__(self):
        super(PSNR, self).__init__()
        self.mse = th.nn.MSELoss()
    def forward(self, out, ref):
        mse = self.mse(out, ref)
        return -10*th.log10(mse)

def main(args):
    """Entrypoint to the training."""

    # Load model parameters from checkpoint, if any
    meta = demosaicnet.utils.Checkpointer.load_meta(args.checkpoint_dir)
    if meta is None:
        log.warning("No checkpoint found at %s, aborting.", args.checkpoint_dir)
        return

    data = demosaicnet.Dataset(args.data, download=False,
                               mode=meta["mode"],
                               subset=demosaicnet.TEST_SUBSET)
    dataloader = DataLoader(
        data, batch_size=1, num_workers=4, pin_memory=True, shuffle=True)

    if meta["mode"] == demosaicnet.BAYER_MODE:
        model = demosaicnet.BayerDemosaick(depth=meta["depth"],
                                           width=meta["width"],
                                           pretrained=True,
                                           pad=False)
    elif meta["mode"] == demosaicnet.XTRANS_MODE:
        model = demosaicnet.XTransDemosaick(depth=meta["depth"],
                                            width=meta["width"],
                                            pretrained=True,
                                            pad=False)

    checkpointer = demosaicnet.utils.Checkpointer(args.checkpoint_dir, model, meta=meta)
    checkpointer.load_latest()  # Resume from checkpoint, if any.

    # No need for gradients
    for p in model.parameters():
        p.requires_grad = False

    mse_fn = th.nn.MSELoss()
    psnr_fn = PSNR()

    device = "cpu"
    if th.cuda.is_available():
        device = "cuda"
        log.info("Using CUDA")

    count = 0
    mse = 0.0
    psnr = 0.0
    for idx, batch in enumerate(dataloader):
        mosaic = batch[0].to(device)
        target = batch[1].to(device)
        output = model(mosaic)

        target = demosaicnet.utils.crop_like(target, output)

        output = th.clamp(output, 0, 1)

        psnr_ = psnr_fn(output, target).item()
        mse_ = mse_fn(output, target).item()

        psnr += psnr_
        mse += mse_
        count += 1

        log.info("Image %04d, PSNR = %.1f dB, MSE = %.5f", idx, psnr_, mse_)

    mse /= count
    psnr /= count

    log.info("-----------------------------------")
    log.info("Average, PSNR = %.1f dB, MSE = %.5f", psnr, mse)



if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("data", help="root directory for the demosaicnet dataset.")
    parser.add_argument("checkpoint_dir", help="directory with the model checkpoints.")
    args = parser.parse_args()
    main(args)
