"""Pointnet++ Implementation."""

from __future__ import annotations

import torch
from torch import Tensor, nn

from vis4d.common.ckpt import load_model_checkpoint
from vis4d.common.typing import LossesType, ModelOutput
from vis4d.data.const import CommonKeys as K
from vis4d.op.base.pointnetpp import (
    PointNet2Segmentation,
    PointNet2SegmentationOut,
)


class PointNet2SegmentationModel(nn.Module):
    """PointNet++ Segmentation Model implementaiton."""

    def __init__(
        self,
        num_classes: int,
        in_dimensions: int = 3,
        weights: str | None = None,
    ):
        """Creates a Pointnet+++ Model.

        Args:
            num_classes (int): Number of classes
            in_dimensions (int, optional): Input dimensions. Defaults to 3.
            weights (str, optional): Path to weights. Defaults to None.
        """
        super().__init__()

        self.segmentation_model = PointNet2Segmentation(
            num_classes, in_dimensions
        )

        if weights is not None:
            load_model_checkpoint(self, weights)

    def forward(
        self, points3d: Tensor, semantics3d: Tensor | None = None
    ) -> PointNet2SegmentationOut | ModelOutput:
        """Forward pass of the model. Extract semantic predictions.

        Args:
            points3d (Tensor): Input point shape [b, N, C].
            semantics3d (torch.Tenosr): Groundtruth semantic labels of
                shape [b, N]. Defaults to None

        Returns:
            ModelOutput: Semantic predictions of the model.
        """
        x = self.segmentation_model(points3d)
        if semantics3d is not None:
            return x
        class_pred = torch.argmax(x.class_logits, dim=1)
        return {K.semantics3d: class_pred}


class Pointnet2SegmentationLoss(nn.Module):
    """Pointnet2SegmentationLoss Loss."""

    def __init__(
        self,
        ignore_index: int = 255,
        semantic_weights: Tensor | None = None,
    ) -> None:
        """Creates an instance of the class.

        Args:
            ignore_index (int, optional): Class Index that should be ignored.
                Defaults to 255.
            semantic_weights (Tensor, optional): Weights for each class.
        """
        super().__init__()
        self.segmentation_loss = nn.CrossEntropyLoss(
            weight=semantic_weights, ignore_index=ignore_index
        )

    def forward(
        self, outputs: PointNet2SegmentationOut, semantics3d: Tensor
    ) -> LossesType:
        """Calculates the loss.

        Args:
            outputs (PointNet2SegmentationOut): Model outputs.
            semantics3d (Tensor): Groundtruth semantic labels.
        """
        return dict(
            segmentation_loss=self.segmentation_loss(
                outputs.class_logits, semantics3d
            ),
        )