import torch

from utils.flow_viz import flow_tensor_to_image
from .visualization import viz_depth_tensor


class Logger:
    def __init__(self, lr_scheduler,
                 summary_writer,
                 summary_freq=100,
                 start_step=0,
                 img_mean=None,
                 img_std=None,
                 ):
        self.lr_scheduler = lr_scheduler
        self.total_steps = start_step
        self.running_loss = {}
        self.summary_writer = summary_writer
        self.summary_freq = summary_freq

        self.img_mean = img_mean
        self.img_std = img_std

    def print_training_status(self, mode='train', is_depth=False):
        if is_depth:
            print('step: %06d \t loss: %.3f' % (self.total_steps, self.running_loss['total_loss'] / self.summary_freq))
        else:
            print('step: %06d \t epe: %.3f' % (self.total_steps, self.running_loss['epe'] / self.summary_freq))

        for k in self.running_loss:
            self.summary_writer.add_scalar(mode + '/' + k,
                                           self.running_loss[k] / self.summary_freq, self.total_steps)
            self.running_loss[k] = 0.0

    def lr_summary(self):
        lr = self.lr_scheduler.get_last_lr()[0]
        self.summary_writer.add_scalar('lr', lr, self.total_steps)

    def add_image_summary(self, img1, img2, flow_preds=None, flow_gt=None, mode='train',
                          is_depth=False,
                          ):
        if self.total_steps % self.summary_freq == 0:
            if is_depth:
                img1 = self.unnormalize_image(img1.detach().cpu())  # [3, H, W], range [0, 1]
                img2 = self.unnormalize_image(img2.detach().cpu())

                concat = torch.cat((img1, img2), dim=-1)  # [3, H, W*2]

                self.summary_writer.add_image(mode + '/img', concat, self.total_steps)
            else:
                img_concat = torch.cat((img1[0].detach().cpu(), img2[0].detach().cpu()), dim=-1)
                img_concat = img_concat.type(torch.uint8)  # convert to uint8 to visualize in tensorboard

                flow_pred = flow_tensor_to_image(flow_preds[-1][0])
                forward_flow_gt = flow_tensor_to_image(flow_gt[0])
                flow_concat = torch.cat((torch.from_numpy(flow_pred),
                                         torch.from_numpy(forward_flow_gt)), dim=-1)

                concat = torch.cat((img_concat, flow_concat), dim=-2)

                self.summary_writer.add_image(mode + '/img_pred_gt', concat, self.total_steps)

    def add_depth_summary(self, depth_pred, depth_gt, mode='train'):
        # assert depth_pred.dim() == 2  # [H, W]
        if self.total_steps % self.summary_freq == 0 or 'val' in mode:
            pred_viz = viz_depth_tensor(depth_pred.detach().cpu())  # [3, H, W]
            gt_viz = viz_depth_tensor(depth_gt.detach().cpu())

            concat = torch.cat((pred_viz, gt_viz), dim=-1)  # [3, H, W*2]

            self.summary_writer.add_image(mode + '/depth_pred_gt', concat, self.total_steps)

    def unnormalize_image(self, img):
        # img: [3, H, W], used for visualizing image
        mean = torch.tensor(self.img_mean).view(3, 1, 1).type_as(img)
        std = torch.tensor(self.img_std).view(3, 1, 1).type_as(img)

        out = img * std + mean

        return out

    def push(self, metrics, mode='train', is_depth=False, ):
        self.total_steps += 1

        self.lr_summary()

        for key in metrics:
            if key not in self.running_loss:
                self.running_loss[key] = 0.0

            self.running_loss[key] += metrics[key]

        if self.total_steps % self.summary_freq == 0:
            self.print_training_status(mode, is_depth=is_depth)
            self.running_loss = {}

    def write_dict(self, results):
        for key in results:
            tag = key.split('_')[0]
            tag = tag + '/' + key
            self.summary_writer.add_scalar(tag, results[key], self.total_steps)

    def close(self):
        self.summary_writer.close()