from ._tp_scatter_base import TensorProductScatter


class OpenEquivarianceTensorProductScatter(TensorProductScatter):
    _nequip_custom_ops_libs = ("openequivariance",)

    def __init__(
        self,
        feature_irreps_in,
        irreps_edge_attr,
        irreps_mid,
        instructions,
        use_opaque: bool,
    ) -> None:
        super().__init__(
            feature_irreps_in=feature_irreps_in,
            irreps_edge_attr=irreps_edge_attr,
            irreps_mid=irreps_mid,
            instructions=instructions,
        )
        # ^ we ensure that the base class keeps around a `self.tp` that carries its own set of persistent buffers
        # even though `self.tp` is not used, having its (persistent) buffers always around ensures state dict compatibility when adding on or removing this subclass module

        # === OEQ ===

        # we do lazy imports of oeq to allow `nequip-package` to pick this file up even if oeq is not installed
        # since `nequip-package` ignores files if it errors on loading the file

        from openequivariance import (
            TensorProductConv,
            TPProblem,
            torch_to_oeq_dtype,
        )

        tpp = TPProblem(
            feature_irreps_in,
            irreps_edge_attr,
            irreps_mid,
            instructions,
            irrep_dtype=torch_to_oeq_dtype(self.model_dtype),
            weight_dtype=torch_to_oeq_dtype(self.model_dtype),
            shared_weights=False,
            internal_weights=False,
        )
        self.tp_conv = TensorProductConv(
            tpp, torch_op=True, deterministic=False, use_opaque=use_opaque
        )

    def forward(self, x, edge_attr, edge_weight, edge_dst, edge_src):
        # explicit cast to account for AMP
        return self.tp_conv(
            x.to(self.model_dtype),
            edge_attr.to(self.model_dtype),
            edge_weight.to(self.model_dtype),
            edge_dst,
            edge_src,
        )