import numpy as np
import tensorflow as tf


def crop_or_pad_to_target_size(image, target_height, target_width):
    height = tf.shape(image)[0]
    width = tf.shape(image)[1]
    # n_ch = tf.shape(image)[2]

    d_h = height - target_height
    d_w = width - target_width
    if d_h > 0:
        h0 = tf.random.uniform((), 0, d_h + 1, tf.int32)
        image = tf.image.crop_to_bounding_box(image, h0, 0, target_height, width)
    else:
        h0 = 0
        image = tf.image.pad_to_bounding_box(image, 0, 0, target_height, width)
    if d_w > 0:
        w0 = tf.random.uniform((), 0, d_w + 1, tf.int32)
        image = tf.image.crop_to_bounding_box(image, 0, w0, target_height, target_width)
    else:
        w0 = 0
        image = tf.image.pad_to_bounding_box(image, 0, 0, target_height, target_width)

    image = tf.ensure_shape(image, [target_height, target_width, None])
    return image, h0, w0


def parse_train_data_func_full_annotation(
    data, augment=True, min_mask_area=16, target_height=512, target_width=512
):
    image = data["image"]
    masks = data["mask_indices"]
    mask_row_ids = masks.value_rowids()
    mask_values = masks.values

    target_width = int(target_width)
    target_height = int(target_height)

    image, h0, w0 = crop_or_pad_to_target_size(image, target_height, target_width)
    mask_values = mask_values - tf.cast([h0, w0], mask_values.dtype)
    if augment:
        if tf.random.uniform(()) >= 0.5:
            image = tf.image.flip_left_right(image)
            mask_values = mask_values * [1, -1] + [0, target_width - 1]
        if tf.random.uniform(()) >= 0.5:
            image = tf.image.flip_up_down(image)
            mask_values = mask_values * [-1, 1] + [target_height - 1, 0]

    valid_rows = tf.logical_and(
        tf.logical_and(mask_values[:, 0] >= 0, mask_values[:, 0] < target_height),
        tf.logical_and(mask_values[:, 1] >= 0, mask_values[:, 1] < target_width),
    )
    mask_values = tf.boolean_mask(mask_values, valid_rows)
    mask_row_ids = tf.boolean_mask(mask_row_ids, valid_rows)
    masks = tf.RaggedTensor.from_value_rowids(mask_values, mask_row_ids)
    mask_row_lengths = masks.row_lengths()
    valid_masks = mask_row_lengths >= min_mask_area
    masks = tf.ragged.boolean_mask(masks, valid_masks)

    locations = tf.cast(tf.reduce_mean(masks, axis=1), tf.float32)

    row_ids = masks.value_rowids() + 1
    labels = tf.zeros([target_height, target_width], row_ids.dtype)
    labels = tf.tensor_scatter_nd_update(labels, masks.values, row_ids)

    return {
        "image": image,
        "locations": locations,
        "mask_labels": labels,
    }


def parse_train_data_func(
    data, augment=True, size_jitter=None, target_height=512, target_width=512
):
    image = data["image"]
    binary_mask = tf.expand_dims(data["binary_mask"], -1)
    locations = data["locations"]

    img_and_label = tf.concat([image, tf.cast(binary_mask, tf.float32)], -1)
    if size_jitter is not None:
        h = tf.cast(tf.shape(img_and_label)[-3], tf.float32)
        w = tf.cast(tf.shape(img_and_label)[-2], tf.float32)
        scaling = tf.random.uniform([], size_jitter[0], size_jitter[1])
        img_and_label = tf.image.resize(
            img_and_label, [int(h * scaling), int(w * scaling)], antialias=True
        )
        locations = locations * scaling

    img_and_label, h0, w0 = crop_or_pad_to_target_size(
        img_and_label, target_height, target_width
    )
    locations = locations - tf.cast([h0, w0], tf.float32)

    if augment:
        if tf.random.uniform([]) >= 0.5:
            img_and_label = tf.image.flip_left_right(img_and_label)
            locations = locations * [1.0, -1.0] + [0, target_width]
        if tf.random.uniform([]) >= 0.5:
            img_and_label = tf.image.flip_up_down(img_and_label)
            locations = locations * [-1.0, 1.0] + [target_height, 0]

        if target_height == target_width:
            if tf.random.uniform([]) >= 0.5:
                img_and_label = tf.image.transpose(img_and_label)
                locations = locations[:, ::-1]

    image = img_and_label[..., :-1]
    if size_jitter is not None:
        binary_mask = tf.cast(img_and_label[..., -1] > 0.5, tf.float32)
    else:
        binary_mask = img_and_label[..., -1]

    # remove out-of-bound locations
    is_valid = tf.logical_and(
        tf.logical_and(locations[:, 0] > 1, locations[:, 0] < target_height - 1),
        tf.logical_and(locations[:, 1] > 1, locations[:, 1] < target_width - 1),
    )
    locations = tf.boolean_mask(locations, is_valid)

    # if augment:
    #     image = tf.image.random_brightness(image, 0.3)
    #     image = tf.image.random_contrast(image, 0.7, 1.3)

    return {
        "image": image,
        "locations": locations,
        "binary_mask": binary_mask,
    }


def parse_test_data_func(data, dim_multiple=32):
    image = data["image"]
    binary_mask = tf.expand_dims(data["binary_mask"], -1)
    bboxes = data["bboxes"]
    locations = data["locations"]
    masks = data["mask_indices"]

    height = tf.shape(image)[0]
    width = tf.shape(image)[1]
    target_height = (height - 1) // dim_multiple * dim_multiple + dim_multiple
    target_width = (width - 1) // dim_multiple * dim_multiple + dim_multiple

    image = tf.image.pad_to_bounding_box(image, 0, 0, target_height, target_width)
    binary_mask = tf.cast(
        tf.image.pad_to_bounding_box(binary_mask, 0, 0, target_height, target_width),
        tf.float32,
    )
    binary_mask = tf.squeeze(binary_mask, -1)

    scaling = data["scaling"] if "scaling" in data else tf.constant(1.0, tf.float32)

    row_ids = masks.value_rowids() + 1
    labels = tf.zeros([target_height, target_width], row_ids.dtype)
    labels = tf.tensor_scatter_nd_update(labels, masks.values, row_ids)

    return {
        "image": image,
        "bboxes": bboxes,
        "locations": locations,
        "img_id": data["img_id"],
        "scaling": scaling,
        "binary_mask": binary_mask,
        "mask_indices": masks,
        "mask_labels": labels,
    }
