"""Predicted dependency and constituent structure over induced spans."""

from __future__ import annotations

from dataclasses import dataclass
import math

import torch
from torch import nn

from strata.modeling.ph_pat.config import PHPATConfig
from strata.modeling.ph_pat.span_compiler import SpanCompilerOutput


@dataclass(frozen=True, slots=True)
class DependencyChartOutput:
    nodes: torch.Tensor
    head_logits: torch.Tensor
    relation_logits: torch.Tensor
    constituent_logits: torch.Tensor
    segment_ids: torch.Tensor
    valid_mask: torch.Tensor


class DependencyChart(nn.Module):
    def __init__(self, config: PHPATConfig) -> None:
        super().__init__()
        self.config = config
        self.head_query = nn.Linear(config.d_model, config.d_model, bias=False)
        self.head_key = nn.Linear(config.d_model, config.d_model, bias=False)
        self.relation = nn.Linear(2 * config.d_model, config.dependency_relations)
        self.constituent = nn.Linear(config.d_model, config.constituent_categories)
        self.node_update = nn.Sequential(
            nn.Linear(config.d_model, config.d_model),
            nn.SiLU(),
            nn.Linear(config.d_model, config.d_model),
        )

    def forward(self, spans: SpanCompilerOutput) -> DependencyChartOutput:
        nodes = spans.nodes + self.node_update(spans.nodes)
        query = self.head_query(nodes)
        key = self.head_key(nodes)
        logits = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.config.d_model)
        same_segment = spans.segment_ids.unsqueeze(-1) == spans.segment_ids.unsqueeze(-2)
        valid_pair = spans.valid_mask.unsqueeze(-1) & spans.valid_mask.unsqueeze(-2) & same_segment
        logits = logits.masked_fill(~valid_pair, torch.finfo(logits.dtype).min)

        head_weights = torch.softmax(logits.float(), dim=-1).to(nodes.dtype)
        head_weights = torch.nan_to_num(head_weights)
        parent = torch.matmul(head_weights, nodes)
        relation_logits = self.relation(torch.cat((nodes, parent), dim=-1))
        return DependencyChartOutput(
            nodes=nodes,
            head_logits=logits,
            relation_logits=relation_logits,
            constituent_logits=self.constituent(nodes),
            segment_ids=spans.segment_ids,
            valid_mask=spans.valid_mask,
        )


__all__ = ["DependencyChart", "DependencyChartOutput"]