import torch as th
import torchvision as tv
from torch import nn
from utils.utils import BatchToSharedObjects, SharedObjectsToBatch, LambdaModule
from pytorch_msssim import ms_ssim, ssim
from einops import rearrange, repeat, reduce



class FlowSmoothnessLoss(nn.Module):
    def __init__(self, depth = 4):
        super(FlowSmoothnessLoss, self).__init__()
        self.depth = depth

    def gradx(self, img):
        return img[:,:,:,1:] - img[:,:,:,:-1]                                                         

    def grady(self, img):
        return img[:,:,1:,:] - img[:,:,:-1,:]                                                         

    def grad2_error(self, flow, img, mask):
        img_grad_x = self.gradx(img)
        img_grad_y = self.grady(img)

        w_x = th.exp(-10.0 * reduce(th.abs(img_grad_x), 'b c h w -> b 1 h w', 'mean'))
        w_y = th.exp(-10.0 * reduce(th.abs(img_grad_y), 'b c h w -> b 1 h w', 'mean'))

        dx = self.gradx(flow)
        dy = self.grady(flow)

        dx2 = self.gradx(dx)
        dy2 = self.grady(dy)

        maskx = mask[:,:,:,2:] * mask[:,:,:,:-2]
        masky = mask[:,:,2:]   * mask[:,:,:-2]

        error = reduce(w_x[:,:,:,1:] * th.abs(dx2) * maskx, 'b c h w -> b', 'mean')
        error = error + reduce(w_y[:,:,1:] * th.abs(dy2) * masky, 'b c h w -> b', 'mean')

        return error

    def forward(self, flow, img, flow_mask):
        H, W = img.shape[2:]
        loss = self.grad2_error(flow, img, flow_mask)
        for i in range(1, self.depth):
            loss = loss + self.grad2_error(
                th.nn.functional.interpolate(flow, [H // 2**i, W // 2**i], mode='bicubic', align_corners=True),
                th.nn.functional.adaptive_avg_pool2d(img, [H // 2**i, W // 2**i]),
                th.nn.functional.adaptive_max_pool2d(flow_mask, [H // 2**i, W // 2**i])
            )

        return th.mean(loss)
                

class SSIMLoss(nn.Module):
    def __init__(self):
        super(SSIMLoss, self).__init__()

    def forward(self, output: th.Tensor, target: th.Tensor):
        return -ssim(output, target)

class MSSSIMLoss(nn.Module):
    def __init__(self):
        super(MSSSIMLoss, self).__init__()

    def forward(self, output: th.Tensor, target: th.Tensor):
        return -ms_ssim(output, target) #, normalize="relu")

class PositionLoss(nn.Module):
    def __init__(self, num_objects: int):
        super(PositionLoss, self).__init__()

        self.to_batch  = LambdaModule(lambda x: rearrange(x, 'b (o c) -> (b o) c', o = num_objects))

    def forward(self, position, position_last, slot_mask):
        
        slot_mask = rearrange(slot_mask, 'b o -> (b o) 1 1 1')
        position      = self.to_batch(position)
        position_last = self.to_batch(position_last).detach()

        return th.mean(slot_mask * (position - position_last)**2)


class MaskModulatedObjectLoss(nn.Module):
    def __init__(self, num_objects: int):
        super(MaskModulatedObjectLoss, self).__init__()

        self.to_batch  = SharedObjectsToBatch(num_objects)
        self.move_to_center  = MoveToCenter(num_objects)
    
    def forward(
        self, 
        object_output,
        object_target,
        slot_mask: th.Tensor
    ):
        slot_mask = rearrange(slot_mask, 'b o -> (b o) 1 1 1')
        object_output = th.sigmoid(self.to_batch(object_output) - 2.5)
        object_target = th.sigmoid(self.to_batch(object_target) - 2.5)

        return th.mean(slot_mask * (object_output - object_target)**2)

class ObjectModulator(nn.Module):
    def __init__(self, num_objects: int): 
        super(ObjectModulator, self).__init__()
        self.to_batch  = LambdaModule(lambda x: rearrange(x, 'b (o c) -> (b o) c', o = num_objects))
        self.to_shared = LambdaModule(lambda x: rearrange(x, '(b o) c -> b (o c)', o = num_objects))
        self.position  = None
        self.gestalt   = None

    def reset_state(self):
        self.position = None
        self.gestalt  = None

    def forward(self, position: th.Tensor, gestalt: th.Tensor, slot_mask: th.Tensor):

        position = self.to_batch(position)
        gestalt  = self.to_batch(gestalt)
        slot_mask = self.to_batch(slot_mask)

        if self.position is None or self.gestalt is None:
            self.position = position.detach()
            self.gestalt  = gestalt.detach()
            return self.to_shared(position), self.to_shared(gestalt)

        _position = slot_mask * position + (1 - slot_mask) * self.position
        position  = th.cat((position[:,:-1], _position[:,-1:]), dim=1) # keep the position of the objects fixed
        gestalt   = slot_mask * gestalt  + (1 - slot_mask) * self.gestalt

        self.gestalt = gestalt.detach()
        self.position = position.detach()

        return self.to_shared(position), self.to_shared(gestalt)

class MoveToCenter(nn.Module):
    def __init__(self, num_objects: int):
        super(MoveToCenter, self).__init__()

        self.to_batch2d = SharedObjectsToBatch(num_objects)
        self.to_batch  = LambdaModule(lambda x: rearrange(x, 'b (o c) -> (b o) c', o = num_objects))
    
    def forward(self, input: th.Tensor, position: th.Tensor):
        
        input    = self.to_batch2d(input) # b (o c) h w -> (b o) c h w
        position = self.to_batch(position).detach()
        position = th.stack((position[:,1], position[:,0]), dim=1)

        theta = th.tensor([1, 0, 0, 1], dtype=th.float, device=input.device).view(1,2,2)
        theta = repeat(theta, '1 a b -> n a b', n=input.shape[0])

        position = rearrange(position, 'b c -> b c 1')
        theta    = th.cat((theta, position), dim=2)

        grid   = nn.functional.affine_grid(theta, input.shape, align_corners=False)
        output = nn.functional.grid_sample(input, grid, align_corners=False)

        return output

class TranslationInvariantObjectLoss(nn.Module):
    def __init__(self, num_objects: int):
        super(TranslationInvariantObjectLoss, self).__init__()

        self.move_to_center  = MoveToCenter(num_objects)
        self.to_batch        = SharedObjectsToBatch(num_objects)
    
    def forward(
        self, 
        slot_mask: th.Tensor,
        object1: th.Tensor, 
        position1: th.Tensor,
        object2: th.Tensor, 
        position2: th.Tensor,
    ):
        slot_mask = rearrange(slot_mask, 'b o -> (b o) 1 1 1')
        object1 = self.move_to_center(th.sigmoid(object1 - 2.5), position1)
        object2 = self.move_to_center(th.sigmoid(object2 - 2.5), position2)

        return th.mean(slot_mask * (object1 - object2)**2)

