import numpy as np

from .config import cfg
from .dropgenerator import generate_label, generateDrops_np


class RainDrop_Augmentor():

    def __init__(self, height, width):
        drops_list, label_map = generate_label(height, width, cfg)
        self.drops_list = drops_list
        self.label_map = label_map

    def generate_one(self, img_np, mode='rgb'):

        assert mode in ['gray', 'rgb']

        # requirement input [H, W, 3]
        if (mode == 'gray'):
            img_np = np.squeeze(img_np)
            img_np = np.expand_dims(img_np, axis=-1)
            img_np = np.repeat(img_np, 3, axis=-1)

        output_img, output_label, mask = generateDrops_np(img_np, cfg, self.drops_list)
        output_img = np.asarray(output_img)

        if (mode == 'gray'):
            output_img = output_img[:, :, 0]

        return output_img