from torch.nn.attention.bias import causal_lower_right

from priceforecast.models.nn.MAE_utils.components import *


class DatetimeInfoEmbedder(nn.Module):
    def __init__(self, dim_input: int, num_input: int, dim_pos_embedding: int, dim_mod_embedding: int):
        super().__init__()
        hoursofday, daysofweek, weeksofyear = torch.arange(24)/24*2*torch.pi, torch.arange(7)/7*2*torch.pi, torch.arange(54)/54*2*torch.pi
        hoursofday_sincos_embedding = torch.stack([torch.sin(hoursofday), torch.cos(hoursofday)], dim=1)
        daysofweek_sincos_embedding = torch.stack([torch.sin(daysofweek), torch.cos(daysofweek)], dim=1)
        weeksofyear_sincos_embedding = torch.stack([torch.sin(weeksofyear), torch.cos(weeksofyear)], dim=1)
        self.register_buffer('hoursofday_sincos_embedding', hoursofday_sincos_embedding)
        self.register_buffer('daysofweek_sincos_embedding', daysofweek_sincos_embedding)
        self.register_buffer('weeksofyear_sincos_embedding', weeksofyear_sincos_embedding)

        self.linear_projection = nn.Linear(2 * 3 * dim_input, dim_pos_embedding)
        self.position_embedding = nn.Embedding(num_input, dim_pos_embedding) #if learned_pos_embedding else SinusoidalEmbedding(num_input, dim_pos_embedding)
        self.modality_embedding = nn.Embedding(1, dim_mod_embedding)


    def forward(self, datetime_info: dict, input_idxs: torch.Tensor) -> torch.Tensor:

        hourofday, dayofweek, weekofyear = datetime_info['hourofday'].int(), datetime_info['dayofweek'].int(), datetime_info['weekofyear'].int()

        hourofday_embed = self.hoursofday_sincos_embedding[hourofday]
        dayofweek_embed = self.daysofweek_sincos_embedding[dayofweek]
        weekofyear_embed = self.weeksofyear_sincos_embedding[weekofyear]

        hourofday_embed = einops.rearrange(hourofday_embed, '... sin cos -> ... (sin cos)')
        dayofweek_embed = einops.rearrange(dayofweek_embed, '... sin cos -> ... (sin cos)')
        weekofyear_embed = einops.rearrange(weekofyear_embed, '... sin cos -> ... (sin cos)')

        datetime_info_embed = torch.cat((hourofday_embed, dayofweek_embed, weekofyear_embed), dim=-1)
        datetime_info_embed = self.linear_projection(datetime_info_embed)
        pos_embeddings = self.position_embedding(input_idxs)
        # add embeddings
        x = datetime_info_embed + pos_embeddings

        # concatenate with modality embedding
        B, N = input_idxs.shape
        mod_embed = self.modality_embedding.weight.expand(B, N, -1)
        x = torch.cat((x, mod_embed), dim=-1)
        return x


class Embedder(nn.Module):
    '''
    Simple embedding module. Takes tokenised inputs, projects them through a linear layer and concatenates them with positional and modality embeddings.
    '''
    def __init__(self, dim_input: int, num_input: int, dim_pos_embedding: int, dim_mod_embedding: int,
                 learned_pos_embedding: bool = True, use_datetime_information = False, **kwargs):
        super().__init__()

        # self.use_datetime_information = use_datetime_information
        # if self.use_datetime_information:
        #     self.datetime_info_embedding = DatetimeInfoEmbedding()
        # else:
        #     self.is_cat = not learned_pos_embedding #meow
        #     self.linear_projection = nn.Linear(dim_input, dim_input if self.is_cat else dim_pos_embedding)
        #     self.position_embedding = nn.Embedding(num_input, dim_pos_embedding) if learned_pos_embedding else SinusoidalEmbedding(num_input, dim_pos_embedding)
        self.is_cat = not learned_pos_embedding  # meow
        self.linear_projection = nn.Linear(dim_input, dim_input if self.is_cat else dim_pos_embedding)
        self.position_embedding = nn.Embedding(num_input, dim_pos_embedding) if learned_pos_embedding else SinusoidalEmbedding(num_input, dim_pos_embedding)
        self.modality_embedding = nn.Embedding(1, dim_mod_embedding)

    def forward(self, inputs: torch.Tensor, input_idxs: torch.Tensor):
        # src: (B, N, D), src_coo: (B, N)
        # if self.use_datetime_information: #TODO: change this: 1) it doesn't improve results 2) it doesn't work with multimodal models
        #     #separate inputs
        #     input_ts = inputs[0]
        #     inputs_datetime = inputs[1]['Datetime_info']
        #
        #     # concatenate embedded datetime information with input
        #     datetime_embedding = self.datetime_info_embedding(inputs_datetime)
        #     x = torch.cat((input_ts, datetime_embedding), dim=-1)
        #
        # else:
        #     # embed source data and coordinates
        #     projected_inputs = self.linear_projection(inputs)
        #     pos_embeddings = self.position_embedding(input_idxs)
        #     # concatenate or add embeddings
        #     x = torch.cat((projected_inputs, pos_embeddings), dim=-1) if self.is_cat else projected_inputs + pos_embeddings

        # embed source data and coordinates
        projected_inputs = self.linear_projection(inputs)
        pos_embeddings = self.position_embedding(input_idxs)
        # concatenate or add embeddings
        x = torch.cat((projected_inputs, pos_embeddings), dim=-1) if self.is_cat else projected_inputs + pos_embeddings

        # concatenate with modality embedding
        B, N = input_idxs.shape
        mod_embed = self.modality_embedding.weight.expand(B, N, -1)
        x = torch.cat((x, mod_embed), dim=-1)
        return x


class Query(nn.Module):
    def __init__(self, dim_query: int, num_outputs: int, dim_tgt_embedding: int, learned_tgt_embedding: bool = True,
                 **kwargs):
        super().__init__()
        # self.tgt_embed = nn.Embedding(num_outputs, embedding_dim=dim_tgt_embedding) if learned_tgt_embedding else SinusoidalEmbedding(num_outputs, dim_tgt_embedding)
        # self.mod_embed = nn.Embedding(1, embedding_dim=dim_query - dim_tgt_embedding)
        self.target_embedding = nn.Embedding(num_outputs, dim_query)

    def forward(self, output_idxs: torch.Tensor):
        # # tgt_coo: (B, N)
        # B, N = output_idxs.shape
        # tgt_embed = self.tgt_embed(output_idxs)
        # mod_embed = self.mod_embed.weight.expand(B, N, -1)
        # return torch.cat((tgt_embed, mod_embed), dim=-1)
        return self.target_embedding(output_idxs)


class CrossAttentionBlock(nn.Module):
    """
    Simplest version of Perceiver encoder, 1 cross-attention block
    """

    def __init__(self, dim_input: int, dim_latent: int = None, dim_head: int = 64, num_heads: int = None,
                 is_causal=False, **kwargs):

        super().__init__()
        dim_tgt = dim_latent if dim_latent is not None else dim_input
        dim_src = dim_input
        self.is_causal = is_causal
        self.norm_tgt = RMSNorm(dim_tgt)
        self.norm_src = RMSNorm(dim_src)
        self.norm_ff = RMSNorm(dim_tgt)
        self.attn = Attention(dim_q = dim_tgt, dim_kv = dim_src, dim_head = dim_head, num_heads = num_heads)
        self.ff = MLP(dim_tgt)

    def forward(self, inputs, latent, is_causal=None, **kwargs):
        # src: (B, *, N, D), tgt: (B, *, M, D)
        is_causal = is_causal or self.is_causal
        # cross-attention
        latent = self.norm_tgt(latent)
        inputs = self.norm_src(inputs)
        latent = self.attn(latent, inputs, inputs, is_causal=is_causal, **kwargs) # (B, *, M, D) x (B, *, N, D) -> dim: int(B, *, M, D)

        # feed-forward
        skip = latent
        latent = self.norm_ff(latent)
        latent = self.ff(latent)
        return latent + skip


class PerceiverIO_prob_ensemble_v4(nn.Module):

    def __init__(self, dim_tokens: int, num_inputs: int, dim_pos_embedding: int, dim_mod_embedding: int,
                 num_ensemble_predictions: int, dim_outputs=None, num_outputs=None, dim_latents=64, num_latents=32,
                 num_attention_heads=None, dim_attention_head=64, num_transformer_layers=5, transformer_drop_prob=0., **kwargs):
        super().__init__()

        if dim_outputs is None: dim_outputs = dim_tokens
        if num_outputs is None: num_outputs = num_inputs
        dim_embedded_inputs = dim_pos_embedding + dim_mod_embedding

        self.input_embedders = Embedder(dim_tokens, num_inputs, dim_pos_embedding, dim_mod_embedding, learned_pos_embedding= True, **kwargs)
        self.output_queries = Query(dim_latents, num_outputs, dim_pos_embedding)
        self.latents = nn.Embedding(num_latents, dim_latents)
        self.encoder = CrossAttentionBlock(dim_embedded_inputs, dim_latents, dim_attention_head, num_attention_heads)
        self.transformer = nn.Sequential(*[TransformerBlock(dim_latents, dim_attention_head, num_attention_heads, drop_path=transformer_drop_prob) for _ in range(num_transformer_layers)])
        self.decoder = CrossAttentionBlock(dim_latents, dim_latents, dim_attention_head, num_attention_heads)

        self.output_projections = nn.ModuleList([MLP(dim_latents, dim_outputs, dim_hidden=3*dim_outputs, bias=True) for _ in range(num_ensemble_predictions)])


    def forward(self, inputs: torch.Tensor, input_idxs: torch.Tensor, output_idxs: torch.Tensor):

        embedded_inputs = self.input_embedders(inputs, input_idxs)
        latents = self.latents.weight.expand(embedded_inputs.size(0), -1, -1)
        latents = self.encoder(embedded_inputs, latents)
        latents = self.transformer(latents)
        outputs = self.decoder(latents, self.output_queries(output_idxs))

        outputs = [output_projection(outputs) for output_projection in self.output_projections]
        outputs = torch.stack(outputs, dim=1)

        return outputs


class PerceiverAR_prob_ensemble_v2(nn.Module):
    def __init__(self, dim_tokens: int, num_inputs: int, num_outputs: int, dim_pos_embedding: int, dim_mod_embedding: int,
                 num_ensemble_predictions: int, num_attention_heads=None, dim_attention_head=64, num_transformer_layers=5, transformer_drop_prob=0., **kwargs):
        super().__init__()

        dim_embedded_inputs = dim_pos_embedding + dim_mod_embedding
        self.num_outputs = num_outputs

        self.input_embedders = Embedder(dim_tokens, num_inputs, dim_pos_embedding, dim_mod_embedding, learned_pos_embedding= True, **kwargs)
        self.encoder = CrossAttentionBlock(dim_input=dim_embedded_inputs, dim_head=dim_attention_head,
                                           num_heads=num_attention_heads)
        self.transformer = nn.Sequential(
            *[TransformerBlock(dim_embedded_inputs, dim_attention_head, num_attention_heads, is_causal=True, drop_path=transformer_drop_prob)
              for _ in range(num_transformer_layers)])

        self.output_projections = nn.ModuleList([MLP(dim_embedded_inputs, dim_tokens, dim_hidden=2*dim_embedded_inputs, bias=True) for _ in range(num_ensemble_predictions)])

    def forward(self, inputs: torch.Tensor, num_outputs: int = None):

        # the number of latents/outputs can be set independently at train and test time, to vary compute and performance
        # as explained in the PerceiverAR paper (https://arxiv.org/abs/2202.07765)
        num_outputs = num_outputs or self.num_outputs

        batch_size, num_inputs = (inputs.size(0), inputs.size(1)) if not isinstance(inputs, list) else (inputs[0].size(0), inputs[0].size(1))
        input_idxs = torch.arange(num_inputs, device=inputs[0].device).expand(batch_size, num_inputs)
        embedded_inputs = self.input_embedders(inputs, input_idxs)
        # divide inputs in context and latents
        latents = embedded_inputs[:, -num_outputs:, :]
        cross_attention_causal_mask = nn.attention.bias.causal_lower_right(num_outputs, num_inputs)
        latents = self.encoder(embedded_inputs, latents, attn_mask=cross_attention_causal_mask)
        latents = self.transformer(latents)
        outputs = [output_projection(latents) for output_projection in self.output_projections]
        outputs = torch.stack(outputs, dim=1)

        return outputs


class PerceiverIO_multimod_prob_ensemble_v3(nn.Module):

    def __init__(self, dim_tokens: int, num_inputs: int, dim_pos_embedding: int, dim_mod_embedding: int, num_ensemble_predictions: int,
                 dim_outputs=None, num_outputs=None, dim_latents=64, num_latents=32, num_attention_heads=None, dim_attention_head=64,
                 num_transformer_layers=5, transformer_drop_prob=0., additional_columns=None, use_datetime_information=False, **kwargs):
        super().__init__()

        if dim_outputs is None: dim_outputs = dim_tokens
        if num_outputs is None: num_outputs = num_inputs
        dim_embedded_inputs = dim_pos_embedding + dim_mod_embedding

        self.input_embedder = Embedder(dim_tokens, num_inputs, dim_pos_embedding, dim_mod_embedding, learned_pos_embedding= True, **kwargs)
        if use_datetime_information:
            self.datetime_embedder = DatetimeInfoEmbedder(dim_tokens, num_inputs, dim_pos_embedding, dim_mod_embedding)
        if additional_columns:
            self.additional_embedders = nn.ModuleDict({k: Embedder(dim_outputs, num_outputs, dim_pos_embedding, dim_mod_embedding) for k in additional_columns})
        self.output_queries = Query(dim_latents, num_outputs, dim_pos_embedding)
        self.latents = nn.Embedding(num_latents, dim_latents)
        self.encoder = CrossAttentionBlock(dim_embedded_inputs, dim_latents, dim_attention_head, num_attention_heads)
        self.transformer = nn.Sequential(*[TransformerBlock(dim_latents, dim_attention_head, num_attention_heads, drop_path=transformer_drop_prob) for _ in range(num_transformer_layers)])
        self.decoder = CrossAttentionBlock(dim_latents, dim_latents, dim_attention_head, num_attention_heads)

        self.output_projections = nn.ModuleList([MLP(dim_latents, dim_outputs, dim_hidden=3 * dim_outputs, bias=True) for _ in range(num_ensemble_predictions)])


    def forward(self, inputs: torch.Tensor or list, input_idxs: torch.Tensor or list, output_idxs: torch.Tensor):

        if isinstance(inputs, list):
            price_inputs, additional_inputs = inputs
            price_input_idxs, additional_input_idxs = input_idxs
            embedded_inputs = self.input_embedder(price_inputs, price_input_idxs)
            if 'Datetime_info' in additional_inputs:
                datetime_info = additional_inputs.pop('Datetime_info')
                embedded_datetime_info = self.datetime_embedder(datetime_info, additional_input_idxs['Datetime_info'])
                embedded_inputs = torch.cat([embedded_inputs, embedded_datetime_info], dim=1)
            if additional_inputs.keys():
                embedded_additional_inputs = torch.cat([self.additional_embedders[k](v, additional_input_idxs[k]) for k,v in additional_inputs.items()], dim=1)
                embedded_inputs = torch.cat((embedded_inputs, embedded_additional_inputs), dim=1)
        else:
            embedded_inputs = self.input_embedder(inputs, input_idxs)


        latents = self.latents.weight.expand(embedded_inputs.size(0), -1, -1)
        latents = self.encoder(embedded_inputs, latents)
        latents = self.transformer(latents)
        outputs = self.decoder(latents, self.output_queries(output_idxs))

        outputs = [output_projection(outputs) for output_projection in self.output_projections]
        outputs = torch.stack(outputs, dim=1)

        return outputs


class PerceiverAR_multimod_prob_ensemble(nn.Module):
    def __init__(self, dim_tokens: int, num_inputs: int, num_outputs: int, dim_pos_embedding: int, dim_mod_embedding: int,
                 num_ensemble_predictions: int, num_attention_heads=None, dim_attention_head=64, num_transformer_layers=5, transformer_drop_prob=0., additional_columns=None, **kwargs):
        super().__init__()

        dim_embedded_inputs = dim_pos_embedding + dim_mod_embedding
        self.num_outputs = num_outputs

        self.input_embedders = Embedder(dim_tokens, num_inputs, dim_pos_embedding, dim_mod_embedding, learned_pos_embedding= True, **kwargs)
        self.encoder = CrossAttentionBlock(dim_input=dim_embedded_inputs, dim_head=dim_attention_head,
                                           num_heads=num_attention_heads)
        if additional_columns:
            self.additional_embedders = nn.ModuleDict({k: Embedder(dim_tokens, num_inputs, dim_pos_embedding, dim_mod_embedding) for k in additional_columns})
            self.additional_encoders = nn.ModuleDict({k: CrossAttentionBlock(dim_input=dim_embedded_inputs, dim_head=dim_attention_head, num_heads=num_attention_heads) for k in additional_columns})

        self.transformer = nn.Sequential(
            *[TransformerBlock(dim_embedded_inputs, dim_attention_head, num_attention_heads, is_causal=True, drop_path=transformer_drop_prob)
              for _ in range(num_transformer_layers)])

        self.output_projections = nn.ModuleList([MLP(dim_embedded_inputs, dim_tokens, dim_hidden=2 * dim_embedded_inputs, bias=True) for _ in range(num_ensemble_predictions)])

    def forward(self, inputs: torch.Tensor, num_outputs: int = None):

        # the number of latents/outputs can be set independently at train and test time, to vary compute and performance
        # as explained in the PerceiverAR paper (https://arxiv.org/abs/2202.07765)
        num_outputs = num_outputs or self.num_outputs

        if isinstance(inputs, list):
            inputs, additional_inputs = inputs
        else:
            additional_inputs = None

        batch_size, num_inputs = (inputs.size(0), inputs.size(1))
        input_idxs = torch.arange(num_inputs, device=inputs.device).expand(batch_size, num_inputs)
        embedded_inputs = self.input_embedders(inputs, input_idxs)
        # divide inputs in context and latents
        latents = embedded_inputs[:, -num_outputs:, :]
        cross_attention_causal_mask = nn.attention.bias.causal_lower_right(num_outputs, num_inputs)
        latents = self.encoder(embedded_inputs, latents, attn_mask=cross_attention_causal_mask)

        if additional_inputs is not None:
            for key, additional_input in additional_inputs.items():
                embedded_additional_input = self.additional_embedders[key](additional_input, input_idxs)
                latents = self.additional_encoders[key](embedded_additional_input, latents)

        latents = self.transformer(latents)
        outputs = [output_projection(latents) for output_projection in self.output_projections]
        outputs = torch.stack(outputs, dim=1)

        return outputs


class PerceiverAR_prob_ensemble_wout_embed(nn.Module):
    def __init__(self, dim_tokens: int,
                 # num_inputs: int,
                 num_outputs: int,
                 # dim_pos_embedding: int, dim_mod_embedding: int,
                 num_ensemble_predictions: int, num_attention_heads=None, dim_attention_head=64, num_transformer_layers=5, transformer_drop_prob=0., additional_columns=None, **kwargs):
        super().__init__()

        # dim_embedded_inputs = dim_pos_embedding + dim_mod_embedding
        self.num_outputs = num_outputs

        # self.input_embedders = Embedder(dim_tokens, num_inputs, dim_pos_embedding, dim_mod_embedding, learned_pos_embedding= True, **kwargs)
        self.encoder = CrossAttentionBlock(dim_input=dim_tokens, dim_head=dim_attention_head,
                                           num_heads=num_attention_heads)
        if additional_columns:
            # self.additional_embedders = nn.ModuleDict({k: Embedder(dim_tokens, num_inputs, dim_pos_embedding, dim_mod_embedding) for k in additional_columns})
            self.additional_encoders = nn.ModuleDict({k: CrossAttentionBlock(dim_input=dim_tokens, dim_head=dim_attention_head, num_heads=num_attention_heads) for k in additional_columns})

        self.transformer = nn.Sequential(
            *[TransformerBlock(dim_tokens, dim_attention_head, num_attention_heads, is_causal=True, drop_path=transformer_drop_prob)
              for _ in range(num_transformer_layers)])

        self.output_projections = nn.ModuleList([MLP(dim_tokens, dim_tokens, dim_hidden=2 * dim_tokens, bias=True) for _ in range(num_ensemble_predictions)])

    def forward(self, inputs: torch.Tensor, num_outputs: int = None):

        # the number of latents/outputs can be set independently at train and test time, to vary compute and performance
        # as explained in the PerceiverAR paper (https://arxiv.org/abs/2202.07765)
        num_outputs = num_outputs or self.num_outputs

        if isinstance(inputs, list):
            inputs, additional_inputs = inputs
        else:
            additional_inputs = None

        batch_size, num_inputs = (inputs.size(0), inputs.size(1))
        # input_idxs = torch.arange(num_inputs, device=inputs.device).expand(batch_size, num_inputs)
        # embedded_inputs = self.input_embedders(inputs, input_idxs)
        # divide inputs in context and latents
        latents = inputs[:, -num_outputs:, :]
        cross_attention_causal_mask = nn.attention.bias.causal_lower_right(num_outputs, num_inputs)
        latents = self.encoder(inputs, latents, attn_mask=cross_attention_causal_mask)

        if additional_inputs is not None:
            for key, additional_input in additional_inputs.items():
                # embedded_additional_input = self.additional_embedders[key](additional_input, input_idxs)
                latents = self.additional_encoders[key](additional_input, latents)

        latents = self.transformer(latents)
        outputs = [output_projection(latents) for output_projection in self.output_projections]
        outputs = torch.stack(outputs, dim=1)

        return outputs