from __future__ import annotations

import torch
from torch import Tensor, nn


class RawExpert(nn.Module):
    def __init__(self, hidden_size: int, decision_size: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(hidden_size * 2, decision_size),
            nn.GELU(),
            nn.LayerNorm(decision_size),
            nn.Linear(decision_size, decision_size),
        )

    def forward(self, state: Tensor, question: Tensor) -> Tensor:
        return self.net(torch.cat((state, question), dim=-1))


class CandidateScorer(nn.Module):
    """Shared pointwise scorer; candidate permutation only permutes its outputs."""

    def __init__(self, decision_size: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(decision_size * 4, decision_size),
            nn.GELU(),
            nn.Linear(decision_size, 1),
        )

    def forward(self, decision: Tensor, candidates: Tensor) -> Tensor:
        if decision.ndim == 1:
            decision = decision.unsqueeze(0)
        if candidates.ndim == 2:
            candidates = candidates.unsqueeze(0).expand(decision.size(0), -1, -1)
        decision = decision.unsqueeze(1).expand(-1, candidates.size(1), -1)
        features = torch.cat(
            (decision, candidates, decision * candidates, torch.abs(decision - candidates)), dim=-1
        )
        return self.net(features).squeeze(-1).squeeze(0)


class NoulHead(nn.Module):
    def __init__(self, decision_size: int):
        super().__init__()
        self.linear = nn.Linear(decision_size, 1)

    def forward(self, decision: Tensor) -> Tensor:
        return self.linear(decision).squeeze(-1)