import cv2
import torch
from .custom_image_param import CustomImageParam
from .transforms import imagenet_transform


class MaskedImageParam(CustomImageParam):
    def __init__(self, mask_tensor, image=None, device="cuda"):
        """Custom image param, but with a mask over the original image.

        The mask helps update only certain parts of the image
        and leave the rest untouched by training. Ideal for
        things like: "differentiable" backgrounds.

        Args:
            image (str or torch.tensor): "path/to/image.jpg" or NCHW tensor
            mask_tensor (torch.tensor): NCHW tensor whose values are clipped between 0,1
            device (str): 'cpu' or 'cuda'
        """

        self.width = mask_tensor.shape[-1]
        self.height = mask_tensor.shape[-2]

        if image is None:
            from .utils import init_image_param, fft_to_rgb, lucid_colorspace_to_rgb

            if self.width % 2 == 1:
                self.param = init_image_param(
                    height=self.height, width=self.width + 1, sd=0.01, device=device
                )
            else:
                self.param = init_image_param(
                    height=self.height, width=self.width, sd=0.01, device=device
                )
            img = fft_to_rgb(
                height=self.height,
                width=self.width,
                image_parameter=self.param,
                device=device,
            )
            img = lucid_colorspace_to_rgb(t=img, device=device)
            image = torch.sigmoid(img)

            # image = torch.ones(1,3,self.height,self.width)

        super().__init__(image=image, device=device)

        self.mask = mask_tensor.to(self.device)

        if isinstance(image, str):
            image = cv2.cvtColor(cv2.imread(image), cv2.COLOR_BGR2RGB) / 255.0
            image = torch.tensor(image).permute(-1, 0, 1).unsqueeze(0)

        self.original_nchw_image_tensor = image.to(device)

        assert (
            self.mask.shape[-2:] == self.to_nchw_tensor(device=self.device).shape[-2:]
        ), (
            "The height and width of the input image"
            + str(self.to_nchw_tensor(device=self.device).shape[-2:])
            + "and the mask"
            + str(self.mask.shape[-2:])
            + "do not match."
        )

    def to_chw_tensor(self, device="cpu"):
        t = self.forward(device=device).squeeze(0).clamp(
            0, 1
        ).detach() + self.original_nchw_image_tensor.to(device) * (
            1 - self.mask.to(device)
        )
        return t.squeeze(0)

    def forward(self, device):
        return self.normalize(self.postprocess(device=device), device=device).clamp(
            0, 1
        ) * self.mask.to(device)

    def update_mask(self, mask):
        """updates the mask to have a new value

        Warning: it also updates the original image

        Args:
            mask (torch.tensor): new mask to be used
        """

        self.original_nchw_image_tensor = self.to_chw_tensor(
            device=self.device
        ).unsqueeze(0)
        self.mask = mask.to(self.device)

    def to_hwc_tensor(self, device="cpu"):
        t = self.forward(device=device).squeeze(0).clamp(
            0, 1
        ).detach() + self.original_nchw_image_tensor.to(device) * (
            1 - self.mask.to(device)
        )
        return t.squeeze(0).permute(1, 2, 0)
