"""Define loss functions that allow mixed outcome types and masked loss computation."""

import torch
import torch.nn.functional as F


class MixedLoss:

    def __init__(
        self, col_types, num_adj=1.0, cat_adj=1.0, bin_adj=1.0, pos_adj=1.0, device=None
    ):
        """
        Initialize the MixedLoss class with column
        """
        self.num_adj = num_adj
        self.cat_adj = cat_adj
        self.bin_adj = bin_adj
        self.pos_adj = pos_adj

        num, bin_, pos, cat = [], [], [], []
        c = 0
        for col in col_types:
            if col == "num":
                num.append(c)
                c += 1
            elif col == "bin":
                bin_.append(c)
                c += 1
            elif col == "pos":
                pos.append(c)
                c += 1
            elif isinstance(col, int):
                cat.append((c, col))
                c += col
            else:
                raise ValueError(f"Unknown column type: {col}")

        self.num_idx = torch.tensor(num, device=device, dtype=torch.long)
        self.bin_idx = torch.tensor(bin_, device=device, dtype=torch.long)
        self.pos_idx = torch.tensor(pos, device=device, dtype=torch.long)
        self.cat_slices = cat  # list of tuples

    @torch.no_grad()
    def _scratch(self, like):
        return like.new_zeros(like.shape)

    def __call__(self, pred, target):

        loss = self._scratch(pred)

        if self.num_idx.numel():
            loss[:, self.num_idx] = self.num_adj * F.mse_loss(
                pred[:, self.num_idx], target[:, self.num_idx], reduction="none"
            )

        if self.bin_idx.numel():
            loss[:, self.bin_idx] = self.bin_adj * F.binary_cross_entropy_with_logits(
                pred[:, self.bin_idx], target[:, self.bin_idx], reduction="none"
            )

        if self.pos_idx.numel():
            loss[:, self.pos_idx] = self.pos_adj * F.mse_loss(
                pred[:, self.pos_idx], target[:, self.pos_idx], reduction="none"
            )

        for start, k in self.cat_slices:
            # target expected as class index, not one-hot:
            tgt = target[:, start : start + k].argmax(1)  # (B,)
            prd = pred[:, start : start + k]  # (B, K)
            ce = self.cat_adj * F.cross_entropy(prd, tgt, reduction="none")  # (B,)
            loss[:, start : start + k] = ce.unsqueeze(1).expand(-1, k)

        return loss

    def masked(self, pred, target, mask):
        """
        Compute the masked mean loss directly without materializing the per-element matrix.
        This matches `_masked_loss(self(pred, target), mask)` but avoids the intermediate tensor.
        """
        if mask.shape != pred.shape:
            raise ValueError("Mixed losses and mask must have the same shape.")

        mask_f = mask.float()
        denom = mask_f.sum()
        if denom == 0:
            raise ValueError("Mask contains no observed entries; cannot compute masked loss.")

        total = pred.new_tensor(0.0)

        if self.num_idx.numel():
            mse = F.mse_loss(
                pred[:, self.num_idx], target[:, self.num_idx], reduction="none"
            )
            total = total + (self.num_adj * mse * mask_f[:, self.num_idx]).sum()

        if self.bin_idx.numel():
            bce = F.binary_cross_entropy_with_logits(
                pred[:, self.bin_idx], target[:, self.bin_idx], reduction="none"
            )
            total = total + (self.bin_adj * bce * mask_f[:, self.bin_idx]).sum()

        if self.pos_idx.numel():
            pos = F.mse_loss(
                pred[:, self.pos_idx], target[:, self.pos_idx], reduction="none"
            )
            total = total + (self.pos_adj * pos * mask_f[:, self.pos_idx]).sum()

        for start, k in self.cat_slices:
            tgt = target[:, start : start + k].argmax(1)  # (B,)
            prd = pred[:, start : start + k]  # (B, K)
            ce = F.cross_entropy(prd, tgt, reduction="none")  # (B,)
            # In the original computation, CE is broadcast across the block and multiplied
            # elementwise by the mask. Summing over the block mask gives the same scaling.
            block_mask_sum = mask_f[:, start : start + k].sum(1)  # (B,)
            total = total + (self.cat_adj * ce * block_mask_sum).sum()

        return total / denom


def _masked_loss(mixed_losses, mask):
    """

    Consolidate loss for only observed data points.

    """
    if mixed_losses.shape != mask.shape:
        raise ValueError("Mixed losses and mask must have the same shape.")
    denom = torch.sum(mask)
    if denom == 0:
        raise ValueError("Mask contains no observed entries; cannot compute masked loss.")
    loss = torch.sum(mixed_losses * mask.float()) / denom
    return loss
