"""torch.nn.Module for mixed type outputs"""

import torch
from torch import nn
import torch.nn.functional as F


class MixedActivation(torch.nn.Module):
    """
    Mix of column types in output layer.
    """

    def __init__(self, col_types, device=None):

        super().__init__()
        self.col_types = col_types

        pos = []
        c = 0
        for spec in col_types:
            if spec == "pos":
                pos.append(c)
                c += 1
            elif isinstance(spec, int):  # categorical block
                c += spec
            else:  # 'num' / 'bin'
                c += 1

        # moves with .to(device)
        self.register_buffer(
            "pos_idx", torch.tensor(pos, dtype=torch.long, device=device)
        )

    def __repr__(self):
        return f"MixedActivation(features={self.col_types})"

    def forward(self, x):
        """Forward pass through MAL"""
        if self.pos_idx.numel() == 0:
            return x

        out = x.clone()
        out[:, self.pos_idx] = F.relu(out[:, self.pos_idx])

        return out
