from .model import *


class sep_MLP(nn.Module):
    def __init__(self, dim, len_feats, categories):
        super(sep_MLP, self).__init__()
        self.len_feats = len_feats
        self.layers = nn.ModuleList([])
        for i in range(len_feats):
            self.layers.append(simple_MLP([dim, 5 * dim, categories[i]]))

    def forward(self, x):
        y_pred = list([])
        for i in range(self.len_feats):
            x_i = x[:, i, :]
            pred = self.layers[i](x_i)
            y_pred.append(pred)
        return y_pred


class SAINT_vision(nn.Module):
    def __init__(
        self,
        *,
        categories,
        num_continuous,
        dim,
        depth,
        heads,
        dim_head=16,
        dim_out=1,
        mlp_hidden_mults=(4, 2),
        mlp_act=None,
        num_special_tokens=0,
        continuous_mean_std=None,
        attn_dropout=0.0,
        ff_dropout=0.0,
        cont_embeddings="MLP",
        scalingfactor=10,
        attentiontype="col",
        final_mlp_style="common",
        y_dim=2,
    ):
        super().__init__()
        assert all(
            map(lambda n: n > 0, categories)
        ), "number of each category must be positive"

        # categories related calculations

        self.num_categories = len(categories)
        self.num_unique_categories = sum(categories)

        # create category embeddings table

        self.num_special_tokens = num_special_tokens
        self.total_tokens = categories[-1] + 256

        # for automatically offsetting unique category ids to the correct position in the categories embedding table

        categories_offset = torch.tensor(
            np.append(np.repeat(0, self.num_categories - 1), [256])
        )
        self.register_buffer("categories_offset", categories_offset)

        self.norm = nn.LayerNorm(num_continuous)
        self.num_continuous = num_continuous
        self.dim = dim
        self.cont_embeddings = cont_embeddings
        self.attentiontype = attentiontype
        self.final_mlp_style = final_mlp_style

        if self.cont_embeddings == "MLP":
            self.simple_MLP = nn.ModuleList(
                [
                    simple_MLP([1, 100, self.dim])
                    for _ in range(self.num_continuous)
                ]
            )
            input_size = (dim * self.num_categories) + (dim * num_continuous)
            nfeats = self.num_categories + num_continuous
        else:
            print("Continous features are not passed through attention")
            input_size = (dim * self.num_categories) + num_continuous
            nfeats = self.num_categories

            # transformer
        if attentiontype == "col":
            self.transformer = Transformer(
                num_tokens=self.total_tokens,
                dim=dim,
                depth=depth,
                heads=heads,
                dim_head=dim_head,
                attn_dropout=attn_dropout,
                ff_dropout=ff_dropout,
            )
        elif attentiontype in ["row", "colrow"]:
            self.transformer = RowColTransformer(
                num_tokens=self.total_tokens,
                dim=dim,
                nfeats=nfeats,
                depth=depth,
                heads=heads,
                dim_head=dim_head,
                attn_dropout=attn_dropout,
                ff_dropout=ff_dropout,
                style=attentiontype,
            )

        l = input_size // 8
        hidden_dimensions = list(map(lambda t: l * t, mlp_hidden_mults))
        all_dimensions = [input_size, *hidden_dimensions, dim_out]

        self.mlp = MLP(all_dimensions, act=mlp_act)
        self.embeds = nn.Embedding(self.total_tokens, self.dim)

        cat_mask_offset = torch.tensor(
            np.append(np.repeat(0, self.num_categories - 1), [2])
        )
        con_mask_offset = torch.empty(0)

        self.register_buffer("cat_mask_offset", cat_mask_offset)
        self.register_buffer("con_mask_offset", con_mask_offset)

        self.mask_embeds_cat = nn.Embedding(4, self.dim)
        self.mask_embeds_cont = nn.Embedding(4, self.dim)
        self.pos_encodings = nn.Embedding(self.num_categories, self.dim)
        if self.final_mlp_style == "common":
            self.mlp1 = simple_MLP(
                [dim, (self.total_tokens) * 2, self.total_tokens]
            )
            self.mlp2 = simple_MLP([dim, (self.num_continuous), 1])

        else:
            self.mlp1 = sep_MLP(dim, self.num_categories, categories)
            self.mlp2 = sep_MLP(
                dim,
                self.num_continuous,
                np.ones(self.num_continuous).astype(int),
            )

        self.mlpfory = simple_MLP([dim, 100, y_dim])

    def forward(self, x_categ, x_cont):
        x = self.transformer(x_categ, x_cont)
        y_reps = x[:, self.num_categories - 1, :]
        y_outs = self.mlpfory(y_reps)
        return y_outs
