import numpy as np
import cv2


def read_image(name):
    img = cv2.imread(name)
    return (img / 255.0)[:, :, ::-1]


def read_trimap(name):
    trimap_im = cv2.imread(name, 0) / 255.0
    h, w = trimap_im.shape
    trimap = np.zeros((h, w, 2))
    trimap[trimap_im == 1, 1] = 1
    trimap[trimap_im == 0, 0] = 1
    return trimap


class Preprocessing:
    def __init__(self, kernel, lower_bound=0.1, upper_bound=0.9, dilate_iter=10, erode_iter=10):
        self.kernel = kernel
        self.low_thresh = lower_bound
        self.high_thresh = upper_bound
        self.dilate_iter = dilate_iter
        self.erode_iter = erode_iter

    def get_target_mask(self, masks):
        out = np.zeros(masks[0].shape)
        for mask in masks:
            out += mask
        out = np.clip(out, 0, 1)
        return out

    def get_trimap(self, masks):
        target_mask = self.get_target_mask(masks)
        foreground = target_mask >= self.high_thresh
        ambiguous = (target_mask < self.high_thresh) * (target_mask >= self.low_thresh)

        erode = cv2.erode(foreground.astype("uint8"), self.kernel, iterations=self.erode_iter)
        dilate = cv2.dilate(ambiguous.astype("uint8"), self.kernel, iterations=self.dilate_iter)
        h, w = target_mask.shape

        bg_giver = np.clip((erode + dilate), 0, 1)
        trimap = np.zeros((h, w, 2))
        trimap[erode == 1, 1] = 1
        trimap[bg_giver == 0, 0] = 1

        return trimap
