"""
Loss functions for PAT (Policy Action Tokenizer)

This module provides various loss functions for training action tokenizers,
including VQ (Vector Quantization) and RVQ (Residual Vector Quantization) losses.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, Optional, Union


class VQLoss(nn.Module):
    """Vector Quantization Loss"""

    def __init__(self, commitment_weight: float = 0.25, codebook_weight: float = 1.0):
        super().__init__()
        self.commitment_weight = commitment_weight
        self.codebook_weight = codebook_weight

    def forward(self, z_e: torch.Tensor, z_q: torch.Tensor) -> Dict[str, torch.Tensor]:
        commitment_loss = F.mse_loss(z_e, z_q.detach())
        codebook_loss = F.mse_loss(z_q, z_e.detach())
        total_loss = (
            self.commitment_weight * commitment_loss
            + self.codebook_weight * codebook_loss
        )

        return {
            "total": total_loss,
            "commitment": commitment_loss,
            "codebook": codebook_loss,
        }


class RVQLoss(nn.Module):
    """Residual Vector Quantization Loss"""

    def __init__(
        self,
        n_codebooks: int = 4,
        commitment_weight: float = 0.25,
        codebook_weight: float = 1.0,
        quantizer_dropout: float = 0.0,
    ):
        super().__init__()
        self.n_codebooks = n_codebooks
        self.commitment_weight = commitment_weight
        self.codebook_weight = codebook_weight
        self.quantizer_dropout = quantizer_dropout

    def forward(
        self, z_e_list: list, z_q_list: list, n_quantizers: Optional[int] = None
    ) -> Dict[str, torch.Tensor]:
        batch_size = z_e_list[0].shape[0]
        device = z_e_list[0].device

        if n_quantizers is None:
            n_quantizers = self.n_codebooks

        if self.training and self.quantizer_dropout > 0:
            n_quantizers_tensor = (
                torch.ones((batch_size,), device=device) * self.n_codebooks + 1
            )
            dropout = torch.randint(
                1, self.n_codebooks + 1, (batch_size,), device=device
            )
            n_dropout = int(batch_size * self.quantizer_dropout)
            n_quantizers_tensor[:n_dropout] = dropout[:n_dropout]
        else:
            n_quantizers_tensor = torch.full((batch_size,), n_quantizers, device=device)

        commitment_loss = 0.0
        codebook_loss = 0.0
        per_quantizer_losses = []

        for i, (z_e_i, z_q_i) in enumerate(zip(z_e_list, z_q_list)):
            if not self.training and i >= n_quantizers:
                break

            mask = (
                torch.full((batch_size,), fill_value=i, device=device)
                < n_quantizers_tensor
            ).float()

            commit_i = F.mse_loss(z_e_i, z_q_i.detach(), reduction="none").mean(
                dim=(1, 2)
            )
            codebook_i = F.mse_loss(z_q_i, z_e_i.detach(), reduction="none").mean(
                dim=(1, 2)
            )

            commitment_loss += (commit_i * mask).mean()
            codebook_loss += (codebook_i * mask).mean()
            per_quantizer_losses.append(
                {"commitment": commit_i.mean(), "codebook": codebook_i.mean()}
            )

        total_loss = (
            self.commitment_weight * commitment_loss
            + self.codebook_weight * codebook_loss
        )

        return {
            "total": total_loss,
            "commitment": commitment_loss,
            "codebook": codebook_loss,
            "per_quantizer": per_quantizer_losses,
        }


class ReconstructionLoss(nn.Module):
    """
    Reconstruction Loss for action sequences with dual-arm robot data

    Data format (20 dims total):
        Left arm (10): xyz(3) + 6D_rotation(6) + gripper(1)
        Right arm (10): xyz(3) + 6D_rotation(6) + gripper(1)
    """

    LEFT_XYZ = slice(0, 3)
    LEFT_6D = slice(3, 9)
    LEFT_GRIPPER = slice(9, 10)
    RIGHT_XYZ = slice(10, 13)
    RIGHT_6D = slice(13, 19)
    RIGHT_GRIPPER = slice(19, 20)

    def __init__(
        self,
        loss_type: str = "geodesic",
        position_weight: float = 1.0,
        rotation_weight: float = 1.0,
        gripper_weight: float = 1.0,
        temporal_weight: float = 0.0,
        temporal_smoothness: bool = False,
    ):
        super().__init__()
        self.loss_type = loss_type
        self.position_weight = position_weight
        self.rotation_weight = rotation_weight
        self.gripper_weight = gripper_weight
        self.temporal_weight = temporal_weight
        self.temporal_smoothness = temporal_smoothness

    def forward(
        self,
        x: torch.Tensor,
        x_recon: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
        if self.loss_type == "mse":
            recon_loss = self._mse_loss(x, x_recon, mask)
        elif self.loss_type == "weighted":
            recon_loss = self._weighted_loss(x, x_recon, mask)
        elif self.loss_type in ["geodesic", "combined"]:
            recon_loss = self._geodesic_loss(x, x_recon, mask)
        else:
            raise ValueError(f"Unknown loss_type: {self.loss_type}")

        if self.temporal_smoothness and self.temporal_weight > 0:
            temporal_loss = self._temporal_smoothness(x_recon, mask)
            total_loss = recon_loss + self.temporal_weight * temporal_loss
            return {
                "total": total_loss,
                "reconstruction": recon_loss,
                "temporal": temporal_loss,
            }

        return recon_loss

    def _temporal_smoothness(
        self, actions: torch.Tensor, mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        batch_size, seq_len, _ = actions.shape
        if seq_len < 3:
            return torch.tensor(0.0, device=actions.device)

        velocity = actions[:, 1:, :] - actions[:, :-1, :]
        acceleration = velocity[:, 1:, :] - velocity[:, :-1, :]

        if mask is not None:
            if mask.dim() == 2:
                mask = mask.unsqueeze(-1)
            if mask.shape[1] >= 3:
                mask_accel = mask[:, 2:, :]
                if mask_accel.dim() == acceleration.dim():
                    acceleration = acceleration * mask_accel

        return F.mse_loss(acceleration, torch.zeros_like(acceleration))

    def _apply_mask(
        self, loss: torch.Tensor, mask: Optional[torch.Tensor]
    ) -> torch.Tensor:
        if mask is None:
            return loss.mean()
        if mask.dim() == 2:
            mask = mask.unsqueeze(-1)
        while mask.dim() < loss.dim():
            mask = mask.unsqueeze(-1)
        return (loss * mask).sum() / mask.sum()

    def _mse_loss(
        self,
        x: torch.Tensor,
        x_recon: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        loss = F.mse_loss(x_recon, x, reduction="none")
        return self._apply_mask(loss, mask)

    def _weighted_loss(
        self,
        x: torch.Tensor,
        x_recon: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        pos_x = torch.cat([x[..., self.LEFT_XYZ], x[..., self.RIGHT_XYZ]], dim=-1)
        pos_recon = torch.cat(
            [x_recon[..., self.LEFT_XYZ], x_recon[..., self.RIGHT_XYZ]], dim=-1
        )
        pos_loss = F.mse_loss(pos_recon, pos_x, reduction="none").mean(
            dim=-1, keepdim=True
        )
        total_loss = self.position_weight * pos_loss

        rot_x = torch.cat([x[..., self.LEFT_6D], x[..., self.RIGHT_6D]], dim=-1)
        rot_recon = torch.cat(
            [x_recon[..., self.LEFT_6D], x_recon[..., self.RIGHT_6D]], dim=-1
        )
        rot_loss = F.mse_loss(rot_recon, rot_x, reduction="none").mean(
            dim=-1, keepdim=True
        )
        total_loss = total_loss + self.rotation_weight * rot_loss

        grip_x = torch.cat(
            [x[..., self.LEFT_GRIPPER], x[..., self.RIGHT_GRIPPER]], dim=-1
        )
        grip_recon = torch.cat(
            [x_recon[..., self.LEFT_GRIPPER], x_recon[..., self.RIGHT_GRIPPER]], dim=-1
        )
        grip_loss = F.mse_loss(grip_recon, grip_x, reduction="none").mean(
            dim=-1, keepdim=True
        )
        total_loss = total_loss + self.gripper_weight * grip_loss

        return self._apply_mask(total_loss, mask)

    def _geodesic_loss(
        self,
        x: torch.Tensor,
        x_recon: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        from pat.utils import rotation_6d_to_matrix_torch, geodesic_loss_torch

        pos_x = torch.cat([x[..., self.LEFT_XYZ], x[..., self.RIGHT_XYZ]], dim=-1)
        pos_recon = torch.cat(
            [x_recon[..., self.LEFT_XYZ], x_recon[..., self.RIGHT_XYZ]], dim=-1
        )

        loss = F.mse_loss(pos_recon, pos_x, reduction="none")
        # xyz, left and right has customized weights
        weights = torch.tensor(
            [0.4, 0.4, 0.4, 0.65, 0.65, 0.65],
            device=loss.device
        )

        pos_loss = (loss * weights).mean(dim=-1, keepdim=True)
        total_loss = self.position_weight * pos_loss

        left_6d_x = x[..., self.LEFT_6D]
        left_6d_recon = x_recon[..., self.LEFT_6D]
        left_R_x = rotation_6d_to_matrix_torch(left_6d_x)
        left_R_recon = rotation_6d_to_matrix_torch(left_6d_recon)
        left_geo_loss = geodesic_loss_torch(left_R_x, left_R_recon).unsqueeze(-1)

        right_6d_x = x[..., self.RIGHT_6D]
        right_6d_recon = x_recon[..., self.RIGHT_6D]
        right_R_x = rotation_6d_to_matrix_torch(right_6d_x)
        right_R_recon = rotation_6d_to_matrix_torch(right_6d_recon)
        right_geo_loss = geodesic_loss_torch(right_R_x, right_R_recon).unsqueeze(-1)

        rot_loss = (left_geo_loss + right_geo_loss) / 2
        total_loss = total_loss + self.rotation_weight * rot_loss

        grip_x = torch.cat(
            [x[..., self.LEFT_GRIPPER], x[..., self.RIGHT_GRIPPER]], dim=-1
        )
        grip_recon = torch.cat(
            [x_recon[..., self.LEFT_GRIPPER], x_recon[..., self.RIGHT_GRIPPER]], dim=-1
        )
        grip_loss = F.mse_loss(grip_recon, grip_x, reduction="none").mean(
            dim=-1, keepdim=True
        )
        total_loss = total_loss + self.gripper_weight * grip_loss

        return self._apply_mask(total_loss, mask)


class CodebookDiversityLoss(nn.Module):
    """Codebook Diversity Loss"""

    def __init__(self, weight: float = 0.1, temperature: float = 1.0):
        super().__init__()
        self.weight = weight
        self.temperature = temperature

    def forward(self, codes: torch.Tensor, codebook_size: int) -> torch.Tensor:
        codes_flat = codes.view(-1)
        usage = torch.bincount(codes_flat, minlength=codebook_size).float()
        probs = usage / (usage.sum() + 1e-10)
        entropy = -(probs * torch.log(probs + 1e-10)).sum()
        max_entropy = torch.log(torch.tensor(codebook_size, dtype=torch.float32))
        diversity_loss = 1.0 - (entropy / max_entropy)
        return self.weight * diversity_loss


class ChunkContinuityLoss(nn.Module):
    """Chunk Continuity Loss"""

    def __init__(
        self,
        boundary_weight: float = 1.0,
        temporal_weight: float = 0.1,
        chunk_size: int = 16,
    ):
        super().__init__()
        self.boundary_weight = boundary_weight
        self.temporal_weight = temporal_weight
        self.chunk_size = chunk_size

    def forward(
        self, actions: torch.Tensor, chunk_boundaries: Optional[torch.Tensor] = None
    ) -> Dict[str, torch.Tensor]:
        boundary_loss = self._boundary_continuity(actions, chunk_boundaries)
        temporal_loss = self._temporal_smoothness(actions)
        total_loss = (
            self.boundary_weight * boundary_loss + self.temporal_weight * temporal_loss
        )

        return {
            "total": total_loss,
            "boundary": boundary_loss,
            "temporal": temporal_loss,
        }

    def _boundary_continuity(
        self, actions: torch.Tensor, chunk_boundaries: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        if chunk_boundaries is None:
            chunk_boundaries = list(
                range(self.chunk_size, actions.shape[1], self.chunk_size)
            )

        if len(chunk_boundaries) == 0:
            return torch.tensor(0.0, device=actions.device)

        boundary_losses = []
        for boundary in chunk_boundaries:
            if boundary < actions.shape[1]:
                prev_end = actions[:, boundary - 1, :]
                next_start = actions[:, boundary, :]
                boundary_losses.append(F.mse_loss(next_start, prev_end))

        return (
            torch.stack(boundary_losses).mean()
            if boundary_losses
            else torch.tensor(0.0, device=actions.device)
        )

    def _temporal_smoothness(self, actions: torch.Tensor) -> torch.Tensor:
        velocity = actions[:, 1:, :] - actions[:, :-1, :]
        acceleration = velocity[:, 1:, :] - velocity[:, :-1, :]
        return F.mse_loss(acceleration, torch.zeros_like(acceleration))


class ActionCodecLoss(nn.Module):
    """Combined Loss for ActionCodec-style training"""

    def __init__(
        self,
        reconstruction_weight: float = 1.0,
        commitment_weight: float = 0.25,
        codebook_weight: float = 1.0,
        diversity_weight: float = 0.0,
        continuity_weight: float = 0.0,
        n_codebooks: int = 1,
    ):
        super().__init__()
        self.reconstruction_weight = reconstruction_weight
        self.commitment_weight = commitment_weight
        self.codebook_weight = codebook_weight
        self.diversity_weight = diversity_weight
        self.continuity_weight = continuity_weight

        self.recon_loss = ReconstructionLoss(loss_type="mse")
        self.vq_loss = VQLoss(
            commitment_weight=commitment_weight, codebook_weight=codebook_weight
        )
        self.rvq_loss = RVQLoss(
            n_codebooks=n_codebooks,
            commitment_weight=commitment_weight,
            codebook_weight=codebook_weight,
        )
        self.diversity_loss = CodebookDiversityLoss(weight=diversity_weight)
        self.continuity_loss = ChunkContinuityLoss(
            boundary_weight=continuity_weight, temporal_weight=continuity_weight * 0.1
        )

    def forward(
        self,
        x: torch.Tensor,
        x_recon: torch.Tensor,
        z_e: Optional[torch.Tensor] = None,
        z_q: Optional[torch.Tensor] = None,
        z_e_list: Optional[list] = None,
        z_q_list: Optional[list] = None,
        codes: Optional[torch.Tensor] = None,
        mask: Optional[torch.Tensor] = None,
    ) -> Dict[str, torch.Tensor]:
        loss_dict = {}
        total_loss = 0.0

        recon_loss = self.recon_loss(x, x_recon, mask)
        total_loss += self.reconstruction_weight * recon_loss
        loss_dict["reconstruction"] = recon_loss

        if z_e is not None and z_q is not None:
            vq_loss_dict = self.vq_loss(z_e, z_q)
            total_loss += vq_loss_dict["total"]
            loss_dict.update({f"vq/{k}": v for k, v in vq_loss_dict.items()})

        if z_e_list is not None and z_q_list is not None:
            rvq_loss_dict = self.rvq_loss(z_e_list, z_q_list)
            total_loss += rvq_loss_dict["total"]
            loss_dict.update({f"rvq/{k}": v for k, v in rvq_loss_dict.items()})

        if codes is not None and self.diversity_weight > 0:
            div_loss = self.diversity_loss(codes, codes.max().item() + 1)
            total_loss += div_loss
            loss_dict["diversity"] = div_loss

        if self.continuity_weight > 0:
            cont_loss_dict = self.continuity_loss(x_recon)
            total_loss += cont_loss_dict["total"]
            loss_dict.update({f"continuity/{k}": v for k, v in cont_loss_dict.items()})

        loss_dict["total"] = total_loss
        return loss_dict


class OATLoss(nn.Module):
    """
    OAT (Occupancy Action Transformer) Loss

    Combines:
        1. Reconstruction loss (default: geodesic for rotation, MSE for position/gripper)
        2. Temporal smoothness loss (second-order derivative / acceleration)
    """

    def __init__(
        self,
        position_weight: float = 1.0,
        rotation_weight: float = 1.0,
        gripper_weight: float = 0.5,
        temporal_weight: float = 0.1,
    ):
        super().__init__()

        self.recon_loss = ReconstructionLoss(
            loss_type="geodesic",
            position_weight=position_weight,
            rotation_weight=rotation_weight,
            gripper_weight=gripper_weight,
            temporal_weight=temporal_weight,
            temporal_smoothness=True,
        )

        self.position_weight = position_weight
        self.rotation_weight = rotation_weight
        self.gripper_weight = gripper_weight
        self.temporal_weight = temporal_weight

    def forward(
        self,
        x: torch.Tensor,
        x_recon: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> Dict[str, torch.Tensor]:
        loss_dict = {}
        total_loss = 0.0

        recon_output = self.recon_loss(x, x_recon, mask)

        if isinstance(recon_output, dict):
            recon_loss = recon_output["reconstruction"]
            temporal_loss = recon_output["temporal"]
            total_recon = recon_output["total"]

            loss_dict["reconstruction"] = recon_loss
            loss_dict["temporal"] = temporal_loss
            loss_dict["reconstruction_total"] = total_recon
        else:
            recon_loss = recon_output
            loss_dict["reconstruction"] = recon_loss
            total_recon = recon_loss

        total_loss += (
            total_recon if isinstance(total_recon, torch.Tensor) else recon_loss
        )
        loss_dict["total"] = total_loss

        return loss_dict


def load_info():
    """Load information about available loss functions"""
    return {
        "VQLoss": "Basic vector quantization loss",
        "RVQLoss": "Residual VQ loss with quantizer dropout",
        "ReconstructionLoss": "MSE and geodesic loss (default: geodesic)",
        "CodebookDiversityLoss": "Encourage diverse codebook usage",
        "ChunkContinuityLoss": "Encourage smooth transitions between chunks",
        "ActionCodecLoss": "Combined loss for ActionCodec-style training",
        "OATLoss": "OAT-style loss with geodesic + temporal smoothness",
    }


__all__ = [
    "VQLoss",
    "RVQLoss",
    "ReconstructionLoss",
    "CodebookDiversityLoss",
    "ChunkContinuityLoss",
    "ActionCodecLoss",
    "OATLoss",
    "load_info",
]
