import numpy as np
import torch


class DepthEvaluation:
    def __init__(self, cfg):

        self.metrics = ('AbsRel', 'SqrRel', 'RMSE', 'RMSElog', 'SILog', 'a1', 'a2', 'a3')

        self.min_depth = cfg.has('min_depth',   0.0)
        self.max_depth = cfg.has('max_depth', 200.0)

        self.crop = cfg.has('crop', '')
        self.scale_output = cfg.has('scale_output', 'resize')

        self.post_process = cfg.has('post_process', False)
        self.median_scaling = cfg.has('median_scaling', False)
        self.shift_scaling = cfg.has('shift_scaling', False)
        self.min_max = cfg.has('min_max', False)
        self.valid_threshold = cfg.has('valid_threshold', 0)
        self.sparse_predictions = cfg.has('sparse_predictions', False)

    def compute(self, gt, pred, median_scaling=False, shift_scaling=False, min_max=False, mask=None):
        from vidar.metrics.utils import create_crop_mask, scale_output
        assert not median_scaling or not shift_scaling
        # Match predicted depth map to ground-truth resolution
        pred = scale_output(pred, gt, self.scale_output)
        # Create crop mask if requested
        crop_mask = create_crop_mask(self.crop, gt)
        # For each batch sample
        metrics = []
        for i, (pred_i, gt_i) in enumerate(zip(pred, gt)):

            # Squeeze GT and PRED
            gt_i, pred_i = torch.squeeze(gt_i), torch.squeeze(pred_i)
            mask_i = torch.squeeze(mask[i]) if mask is not None else None

            # Keep valid pixels (min/max depth and crop)
            valid = (gt_i > self.min_depth) & (gt_i < self.max_depth)
            # Remove invalid predicted pixels as well
            if self.sparse_predictions:
                valid = valid & (pred_i > 0) & (~torch.isnan(pred_i))
            # Apply crop mask if requested
            valid = valid & crop_mask.bool() if crop_mask is not None else valid
            # Apply provided mask if available
            valid = valid & mask_i.bool() if mask is not None else valid

            # Invalid evaluation
            if valid.sum() <= self.valid_threshold:
                metrics.append([-1.] * 8)
                continue

            # Keep only valid pixels
            gt_i, pred_i = gt_i[valid], pred_i[valid]
            # GT median scaling if needed
            if shift_scaling:
                from vidar.utils.depth import scale_and_shift_pred
                _, _, pred_i = scale_and_shift_pred(pred_i, gt_i)
            if median_scaling:
                pred_i = pred_i * torch.median(gt_i) / torch.median(pred_i)
            if min_max:
                min_, max_ = gt_i.min(), gt_i.max()
                pred_i = pred_i * (max_ - min_) + min_
            # Clamp PRED depth values to min/max values
            pred_i = pred_i.clamp(self.min_depth, self.max_depth)

            # Calculate depth metrics

            thresh = torch.max((gt_i / pred_i), (pred_i / gt_i))
            a1 = (thresh < 1.25).float().mean()
            a2 = (thresh < 1.25 ** 2).float().mean()
            a3 = (thresh < 1.25 ** 3).float().mean()

            diff_i = gt_i - pred_i
            abs_rel = torch.mean(torch.abs(diff_i) / gt_i)
            sq_rel = torch.mean(diff_i ** 2 / gt_i)
            rmse = torch.sqrt(torch.mean(diff_i ** 2))
            rmse_log = torch.sqrt(torch.mean((torch.log(gt_i) - torch.log(pred_i)) ** 2))

            err = torch.log(pred_i) - torch.log(gt_i)
            silog = torch.sqrt(torch.mean(err ** 2) - torch.mean(err) ** 2) * 100

            metrics.append([abs_rel, sq_rel, rmse, rmse_log, silog, a1, a2, a3])

        # Return metrics
        return torch.tensor(metrics, dtype=gt.dtype, device=gt.device)

    def evaluate(self, gt, pred):
        return self.compute(gt, pred)