################################################################################
# Authors: Zemin Xu
# License: MIT, see LICENSE.md
################################################################################

from typing import List, Dict


import torch
from torch import nn, Tensor




from .act import ACT
from .utils import expand_dims_to
from ...dataset.quantity import PROPERTY, UNIVERSAL_EMBEDDING_ALLOWED_PROPERTY


def add_rank1_to_left(T: Dict[int, torch.Tensor], rank1: torch.Tensor) -> Dict[int, torch.Tensor]:
    if 1 in T:
        T[1] = T[1] + rank1
    else:
        T[1] = rank1
    return T


def add_rank2_to_left(T: Dict[int, torch.Tensor], rank2: torch.Tensor) -> Dict[int, torch.Tensor]:
    if 2 in T:
        T[2] = T[2] + rank2
    else:
        T[2] = rank2
    return T


def add_rank3_to_left(T: Dict[int, torch.Tensor], rank3: torch.Tensor) -> Dict[int, torch.Tensor]:
    if 3 in T:
        T[3] = T[3] + rank3
    else:
        T[3] = rank3
    return T


ADD_FN = {
    1: add_rank1_to_left,
    2: add_rank2_to_left,
    3: add_rank3_to_left,
}


class UniversalInvariantEmbedding(torch.nn.Module):
    def __init__(
        self,
        out_dim: int,
        invariant_embedding: Dict[str, bool | str | int],
    ):
        super().__init__()

        self.embeddings = nn.ModuleDict()

        total_dim = 0
        for k, v in invariant_embedding.items():
            if v.get('enable', False) and PROPERTY[k]['rank'] == 0:
                p_type = PROPERTY[k]["type"]
                if p_type == "int":
                    self.embeddings[k] = nn.Embedding(v["num_embeddings"], out_dim)
                elif p_type == "float":
                    act = v.get("act", "silu")
                    self.embeddings[k] = nn.Sequential(
                        nn.Linear(1, out_dim, bias=False),
                        ACT[act](),
                        nn.Linear(out_dim, out_dim, bias=False),
                    )
                total_dim += out_dim

        self.project = nn.Sequential(
            nn.Linear(total_dim, out_dim, bias=False),
            nn.SiLU(),
        )

    def forward(
        self,
        batch: Tensor,
        attrs: Dict[str, Tensor],
    ) -> Tensor:
        embeddings = []

        for p, module in self.embeddings.items():
            attr = attrs[p]

            type_ = PROPERTY[p]['type']
            scope = PROPERTY[p]['scope']

            if scope == "per-system":
                attr = attr[batch]

            if type_ == 'float':
                attr = attr.unsqueeze(-1)
    
            embeddings.append(module(attr))

        return self.project(torch.cat(embeddings, dim=-1))


class EquivariantEmbedding(torch.nn.Module):
    def __init__(
        self,
        p: str,
        rank: int,
        scope: str,
        atomic_numbers: List,
        num_channel: int,
        element_trainable: bool = True,
        channel_trainable: bool = True,
        normalizer: float = 1.0
    ):
        super().__init__()
        num_elements = len(atomic_numbers)
        if element_trainable:
            self.element_weights = nn.Parameter(
                torch.ones(num_elements, dtype=torch.get_default_dtype())
            )
        else:
            self.register_buffer(
                "element_weights",
                torch.ones(num_elements, dtype=torch.get_default_dtype()),
            )
        if channel_trainable:
            self.channel_weights = nn.Parameter(
                torch.ones(num_channel, dtype=torch.get_default_dtype())
            )
        else:
            self.register_buffer(
                "channel_weights",
                torch.ones(num_channel, dtype=torch.get_default_dtype()),
            )
        self.p = p
        self.add_fn = ADD_FN[rank]
        self.scope = scope
        self.rank = rank
        self.normalizer = normalizer

    def forward(
        self,
        node_feats: Dict[int, Tensor],
        node_attrs: Tensor,
        batch: Tensor,
        attr: Dict[str, Tensor],
    ):
        element_idx = torch.argmax(node_attrs, dim=-1)
        label = attr * self.normalizer
        if self.scope == "per-system":
            label = label[batch].unsqueeze(1)
        else:
            label = label.unsqueeze(1)

        element_weights = expand_dims_to(
            self.element_weights[element_idx], n_dim=self.rank + 2
        )
        shape = (1, -1) + (1,) * self.rank
        channel_weights = self.channel_weights.view(*shape)
        
        embedding = label * element_weights * channel_weights
        return self.add_fn(node_feats, embedding)
  

class UniversalEquivariantEmbedding(torch.nn.Module):
    def __init__(
        self,
        equivariant_embedding: Dict[str, bool | str | int],
        atomic_numbers: List,
        num_channel: int,
    ):
        super().__init__()

        self.equivariant_embedding = equivariant_embedding
        self.embeddings = nn.ModuleDict()
        for k, v in equivariant_embedding.items():
            if v.get('enable', False) and PROPERTY[k]['rank'] > 0:
                self.embeddings[k] = EquivariantEmbedding(
                    k,
                    PROPERTY[k]["rank"],
                    PROPERTY[k]["scope"],
                    atomic_numbers,
                    num_channel,
                    element_trainable=True,
                    channel_trainable=True,
                    normalizer=float(v.get('normalizer', 1.0))
                )

    def forward(
            self, 
            node_feats: Dict[int, torch.Tensor], 
            node_attrs: torch.Tensor,
            batch: torch.Tensor,
            attrs: Dict[str, torch.Tensor]
        ) -> torch.Tensor:
        for p, module in self.embeddings.items():
            node_feats = module(node_feats, node_attrs, batch, attrs[p])
        return node_feats
