# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/preprocessing.ipynb (unless otherwise specified).

__all__ = ['log_transform', 'min_max_normalize', 'crop_and_pad_to_square', 'compute_per_pixel_stats']

# Cell
from ifcimglib import imglmdb, utils
import numpy
import matplotlib.pyplot as plt
import cv2
from tqdm import trange, tqdm
import seaborn

# Cell
def log_transform(x, m, channels):
    x[channels] = numpy.where(x[channels] == 0.0, 1.0, x[channels])
    x[channels] = numpy.log(x[channels])
    return x

# Cell
def min_max_normalize(x, m, post_norm):
    if not post_norm in ["zero", "mean", "no", "clip"]:
        raise ValueError("%s not recognized." % post_norm)

    if m is None:
        flat_x = x.reshape(x.shape[0], -1)
        mi = flat_x.min(axis=1)
        ma = flat_x.max(axis=1)
    else:
        x_copy = x.copy()
        x_copy[~m] = numpy.nan
        flat_x = x_copy.reshape(x_copy.shape[0], -1)

        nan_check = numpy.isnan(flat_x).all(axis=1)
        if nan_check.any():
            flat_x[nan_check] = 0.0

        mi = numpy.nanmin(flat_x, axis=1)
        ma = numpy.nanmax(flat_x, axis=1)

    denom = ma - mi
    denom[denom <= 0.0] = 1.0

    normed = ((x.T-mi)/denom).T

    if numpy.isnan(normed).any():
        raise ValueError("normed contains nan")
    if numpy.isinf(normed).any():
        raise ValueError("normed contains inf")

    if m is not None:
        if post_norm == "mean":
            for i in range(normed.shape[0]):
                if m[i].any():
                    me = numpy.mean(normed[i][m[i]])
                    if numpy.isnan(me):
                        raise ValueError("here", me)
                    if (~m[i]).any():
                        normed[i][~m[i]] = me
        elif post_norm == "zero":
            normed[~m] = 0.0
        elif post_norm == "clip":
            normed = numpy.clip(normed, 0, 1)

    return normed

# Cell
def crop_and_pad_to_square(x, desired_size):
    x = numpy.swapaxes(x, 0, -1)

    old_size = x.shape[:2] # old_size is in (height, width) format

    ratio = float(desired_size)/max(old_size)
    new_size = tuple([int(a*ratio) for a in old_size])

    x = cv2.resize(x, (new_size[1], new_size[0]))

    delta_w = desired_size - new_size[1]
    delta_h = desired_size - new_size[0]
    top, bottom = delta_h//2, delta_h-(delta_h//2)
    left, right = delta_w//2, delta_w-(delta_w//2)

    color = [0, 0, 0]
    return numpy.swapaxes(cv2.copyMakeBorder(x, top, bottom, left, right, cv2.BORDER_CONSTANT,
        value=color), -1, 0)

# Cell
def compute_per_pixel_stats(db, shape, do_preprocess_steps=False, idx=None):
    # according to http://proceedings.mlr.press/v28/wan13.html

    idx = numpy.arange(len(db)) if idx is None else idx

    if shape is None:
        shape = db.get_image(idx[0])[0].shape

    s = numpy.zeros(shape, dtype=numpy.float32)
    n = 0.0

    for i in tqdm(idx):
        x, m, _ = db.get_image(i)

        if do_preprocess_steps:
            x = x.astype(numpy.float32)
            x = log_transform(x, m, [1])
            x = min_max_normalize(x, m, "clip")
            x = crop_and_pad_to_square(x, 70)

        s += x
        n += 1

    mean = s / n

    s = numpy.zeros(shape, dtype=numpy.float32)
    for i in tqdm(idx):
        x, m, _ = db.get_image(i)

        if do_preprocess_steps:
            x = x.astype(numpy.float32)
            x = min_max_normalize(x, m, "clip")
            x = crop_and_pad_to_square(x, 70)

        s += (x - mean)**2

    s = numpy.sqrt(s/n)

    return mean, s