from dataclasses import dataclass, field
from math import sqrt
from typing import Dict, List, Literal, Optional, Tuple, Union

import torch
from torch import Tensor, nn

from models.base_model import BaseModel
from models.nn_utils import get_conv_output_shape, make_sequential_from_config
from models.shared.nn import PositionalEmbedding


class EncoderConfig(Dict):
    channels: List[int]
    kernels: List[int]
    strides: List[int]
    paddings: List[int]
    width: int
    height: int
    input_channels: int = 3


class DecoderConfig(Dict):
    conv_tranposes: List[bool]
    channels: List[int]
    kernels: List[int]
    strides: List[int]
    paddings: List[int]
    width: int
    height: int
    input_channels: int = 3


class Encoder(nn.Module):
    def __init__(
        self,
        width: int,
        height: int,
        channels: List[int] = (32, 32, 32, 32),
        kernels: List[int] = (5, 5, 5, 5),
        strides: List[int] = (1, 1, 1, 1),
        paddings: List[int] = (2, 2, 2, 2),
        input_channels: int = 3,
        batchnorms: List[bool] = tuple([False] * 4),
    ):
        super().__init__()
        assert len(kernels) == len(strides) == len(paddings) == len(channels)
        self.conv_bone = make_sequential_from_config(
            input_channels,
            channels,
            kernels,
            batchnorms,
            False,
            paddings,
            strides,
            "relu",
            try_inplace_activation=True,
        )
        output_channels = channels[-1]
        output_width, output_height = get_conv_output_shape(
            width, height, kernels, paddings, strides
        )
        self.pos_embedding = PositionalEmbedding(
            output_width, output_height, output_channels
        )
        self.lnorm = nn.GroupNorm(1, output_channels, affine=True, eps=0.001)
        self.conv_1x1 = [
            nn.Conv1d(output_channels, output_channels, kernel_size=1),
            nn.ReLU(inplace=True),
            nn.Conv1d(output_channels, output_channels, kernel_size=1),
        ]
        self.conv_1x1 = nn.Sequential(*self.conv_1x1)

    def forward(self, x: Tensor) -> Tensor:
        conv_output = self.conv_bone(x)
        conv_output = self.pos_embedding(conv_output)
        conv_output = conv_output.flatten(2, 3)  # bs x c x (w * h)
        conv_output = self.lnorm(conv_output)
        return self.conv_1x1(conv_output)


class Decoder(nn.Module):
    def __init__(
        self,
        input_channels: int,
        width: int,
        height: int,
        channels: List[int] = (32, 32, 32, 4),
        kernels: List[int] = (5, 5, 5, 3),
        strides: List[int] = (1, 1, 1, 1),
        paddings: List[int] = (2, 2, 2, 1),
        output_paddings: List[int] = (0, 0, 0, 0),
        conv_transposes: List[bool] = tuple([False] * 4),
        activations: List[str] = tuple(["relu"] * 4),
    ):
        super().__init__()
        self.conv_bone = []
        assert len(channels) == len(kernels) == len(strides) == len(paddings)
        if conv_transposes:
            assert len(channels) == len(output_paddings)
        self.pos_embedding = PositionalEmbedding(width, height, input_channels)
        self.width = width
        self.height = height

        self.conv_bone = make_sequential_from_config(
            input_channels,
            channels,
            kernels,
            False,
            False,
            paddings,
            strides,
            activations,
            output_paddings,
            conv_transposes,
            try_inplace_activation=True,
        )

    def forward(self, x: Tensor) -> Tuple[Tensor, Tensor]:
        x = self.pos_embedding(x)
        output = self.conv_bone(x)
        img, mask = output[:, :3], output[:, -1:]
        return img, mask


class BoqSlotAttentionModule(nn.Module):
    def __init__(self, num_slots, channels_enc, dim, iters=3, eps=1e-8, hidden_dim=128,
                        opt = "bi-level",
                        init = "embedding"):
        super().__init__()
    
        self.num_slots = num_slots
        self.iters = iters
        self.eps = eps
        self.scale = dim**-0.5
        self.opt = opt
        self.init = init

        assert self.opt in ["bi-level", "none"]
        assert self.init in ["embedding", "shared_gaussian", "none"]
        
        if self.init == "shared_gaussian":
            self.slots_mu = nn.Parameter(torch.empty(1, 1, dim))
            self.slots_log_sigma = nn.Parameter(torch.empty(1, 1, dim))
            with torch.no_grad():
                limit = sqrt(6.0 / (1 + dim))
                torch.nn.init.uniform_(self.slots_mu, -limit, limit)
                torch.nn.init.uniform_(self.slots_log_sigma, -limit, limit)
        elif self.init == "embedding":
            self.register_buffer("sigma", torch.tensor(torch.nan))  # Value set by trainer
            self.slots_init = nn.Embedding(num_slots, dim)
            nn.init.xavier_uniform_(self.slots_init.weight)
        else:
            raise NotImplementedError       
            
        self.to_q = nn.Linear(dim, dim, bias=False)
        self.to_k = nn.Linear(channels_enc, dim, bias=False)
        self.to_v = nn.Linear(channels_enc, dim, bias=False)

        self.gru = nn.GRUCell(dim, dim)

        hidden_dim = max(dim, hidden_dim)

        self.mlp = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, dim),
        )

        self.norm_input = nn.LayerNorm(channels_enc, eps=0.001)
        self.norm_slots = nn.LayerNorm(dim, eps=0.001)
        self.norm_pre_ff = nn.LayerNorm(dim, eps=0.001)
        self.dim = dim
        
       

    def forward(self, inputs: Tensor, sigma: float = 0., num_slots: Optional[int] = None) -> Tensor:
        b, n, _ = inputs.shape
        if num_slots is None:
            num_slots = self.num_slots

        if self.init == "shared_gaussian":    
            mu = self.slots_mu.expand(b, num_slots, -1)
            sigma = self.slots_log_sigma.expand(b, num_slots, -1).exp()
            slots = mu + sigma * torch.randn_like(sigma)
        elif self.init == "embedding":
            mu = self.slots_init.weight.expand(b,-1,-1)
            z = torch.randn_like(mu)
            slots = mu + z * self.sigma * mu.detach()
            
        inputs = self.norm_input(inputs)
        k, v = self.to_k(inputs), self.to_v(inputs)
        
        slots_init = slots
        for i in range(self.iters):
            if i == self.iters-1:
                if self.opt == "bi-level":
                    slots = slots.detach() + slots_init - slots_init.detach()
            
            slots_prev = slots

            slots = self.norm_slots(slots)
            q = self.to_q(slots)

            dots = torch.einsum("bid,bjd->bij", q, k) * self.scale
            attn = dots.softmax(dim=1) + self.eps
            attn = attn / attn.sum(dim=-1, keepdim=True)

            updates = torch.einsum("bjd,bij->bid", v, attn)

            slots = self.gru(
                updates.reshape(-1, self.dim), slots_prev.reshape(-1, self.dim)
            )

            slots = slots.reshape(b, -1, self.dim)
            slots = slots + self.mlp(self.norm_pre_ff(slots))

        return slots, attn


@dataclass(eq=False, repr=False)
class BoqSlotAttentionAE(BaseModel):
    latent_size: int

    encoder_params: EncoderConfig
    decoder_params: DecoderConfig
    
    input_channels: int = 3
    eps: float = 1e-8
    mlp_size: int = 128
    attention_iters: int = 3
    optimization_method: str = "bi-level"
    slots_initialization: str = "embedding"
    w_broadcast: Union[int, Literal["dataset"]] = "dataset"
    h_broadcast: Union[int, Literal["dataset"]] = "dataset"

    encoder: Encoder = field(init=False)
    decoder: Decoder = field(init=False)

    def __post_init__(self):
        super().__post_init__()
        self.loss_fn = nn.MSELoss()
        if self.w_broadcast == "dataset":
            self.w_broadcast = self.width
        if self.h_broadcast == "dataset":
            self.h_broadcast = self.height
        self.encoder_params.update(
            width=self.width, height=self.height, input_channels=self.input_channels
        )
        self.encoder = Encoder(**self.encoder_params)
        self.boq_slot_attention = BoqSlotAttentionModule(
            self.num_slots,
            self.encoder_params["channels"][-1],
            self.latent_size,
            self.attention_iters,
            self.eps,
            self.mlp_size,
            self.optimization_method,
            self.slots_initialization
        )
        self.decoder_params.update(
            width=self.w_broadcast,
            height=self.h_broadcast,
            input_channels=self.latent_size,
        )
        self.decoder = Decoder(**self.decoder_params)

    @property
    def slot_size(self) -> int:
        return self.latent_size

    def spatial_broadcast(self, slot: Tensor) -> Tensor:
        slot = slot.unsqueeze(-1).unsqueeze(-1)
        return slot.repeat(1, 1, self.w_broadcast, self.h_broadcast)

    def forward(self, x: Tensor) -> dict:
        with torch.no_grad():
            x = x * 2.0 - 1.0
        encoded = self.encoder(x)
        encoded = encoded.permute(0, 2, 1)
        z, attn = self.boq_slot_attention(encoded)
        bs = z.size(0)
        slots = z.flatten(0, 1)
        slots = self.spatial_broadcast(slots)
        img_slots, masks = self.decoder(slots)
        img_slots = img_slots.view(bs, self.num_slots, 3, self.width, self.height)
        masks = masks.view(bs, self.num_slots, 1, self.width, self.height)
        masks = masks.softmax(dim=1)

        recon_slots_masked = img_slots * masks
        recon_img = recon_slots_masked.sum(dim=1)
        loss = self.loss_fn(x, recon_img)

        _,num_slots,hw = attn.shape
        h = int(sqrt(hw))
        if h == masks.size(-1):
            attn = attn.view(masks.shape)
        else:
            attn = attn.view(bs,num_slots,h,h)
            attn = nn.functional.interpolate(attn, (masks.size(-2), masks.size(-1)), mode="bilinear")
        
        """
        attn = attn.view(bs,num_slots,-1)
        attn_min,_ = torch.min(attn,dim=-1,keepdim=True)
        attn_max,_ = torch.max(attn,dim=-1,keepdim=True)
        attn = (attn - attn_min)/(attn_max - attn_min)
        
        attn = attn.view(masks.shape)
        """
        attn = attn.view(attn.size(0),attn.size(1),-1)
        attn_min,_ = torch.min(attn,dim=-1,keepdim=True)
        attn_max,_ = torch.max(attn,dim=-1,keepdim=True)
        attn = (attn - attn_min)/(attn_max - attn_min)
        attn = nn.functional.softmax(attn,dim=1).view(masks.shape)

        
        with torch.no_grad():
            recon_slots_output = (img_slots + 1.0) / 2.0
        return {
            "loss": loss,  # scalar
            "mask": masks,  # (B, slots, 1, H, W)
            "attn": attn,
            "slot": recon_slots_output,  # (B, slots, 3, H, W)
            "representation": z,  # (B, slots, latent dim)
            #
            "reconstruction": recon_img,  # (B, 3, H, W)
        }
