"""Implementation of Pointnet."""

from __future__ import annotations

import torch
from torch import nn

from vis4d.common.ckpt import load_model_checkpoint
from vis4d.common.typing import LossesType, ModelOutput
from vis4d.data.const import CommonKeys
from vis4d.op.base.pointnet import PointNetSegmentation, PointNetSemanticsOut
from vis4d.op.loss.orthogonal_transform_loss import (
    OrthogonalTransformRegularizationLoss,
)


class PointnetSegmentationModel(nn.Module):
    """Simple Segmentation Model using Pointnet."""

    def __init__(
        self,
        num_classes: int = 11,
        in_dimensions: int = 3,
        weights: str | None = None,
    ) -> None:
        """Simple Segmentation Model using Pointnet.

        Args:
            num_classes: Number of semantic classes
            in_dimensions: Input dimension
            weights: Path to weight file
        """
        super().__init__()
        self.model = PointNetSegmentation(
            n_classes=num_classes, in_dimensions=in_dimensions
        )
        if weights is not None:
            load_model_checkpoint(self, weights)

    def __call__(
        self, data: torch.Tensor, target: torch.Tensor | None = None
    ) -> PointNetSemanticsOut | ModelOutput:
        """Runs the semantic model.

        Args:
            data: Input Tensor Shape [N, C, n_pts]
            target: Target Classes shape [N, n_pts]
        """
        return self._call_impl(data, target)

    def forward(
        self, data: torch.Tensor, target: torch.Tensor | None = None
    ) -> PointNetSemanticsOut | ModelOutput:
        """Runs the semantic model.

        Args:
            data: Input Tensor Shape [N, C, n_pts]
            target: Target Classes shape [N, n_pts]
        """
        if target is not None:
            return self.forward_train(data, target)
        return self.forward_test(data)

    def forward_train(
        self,
        points: torch.Tensor,
        target: torch.Tensor,
    ) -> PointNetSemanticsOut:
        """Forward training stage.

        Args:
            points: Input Tensor Shape [N, C, n_pts]
            target: Target Classes shape [N, n_pts]
        """
        out = self.model(points)
        return out

    def forward_test(
        self,
        points: torch.Tensor,
    ) -> ModelOutput:
        """Forward test stage.

        Args:
            points: Input Tensor Shape [N, C, n_pts]
        """
        return {
            CommonKeys.semantics3d: torch.argmax(
                self.model(points).class_logits, dim=1
            )
        }


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

    def __init__(
        self,
        regularize_transform: bool = True,
        ignore_index: int = 255,
        transform_weight: float = 1e-3,
        semantic_weights: torch.Tensor | None = None,
    ) -> None:
        """Creates an instance of the class.

        Args:
            regularize_transform: If true add transforms to loss
            ignore_index: Semantic class that should be ignored
            transform_weight: Loss weight factor for transform
                              regularization loss
            semantic_weights: Classwise weights for semantic loss
        """
        super().__init__()
        self.segmentation_loss = nn.CrossEntropyLoss(
            weight=semantic_weights, ignore_index=ignore_index
        )
        self.transformation_loss = OrthogonalTransformRegularizationLoss()
        self.regularize_transform = regularize_transform
        self.transform_weight = transform_weight

    def forward(
        self, outputs: PointNetSemanticsOut, target: torch.Tensor
    ) -> LossesType:
        """Calculates the losss.

        Args:
            outputs: Pointnet output
            target: Target Labels
        """
        if not self.regularize_transform:
            dict(
                segmentation_loss=self.segmentation_loss(
                    outputs.class_logits, target
                )
            )

        return dict(
            segmentation_loss=self.segmentation_loss(
                outputs.class_logits, target
            ),
            transform_loss=self.transform_weight
            * self.transformation_loss(outputs.transformations),
        )