from typing import Optional, Tuple, Union
import logging
import torch
import torch.nn as nn

from diffusion_policy.model.common.module_attr_mixin import ModuleAttrMixin
from diffusion_policy.model.diffusion.positional_embedding import SinusoidalPosEmb


logger = logging.getLogger(__name__)


class PMFTransformerForDiffusion(ModuleAttrMixin):
    def __init__(
        self,
        input_dim: int,
        output_dim: int,
        horizon: int,
        n_obs_steps: Optional[int] = None,
        cond_dim: int = 0,
        n_layer: int = 12,
        n_head: int = 12,
        n_emb: int = 768,
        p_drop_emb: float = 0.1,
        p_drop_attn: float = 0.1,
        causal_attn: bool = False,
        obs_as_cond: bool = False,
        n_cond_layers: int = 0,
        n_time_tokens: int = 4,
        n_head_layers: int = 4,
    ) -> None:
        super().__init__()

        if n_obs_steps is None:
            n_obs_steps = horizon
        if n_time_tokens < 1:
            raise ValueError("n_time_tokens must be >= 1")
        if n_head_layers < 0:
            raise ValueError("n_head_layers must be >= 0")
        if n_head_layers >= n_layer:
            raise ValueError(
                "n_head_layers must be smaller than n_layer so shared trunk depth stays positive"
            )

        obs_as_cond = cond_dim > 0
        T = horizon
        n_global_cond_tokens = 2 * n_time_tokens
        T_cond = n_global_cond_tokens + (n_obs_steps if obs_as_cond else 0)
        n_shared_layers = n_layer - n_head_layers

        self.input_emb = nn.Linear(input_dim, n_emb)
        self.pos_emb = nn.Parameter(torch.zeros(1, T, n_emb))
        self.drop = nn.Dropout(p_drop_emb)

        self.t_emb = SinusoidalPosEmb(n_emb)
        self.r_emb = SinusoidalPosEmb(n_emb)
        self.t_tokens = nn.Parameter(torch.zeros(1, n_time_tokens, n_emb))
        self.r_tokens = nn.Parameter(torch.zeros(1, n_time_tokens, n_emb))
        self.cond_obs_emb = nn.Linear(cond_dim, n_emb) if obs_as_cond else None
        self.cond_pos_emb = nn.Parameter(torch.zeros(1, T_cond, n_emb))

        if n_cond_layers > 0:
            encoder_layer = nn.TransformerEncoderLayer(
                d_model=n_emb,
                nhead=n_head,
                dim_feedforward=4 * n_emb,
                dropout=p_drop_attn,
                activation="gelu",
                batch_first=True,
                norm_first=True,
            )
            self.encoder = nn.TransformerEncoder(
                encoder_layer=encoder_layer,
                num_layers=n_cond_layers,
            )
        else:
            self.encoder = nn.Sequential(
                nn.Linear(n_emb, 4 * n_emb),
                nn.Mish(),
                nn.Linear(4 * n_emb, n_emb),
            )

        decoder_layer = nn.TransformerDecoderLayer(
            d_model=n_emb,
            nhead=n_head,
            dim_feedforward=4 * n_emb,
            dropout=p_drop_attn,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.shared_decoder = nn.TransformerDecoder(
            decoder_layer=decoder_layer,
            num_layers=n_shared_layers,
        )
        self.u_decoder = nn.TransformerDecoder(
            decoder_layer=decoder_layer,
            num_layers=n_head_layers,
        )
        self.v_decoder = nn.TransformerDecoder(
            decoder_layer=decoder_layer,
            num_layers=n_head_layers,
        )

        if causal_attn:
            sz = T
            mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
            mask = mask.float().masked_fill(mask == 0, float("-inf")).masked_fill(mask == 1, float(0.0))
            self.register_buffer("mask", mask)

            if obs_as_cond:
                q_idx, c_idx = torch.meshgrid(
                    torch.arange(T),
                    torch.arange(T_cond),
                    indexing="ij",
                )
                obs_offset = n_global_cond_tokens
                visible = c_idx < obs_offset
                visible = visible | (q_idx >= (c_idx - obs_offset))
                memory_mask = visible.float().masked_fill(~visible, float("-inf")).masked_fill(visible, float(0.0))
                self.register_buffer("memory_mask", memory_mask)
            else:
                self.memory_mask = None
        else:
            self.mask = None
            self.memory_mask = None

        self.ln_u = nn.LayerNorm(n_emb)
        self.ln_v = nn.LayerNorm(n_emb)
        self.head_u = nn.Linear(n_emb, output_dim)
        self.head_v = nn.Linear(n_emb, output_dim)

        self.T = T
        self.T_cond = T_cond
        self.horizon = horizon
        self.n_obs_steps = n_obs_steps
        self.obs_as_cond = obs_as_cond
        self.n_global_cond_tokens = n_global_cond_tokens
        self.n_time_tokens = n_time_tokens
        self.n_layer = n_layer
        self.n_head_layers = n_head_layers
        self.n_shared_layers = n_shared_layers

        self.apply(self._init_weights)
        logger.info(
            "number of parameters: %e", sum(p.numel() for p in self.parameters())
        )
        logger.info(
            "PMFTransformerForDiffusion layers: shared=%d u_head=%d v_head=%d",
            self.n_shared_layers,
            self.n_head_layers,
            self.n_head_layers,
        )

    def _init_weights(self, module):
        ignore_types = (
            nn.Dropout,
            SinusoidalPosEmb,
            nn.TransformerEncoderLayer,
            nn.TransformerDecoderLayer,
            nn.TransformerEncoder,
            nn.TransformerDecoder,
            nn.ModuleList,
            nn.Mish,
            nn.Sequential,
        )
        if isinstance(module, (nn.Linear, nn.Embedding)):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.MultiheadAttention):
            for name in ("in_proj_weight", "q_proj_weight", "k_proj_weight", "v_proj_weight"):
                weight = getattr(module, name)
                if weight is not None:
                    torch.nn.init.normal_(weight, mean=0.0, std=0.02)
            for name in ("in_proj_bias", "bias_k", "bias_v"):
                bias = getattr(module, name)
                if bias is not None:
                    torch.nn.init.zeros_(bias)
        elif isinstance(module, nn.LayerNorm):
            torch.nn.init.zeros_(module.bias)
            torch.nn.init.ones_(module.weight)
        elif isinstance(module, PMFTransformerForDiffusion):
            torch.nn.init.normal_(module.pos_emb, mean=0.0, std=0.02)
            torch.nn.init.normal_(module.cond_pos_emb, mean=0.0, std=0.02)
            torch.nn.init.normal_(module.t_tokens, mean=0.0, std=0.02)
            torch.nn.init.normal_(module.r_tokens, mean=0.0, std=0.02)
        elif isinstance(module, ignore_types):
            pass
        else:
            raise RuntimeError("Unaccounted module {}".format(module))

    def get_optim_groups(self, weight_decay: float = 1e-3):
        decay = set()
        no_decay = set()
        whitelist_weight_modules = (torch.nn.Linear, torch.nn.MultiheadAttention)
        blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)
        for mn, m in self.named_modules():
            for pn, _ in m.named_parameters():
                fpn = "%s.%s" % (mn, pn) if mn else pn
                if pn.endswith("bias"):
                    no_decay.add(fpn)
                elif pn.startswith("bias"):
                    no_decay.add(fpn)
                elif pn.endswith("weight") and isinstance(m, whitelist_weight_modules):
                    decay.add(fpn)
                elif pn.endswith("weight") and isinstance(m, blacklist_weight_modules):
                    no_decay.add(fpn)

        no_decay.update(
            {
                "pos_emb",
                "cond_pos_emb",
                "t_tokens",
                "r_tokens",
                "_dummy_variable",
            }
        )

        param_dict = {pn: p for pn, p in self.named_parameters()}
        inter_params = decay & no_decay
        union_params = decay | no_decay
        assert len(inter_params) == 0, "parameters %s made it into both decay/no_decay sets!" % (str(inter_params),)
        assert len(param_dict.keys() - union_params) == 0, (
            "parameters %s were not separated into either decay/no_decay set!" % (str(param_dict.keys() - union_params),)
        )

        return [
            {
                "params": [param_dict[pn] for pn in sorted(list(decay))],
                "weight_decay": weight_decay,
            },
            {
                "params": [param_dict[pn] for pn in sorted(list(no_decay))],
                "weight_decay": 0.0,
            },
        ]

    def configure_optimizers(
        self,
        learning_rate: float = 1e-4,
        weight_decay: float = 1e-3,
        betas: Tuple[float, float] = (0.9, 0.95),
    ):
        optim_groups = self.get_optim_groups(weight_decay=weight_decay)
        return torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas)

    def _broadcast_time(self, value: Union[torch.Tensor, float, int], batch_size: int, device: torch.device):
        if not torch.is_tensor(value):
            value = torch.tensor([value], dtype=torch.float32, device=device)
        elif value.ndim == 0:
            value = value[None].to(device=device, dtype=torch.float32)
        else:
            value = value.to(device=device, dtype=torch.float32)
        return value.expand(batch_size)

    def forward(
        self,
        sample: torch.Tensor,
        t: Union[torch.Tensor, float, int],
        r: Union[torch.Tensor, float, int],
        cond: Optional[torch.Tensor] = None,
    ):
        batch_size = sample.shape[0]
        device = sample.device
        t = self._broadcast_time(t, batch_size, device)
        r = self._broadcast_time(r, batch_size, device)

        input_emb = self.input_emb(sample)

        t_cond = self.t_tokens + self.t_emb(t).unsqueeze(1)
        r_cond = self.r_tokens + self.r_emb(r).unsqueeze(1)
        cond_embeddings = [t_cond, r_cond]
        if self.obs_as_cond:
            cond_embeddings.append(self.cond_obs_emb(cond))
        cond_embeddings = torch.cat(cond_embeddings, dim=1)

        cond_pos = self.cond_pos_emb[:, : cond_embeddings.shape[1], :]
        memory = self.drop(cond_embeddings + cond_pos)
        memory = self.encoder(memory)

        token_pos = self.pos_emb[:, : input_emb.shape[1], :]
        x = self.drop(input_emb + token_pos)
        shared_x = self.shared_decoder(
            tgt=x,
            memory=memory,
            tgt_mask=self.mask,
            memory_mask=self.memory_mask,
        )
        u_x = self.u_decoder(
            tgt=shared_x,
            memory=memory,
            tgt_mask=self.mask,
            memory_mask=self.memory_mask,
        )
        v_x = self.v_decoder(
            tgt=shared_x,
            memory=memory,
            tgt_mask=self.mask,
            memory_mask=self.memory_mask,
        )
        return self.head_u(self.ln_u(u_x)), self.head_v(self.ln_v(v_x))
