from abc import abstractmethod
from typing import Dict, Optional

import torch
import torch.nn as nn

from .perceiver import SimplePerceiver
from .transformer import Transformer


class PointCloudSDFModel(nn.Module):
    @property
    @abstractmethod
    def device(self) -> torch.device:
        """
        Get the device that should be used for input tensors.
        """

    @property
    @abstractmethod
    def default_batch_size(self) -> int:
        """
        Get a reasonable default number of query points for the model.
        In some cases, this might be the only supported size.
        """

    @abstractmethod
    def encode_point_clouds(self, point_clouds: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        Encode a batch of point clouds to cache part of the SDF calculation
        done by forward().

        :param point_clouds: a batch of [batch x 3 x N] points.
        :return: a state representing the encoded point cloud batch.
        """

    def forward(
        self,
        x: torch.Tensor,
        point_clouds: Optional[torch.Tensor] = None,
        encoded: Optional[Dict[str, torch.Tensor]] = None,
    ) -> torch.Tensor:
        """
        Predict the SDF at the coordinates x, given a batch of point clouds.

        Either point_clouds or encoded should be passed. Only exactly one of
        these arguments should be None.

        :param x: a [batch x 3 x N'] tensor of query points.
        :param point_clouds: a [batch x 3 x N] batch of point clouds.
        :param encoded: the result of calling encode_point_clouds().
        :return: a [batch x N'] tensor of SDF predictions.
        """
        assert point_clouds is not None or encoded is not None
        assert point_clouds is None or encoded is None
        if point_clouds is not None:
            encoded = self.encode_point_clouds(point_clouds)
        return self.predict_sdf(x, encoded)

    @abstractmethod
    def predict_sdf(
        self, x: torch.Tensor, encoded: Optional[Dict[str, torch.Tensor]]
    ) -> torch.Tensor:
        """
        Predict the SDF at the query points given the encoded point clouds.

        Each query point should be treated independently, only conditioning on
        the point clouds themselves.
        """


class CrossAttentionPointCloudSDFModel(PointCloudSDFModel):
    """
    Encode point clouds using a transformer, and query points using cross
    attention to the encoded latents.
    """

    def __init__(
        self,
        *,
        device: torch.device,
        dtype: torch.dtype,
        n_ctx: int = 4096,
        width: int = 512,
        encoder_layers: int = 12,
        encoder_heads: int = 8,
        decoder_layers: int = 4,
        decoder_heads: int = 8,
        init_scale: float = 0.25,
    ):
        super().__init__()
        self._device = device
        self.n_ctx = n_ctx

        self.encoder_input_proj = nn.Linear(3, width, device=device, dtype=dtype)
        self.encoder = Transformer(
            device=device,
            dtype=dtype,
            n_ctx=n_ctx,
            width=width,
            layers=encoder_layers,
            heads=encoder_heads,
            init_scale=init_scale,
        )
        self.decoder_input_proj = nn.Linear(3, width, device=device, dtype=dtype)
        self.decoder = SimplePerceiver(
            device=device,
            dtype=dtype,
            n_data=n_ctx,
            width=width,
            layers=decoder_layers,
            heads=decoder_heads,
            init_scale=init_scale,
        )
        self.ln_post = nn.LayerNorm(width, device=device, dtype=dtype)
        self.output_proj = nn.Linear(width, 1, device=device, dtype=dtype)

    @property
    def device(self) -> torch.device:
        return self._device

    @property
    def default_batch_size(self) -> int:
        return self.n_query

    def encode_point_clouds(self, point_clouds: torch.Tensor) -> Dict[str, torch.Tensor]:
        h = self.encoder_input_proj(point_clouds.permute(0, 2, 1))
        h = self.encoder(h)
        return dict(latents=h)

    def predict_sdf(
        self, x: torch.Tensor, encoded: Optional[Dict[str, torch.Tensor]]
    ) -> torch.Tensor:
        data = encoded["latents"]
        x = self.decoder_input_proj(x.permute(0, 2, 1))
        x = self.decoder(x, data)
        x = self.ln_post(x)
        x = self.output_proj(x)
        return x[..., 0]