from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn


@dataclass
class ReadoutOutput:
    hidden: torch.Tensor
    readout_mask: torch.Tensor


class TokenReadout(nn.Module):
    def forward(
        self,
        hidden: torch.Tensor,
        time_seq: torch.Tensor,
        padding_mask: torch.Tensor,
        readout_mask: torch.Tensor | None = None,
    ) -> ReadoutOutput:
        mask = padding_mask if readout_mask is None else readout_mask
        return ReadoutOutput(hidden=hidden, readout_mask=mask.bool())


class SameTimeGroupEndReadout(nn.Module):
    def __init__(self, reduce: str = "mean"):
        super().__init__()
        if reduce not in {"mean", "sum"}:
            raise ValueError("reduce must be either 'mean' or 'sum'")
        self.reduce = reduce

    def forward(
        self,
        hidden: torch.Tensor,
        time_seq: torch.Tensor,
        padding_mask: torch.Tensor,
        readout_mask: torch.Tensor | None = None,
    ) -> ReadoutOutput:
        if readout_mask is None:
            next_is_new_time = torch.ones_like(padding_mask, dtype=torch.bool)
            next_is_new_time[:, :-1] = time_seq[:, 1:] != time_seq[:, :-1]
            readout_mask = padding_mask.bool() & next_is_new_time
        else:
            readout_mask = readout_mask.bool()

        group_start = torch.ones_like(padding_mask, dtype=torch.bool)
        group_start[:, 1:] = time_seq[:, 1:] != time_seq[:, :-1]
        group_start = group_start & padding_mask.bool()

        group_id = group_start.long().cumsum(dim=1) - 1
        group_id = group_id.clamp_min(0)
        max_groups = hidden.size(1)

        group_sum = hidden.new_zeros(hidden.size(0), max_groups, hidden.size(2))
        group_sum.scatter_add_(
            1,
            group_id.unsqueeze(-1).expand_as(hidden),
            hidden * padding_mask.unsqueeze(-1).to(hidden.dtype),
        )

        if self.reduce == "mean":
            group_count = hidden.new_zeros(hidden.size(0), max_groups, 1)
            group_count.scatter_add_(
                1,
                group_id.unsqueeze(-1),
                padding_mask.unsqueeze(-1).to(hidden.dtype),
            )
            group_sum = group_sum / group_count.clamp_min(1.0)

        out = hidden.clone()
        out[readout_mask] = group_sum.gather(
            1,
            group_id.unsqueeze(-1).expand_as(hidden),
        )[readout_mask]
        return ReadoutOutput(hidden=out, readout_mask=readout_mask)


class LastValidReadout(nn.Module):
    def forward(
        self,
        hidden: torch.Tensor,
        time_seq: torch.Tensor,
        padding_mask: torch.Tensor,
        readout_mask: torch.Tensor | None = None,
    ) -> ReadoutOutput:
        batch_size, seq_len = padding_mask.shape
        last_idx = padding_mask.long().sum(dim=1).clamp_min(1) - 1
        out = hidden[torch.arange(batch_size, device=hidden.device), last_idx]
        mask = torch.ones(batch_size, dtype=torch.bool, device=hidden.device)
        return ReadoutOutput(hidden=out, readout_mask=mask)


def build_readout(name: str, **kwargs) -> nn.Module:
    name = name.lower()
    if name == "token":
        return TokenReadout()
    if name in {"same_time_group_end", "same_time"}:
        return SameTimeGroupEndReadout(**kwargs)
    if name == "last_valid":
        return LastValidReadout()
    raise ValueError(
        "Unknown readout {!r}. Available: token, same_time_group_end, last_valid.".format(
            name
        )
    )