"""
Support log functions

TODO: log model using mlflow.pytorch in parallel / addition to checkpointing
"""

import numpy as np
import h5py
import os
import argparse

import torch
import torchvision.utils as vutils
import pytorch_lightning as pl


class ImgCB(pl.Callback):
    def __init__(self, **kwargs):
        parser = ImgCB.add_argparse_args()
        for action in parser._actions:
            if action.dest in kwargs:
                action.default = kwargs[action.dest]
        args = parser.parse_args([])
        self.__dict__.update(vars(args))

    @staticmethod
    def add_argparse_args(parent_parser=None):
        parser = argparse.ArgumentParser(
            prog='ImgCB',
            usage=ImgCB.__doc__,
            parents=[parent_parser] if parent_parser is not None else [],
            add_help=False)

        parser.add_argument('--img_ranges', default=[1300, 1800], nargs='*', help='Scaling range on output image, either pair, or set of pairs')
        parser.add_argument('--err_ranges', default=[0, 50], nargs='*', help='Scaling range on error images, either pair, or set of pairs')

        return parser

    def log_images(self, mfl_logger, y, z, prefix):
        img_ranges = tuple(self.img_ranges)
        err_ranges = tuple(self.err_ranges)
        # 
        for i in range(y.shape[1]):
            if y.shape[1] > 1:
                tag = f'_{i}_'

                if len(self.img_ranges) > 2:
                    img_ranges = tuple(self.img_ranges[2*i, 2*i + 1])
                if len(self.err_ranges) > 2:
                    err_ranges = tuple(self.err_ranges[2*i, 2*i + 1])
            else:
                tag = ''

            mfl_logger.experiment.log_image(
                mfl_logger.run_id,
                (np.array(vutils.make_grid(
                    y[:, [i], ...].detach(),
                    normalize=True, value_range=img_ranges, nrow=6).cpu())[0, ...] * 255.).astype(np.int),
                prefix + tag + '_labels.png')

            mfl_logger.experiment.log_image(
                mfl_logger.run_id,
                (np.array(vutils.make_grid(
                    z[:, [i], ...].detach(),
                    normalize=True, value_range=img_ranges, nrow=6).cpu())[0, ...] * 255.).astype(np.int),
                prefix + tag + '_outputs.png')

            mfl_logger.experiment.log_image(
                mfl_logger.run_id,
                (np.array(vutils.make_grid(
                    torch.abs(y[:, [i], ...].detach() - z[:, [i], ...].detach()),
                    normalize=True, value_range=err_ranges, nrow=6).cpu())[0, ...] * 255.).astype(np.int),
                prefix + tag + '_errors.png')

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        if batch_idx == 0:
            with torch.no_grad():
                x, y = batch

                if pl_module.hparams.rand_output_crop:
                    x = x[..., :-pl_module.hparams.rand_output_crop, :]
                    y = y[..., :-pl_module.hparams.rand_output_crop * 2, :]

                z = pl_module(x.to(pl_module.device))

                if isinstance(z, tuple) or isinstance(z, list):
                    z = z[0]

                self.log_images(pl_module.logger, y.to(pl_module.device), z, 'train_')

    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
        if batch_idx == 0:
            with torch.no_grad():
                x, y = batch

                if pl_module.hparams.rand_output_crop:
                    x = x[..., :-pl_module.hparams.rand_output_crop, :]
                    y = y[..., :-pl_module.hparams.rand_output_crop * 2, :]

                z = pl_module(x.to(pl_module.device))

                if isinstance(z, tuple) or isinstance(z, list):
                    z = z[0]

                self.log_images(pl_module.logger, y.to(pl_module.device), z, 'validate_')


class TestLogger(pl.Callback):
    """
    pytorch_lightning Data saving logger for testing output
    Warning !!! : this function is not multi GPU / multi device safe -- only run on a single gpu / device
    """
    def __init__(self, fname: str = 'output.h5'):
        self.fname = fname

    def on_test_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
        with h5py.File(self.fname, 'a') as f:
            f[f'batch_{batch_idx:05}'] = outputs.to('cpu').numpy()
            if len(batch) > 1:
                f[f'labels_{batch_idx:05}'] = batch[1].to('cpu').numpy()