import torch
import torch.nn as nn


class AdaptiveLossBalancer(nn.Module):
    """
    Multi-task loss weighting with optional adaptive floor.

    Two modes:
      * ``fixed_weights``: directly supplied task weights (e.g., priority weighting
        for the time-resolved objective head). This is the default in PINO because
        the objective head is the only task that consumes the physics trajectory;
        letting it be starved removes the model's ability to learn time dynamics.
      * ``adaptive`` mode: weights are derived from per-task progress ratios
        (current / initial loss) with a per-task floor so that fast-converging
        tasks are not dropped to zero.

    The computed weights are stored in ``self.last_weights`` as a detached CPU list
    so training loops can log them without modifying the loss flow.
    """

    def __init__(
        self,
        num_tasks: int = 4,
        fixed_weights: torch.Tensor | None = None,
        momentum: float = 0.9,
        min_weight: float = 0.20,
    ) -> None:
        super().__init__()
        self.momentum = momentum
        self.min_weight = min_weight
        self.register_buffer("running_losses", torch.ones(num_tasks))
        self.register_buffer("initial_losses", torch.zeros(num_tasks))
        self.register_buffer("_initialized", torch.tensor(0, dtype=torch.bool))
        if fixed_weights is not None:
            fixed = torch.as_tensor(fixed_weights, dtype=torch.float32)
            if fixed.dim() == 1 and fixed.numel() == num_tasks:
                self.register_buffer("fixed_weights", fixed)
            else:
                raise ValueError(
                    f"fixed_weights must be a {num_tasks}-element vector, got shape {fixed.shape}"
                )
        else:
            self.register_buffer("fixed_weights", torch.zeros(num_tasks))

    def forward(self, loss_dict: dict[str, torch.Tensor]) -> torch.Tensor:
        """Combine task losses with either fixed or adaptive weights."""
        losses = torch.stack(list(loss_dict.values()))
        num_tasks = float(len(losses))

        if self.fixed_weights.abs().sum().item() > 0:
            weights = self.fixed_weights.to(losses.device)
            self.last_weights = weights.detach().cpu().tolist()
            return torch.sum(losses * weights)

        # Adaptive fallback.
        if not self._initialized.item() or torch.all(self.initial_losses == 0):
            self.initial_losses.copy_(losses.detach().clamp(min=1e-5))
            self.running_losses.copy_(losses.detach().clamp(min=1e-5))
            self._initialized.fill_(True)

        self.running_losses.copy_(
            self.momentum * self.running_losses + (1 - self.momentum) * losses.detach()
        )

        ratios = self.running_losses / self.initial_losses
        weights = ratios / (torch.sum(ratios) + 1e-8)

        # Apply a floor so that fast-converging tasks are not starved.
        weights = torch.clamp(weights, min=self.min_weight)  # type: ignore[attr-defined]
        weights = weights / (torch.sum(weights) + 1e-8)  # type: ignore[attr-defined]

        self.last_weights = weights.detach().cpu().tolist()
        return torch.sum(losses * weights * num_tasks)