import base64
import math
import re
from io import BytesIO

import matplotlib.cm
import numpy as np
import torch
import torch.nn
from PIL import Image

# Compute edge magnitudes
from scipy import ndimage


class RunningAverage:
    def __init__(self):
        self.avg = 0
        self.count = 0

    def append(self, value):
        self.avg = (value + self.count * self.avg) / (self.count + 1)
        self.count += 1

    def get_value(self):
        return self.avg


def denormalize(x, device='cpu'):
    mean = torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device)
    std = torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device)
    return x * std + mean


class RunningAverageDict:
    def __init__(self):
        self._dict = None

    def update(self, new_dict):
        if self._dict is None:
            self._dict = dict()
            for key, _ in new_dict.items():
                self._dict[key] = RunningAverage()

        for key, value in new_dict.items():
            self._dict[key].append(value)

    def get_value(self):
        return {key: value.get_value() for key, value in self._dict.items()}


def colorize(value, vmin=10, vmax=1000, cmap='magma_r'):
    value = value.cpu().numpy()[0, :, :]
    invalid_mask = value == -1

    # normalize
    vmin = value.min() if vmin is None else vmin
    vmax = value.max() if vmax is None else vmax
    if vmin != vmax:
        value = (value - vmin) / (vmax - vmin)  # vmin..vmax
    else:
        # Avoid 0-division
        value = value * 0.0
    # squeeze last dim if it exists
    # value = value.squeeze(axis=0)
    cmapper = matplotlib.cm.get_cmap(cmap)
    value = cmapper(value, bytes=True)  # (nxmx4)
    value[invalid_mask] = 255
    img = value[:, :, :3]

    #     return img.transpose((2, 0, 1))
    return img


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def compute_errors(gt, pred):
    thresh = np.maximum((gt / pred), (pred / gt))
    a1 = (thresh < 1.25).mean()
    a2 = (thresh < 1.25 ** 2).mean()
    a3 = (thresh < 1.25 ** 3).mean()

    abs_rel = np.mean(np.abs(gt - pred) / gt)
    sq_rel = np.mean(((gt - pred) ** 2) / gt)

    rmse = (gt - pred) ** 2
    rmse = np.sqrt(rmse.mean())

    rmse_log = (np.log(gt) - np.log(pred)) ** 2
    rmse_log = np.sqrt(rmse_log.mean())

    err = np.log(pred) - np.log(gt)
    silog = np.sqrt(np.mean(err ** 2) - np.mean(err) ** 2) * 100

    log_10 = (np.abs(np.log10(gt) - np.log10(pred))).mean()
    return dict(
        a1=a1,
        a2=a2,
        a3=a3,
        abs_rel=abs_rel,
        rmse=rmse,
        log_10=log_10,
        rmse_log=rmse_log,
        silog=silog,
        sq_rel=sq_rel,
    )


# Demo Utilities
def b64_to_pil(b64string):
    image_data = re.sub('^data:image/.+;base64,', '', b64string)
    # image = Image.open(cStringIO.StringIO(image_data))
    return Image.open(BytesIO(base64.b64decode(image_data)))


def edges(d):
    dx = ndimage.sobel(d, 0)  # horizontal derivative
    dy = ndimage.sobel(d, 1)  # vertical derivative
    return np.abs(dx) + np.abs(dy)


class PointCloudHelper:
    def __init__(self, width=640, height=480):
        self.xx, self.yy = self.worldCoords(width, height)

    def worldCoords(self, width=640, height=480):
        hfov_degrees, vfov_degrees = 57, 43
        hFov = math.radians(hfov_degrees)
        vFov = math.radians(vfov_degrees)
        cx, cy = width / 2, height / 2
        fx = width / (2 * math.tan(hFov / 2))
        fy = height / (2 * math.tan(vFov / 2))
        xx, yy = np.tile(range(width), height), np.repeat(range(height), width)
        xx = (xx - cx) / fx
        yy = (yy - cy) / fy
        return xx, yy

    def depth_to_points(self, depth):
        depth[edges(depth) > 0.3] = np.nan  # Hide depth edges
        length = depth.shape[0] * depth.shape[1]
        # depth[edges(depth) > 0.3] = 1e6  # Hide depth edges
        z = depth.reshape(length)

        return np.dstack((self.xx * z, self.yy * z, z)).reshape((length, 3))

[((80, 13, 80, 49), 'numpy.maximum', 'np.maximum', ({(80, 25, 80, 34): 'gt / pred', (80, 38, 80, 47): 'pred / gt'}, {}), '(gt / pred, pred / gt)', True, 'import numpy as np\n'), ((86, 13, 86, 45), 'numpy.mean', 'np.mean', ({(86, 21, 86, 44): '(gt - pred) ** 2 / gt'}, {}), '((gt - pred) ** 2 / gt)', True, 'import numpy as np\n'), ((113, 17, 113, 64), 're.sub', 're.sub', ({(113, 24, 113, 48): '"""^data:image/.+;base64,"""', (113, 50, 113, 52): '""""""', (113, 54, 113, 63): 'b64string'}, {}), "('^data:image/.+;base64,', '', b64string)", False, 'import re\n'), ((119, 9, 119, 28), 'scipy.ndimage.sobel', 'ndimage.sobel', ({(119, 23, 119, 24): 'd', (119, 26, 119, 27): '0'}, {}), '(d, 0)', False, 'from scipy import ndimage\n'), ((120, 9, 120, 28), 'scipy.ndimage.sobel', 'ndimage.sobel', ({(120, 23, 120, 24): 'd', (120, 26, 120, 27): '1'}, {}), '(d, 1)', False, 'from scipy import ndimage\n'), ((94, 10, 94, 22), 'numpy.log', 'np.log', ({(94, 17, 94, 21): 'pred'}, {}), '(pred)', True, 'import numpy as np\n'), ((94, 25, 94, 35), 'numpy.log', 'np.log', ({(94, 32, 94, 34): 'gt'}, {}), '(gt)', True, 'import numpy as np\n'), ((121, 11, 121, 21), 'numpy.abs', 'np.abs', ({(121, 18, 121, 20): 'dx'}, {}), '(dx)', True, 'import numpy as np\n'), ((121, 24, 121, 34), 'numpy.abs', 'np.abs', ({(121, 31, 121, 33): 'dy'}, {}), '(dy)', True, 'import numpy as np\n'), ((130, 15, 130, 41), 'math.radians', 'math.radians', ({(130, 28, 130, 40): 'hfov_degrees'}, {}), '(hfov_degrees)', False, 'import math\n'), ((131, 15, 131, 41), 'math.radians', 'math.radians', ({(131, 28, 131, 40): 'vfov_degrees'}, {}), '(vfov_degrees)', False, 'import math\n'), ((85, 22, 85, 39), 'numpy.abs', 'np.abs', ({(85, 29, 85, 38): 'gt - pred'}, {}), '(gt - pred)', True, 'import numpy as np\n'), ((91, 16, 91, 26), 'numpy.log', 'np.log', ({(91, 23, 91, 25): 'gt'}, {}), '(gt)', True, 'import numpy as np\n'), ((91, 29, 91, 41), 'numpy.log', 'np.log', ({(91, 36, 91, 40): 'pred'}, {}), '(pred)', True, 'import numpy as np\n'), ((115, 30, 115, 58), 'base64.b64decode', 'base64.b64decode', ({(115, 47, 115, 57): 'image_data'}, {}), '(image_data)', False, 'import base64\n'), ((95, 20, 95, 37), 'numpy.mean', 'np.mean', ({(95, 28, 95, 36): '(err ** 2)'}, {}), '(err ** 2)', True, 'import numpy as np\n'), ((133, 26, 133, 44), 'math.tan', 'math.tan', ({(133, 35, 133, 43): '(hFov / 2)'}, {}), '(hFov / 2)', False, 'import math\n'), ((134, 27, 134, 45), 'math.tan', 'math.tan', ({(134, 36, 134, 44): '(vFov / 2)'}, {}), '(vFov / 2)', False, 'import math\n'), ((146, 15, 146, 55), 'numpy.dstack', 'np.dstack', ({(146, 25, 146, 54): '(self.xx * z, self.yy * z, z)'}, {}), '((self.xx * z, self.yy * z, z))', True, 'import numpy as np\n'), ((30, 11, 30, 46), 'torch.Tensor', 'torch.Tensor', ({(30, 24, 30, 45): '[0.485, 0.456, 0.406]'}, {}), '([0.485, 0.456, 0.406])', False, 'import torch\n'), ((31, 10, 31, 45), 'torch.Tensor', 'torch.Tensor', ({(31, 23, 31, 44): '[0.229, 0.224, 0.225]'}, {}), '([0.229, 0.224, 0.225])', False, 'import torch\n'), ((95, 40, 95, 52), 'numpy.mean', 'np.mean', ({(95, 48, 95, 51): 'err'}, {}), '(err)', True, 'import numpy as np\n'), ((97, 21, 97, 33), 'numpy.log10', 'np.log10', ({(97, 30, 97, 32): 'gt'}, {}), '(gt)', True, 'import numpy as np\n'), ((97, 36, 97, 50), 'numpy.log10', 'np.log10', ({(97, 45, 97, 49): 'pred'}, {}), '(pred)', True, 'import numpy as np\n')]