"""Relation head: prune to top-K spans, score every pair, cross-entropy against gold.

Port of `radgraph.dygie.models.relation.RelationExtractor`. Two deliberate simplifications
made possible by batch_size=1 (see dataset.py):

  1. No span mask / no ignore_index: every one of the K pruned spans is a real span (there's
     no cross-document padding to mask out), so every (i, j) pair among them is a legitimate
     loss target -- either a real relation or "null".
  2. The scorer outputs `n_labels - 1` channels (real labels only) with a null column
     prepended, exactly like ner_head.py already does. v1's RelationExtractor scores the
     *full* `n_labels` (which already includes the null vocab slot at index 0) and *also*
     prepends a dummy null column, leaving one dead, never-targeted output channel that gets
     trained toward nothing and just gets suppressed as a side effect of softmax competition.
     Sizing it like NER removes that dead channel and, as a bonus, means predicted/gold label
     ids line up with vocabulary indices directly -- no "+1 shift so null is class 0" dance.

`transformer_params` (medbit_transformer arm, see notes/v2_plan.md item 1) swaps the pairwise
`relation_feedforwards` scorer for a TransformerBlock applied to the pruner's already-pruned
top-K spans, *before* pair formation: the K surviving spans attend to each other, then pairs
are built from the now-contextualized embeddings and scored with a bare Linear (the
transformer's own internal FFN sublayer already did the nonlinear work `relation_feedforwards`
used to do). The pruner's own mention scorer stays a plain FeedForward -- it's a preliminary
filter picking which K spans survive, not the classifier itself.

`use_between_context` (medbit_relations_pooling arm, see notes/v2_plan.md item 2) adds a 4th
block to the pair representation: the mean-pooled *raw* word embeddings strictly between the
two spans (not the pooled span embeddings -- those don't preserve individual word positions).
Computed via a prefix sum over `word_embeddings` (`_between_context`) so all K^2 pairs are done
in one vectorized pass instead of a per-pair loop. Overlapping/adjacent/nested spans have no
words between them -> zero vector for that pair. Orthogonal to `transformer_params`: the arm
forks from the plain medbit baseline, not medbit_transformer, so both flags are never on
together in practice, but the code composes correctly either way (the between-block just adds
`word_emb_dim` to whichever pair representation -- FeedForward or transformer -- is in use).

`relation_feedforward_params` (medbit_relations_pooling arm only, defaults to `feedforward_params`
for every other arm) sizes just the relation pair scorer's hidden layers, kept separate from the
pruner's `mention_ff` -- the pruner still consumes plain `span_emb_dim` (unaffected by
`use_between_context`), so only the layers that actually see the enlarged pair embedding get
widened, mirroring the "pruner untouched" principle already applied to `transformer_params`.
"""
import math
from typing import Dict, List, Tuple

import torch
from torch import nn

from .document import PredictedRelation, Sentence
from .feedforward import FeedForward
from .metrics import RelationMetrics
from .pruner import Pruner
from .transformer_block import TransformerBlock
from .vocab import Vocabulary


class RelationHead(nn.Module):
    def __init__(self, vocab: Vocabulary, span_emb_dim: int, word_emb_dim: int,
                 feedforward_params: dict, spans_per_word: float, transformer_params: dict = None,
                 use_between_context: bool = False, relation_feedforward_params: dict = None):
        super().__init__()
        self.vocab = vocab
        self.spans_per_word = spans_per_word
        self.namespaces = [ns for ns in vocab.get_namespaces() if ns.endswith("relation_labels")]
        self.n_labels = {ns: vocab.get_vocab_size(ns) for ns in self.namespaces}
        self.use_transformer = transformer_params is not None
        self.use_between_context = use_between_context
        between_dim = word_emb_dim if use_between_context else 0
        relation_feedforward_params = relation_feedforward_params or feedforward_params

        self.pruners = nn.ModuleDict()
        self.span_transformers = nn.ModuleDict()
        self.relation_feedforwards = nn.ModuleDict()
        self.relation_scorers = nn.ModuleDict()
        self.metrics: Dict[str, RelationMetrics] = {}

        for ns in self.namespaces:
            mention_ff = FeedForward(span_emb_dim, feedforward_params["hidden_dims"],
                                      feedforward_params["dropout"])
            mention_scorer = nn.Sequential(mention_ff, nn.Linear(mention_ff.output_dim, 1))
            self.pruners[ns] = Pruner(mention_scorer)

            if self.use_transformer:
                block = TransformerBlock(span_emb_dim, **transformer_params)
                self.span_transformers[ns] = block
                self.relation_scorers[ns] = nn.Linear(3 * block.output_dim + between_dim,
                                                       self.n_labels[ns] - 1)
            else:
                relation_ff = FeedForward(3 * span_emb_dim + between_dim,
                                           relation_feedforward_params["hidden_dims"],
                                           relation_feedforward_params["dropout"])
                self.relation_feedforwards[ns] = relation_ff
                self.relation_scorers[ns] = nn.Linear(relation_ff.output_dim, self.n_labels[ns] - 1)
            self.metrics[ns] = RelationMetrics()

        self._loss = nn.CrossEntropyLoss(reduction="sum")

    @staticmethod
    def _pair_embeddings(top_embeddings: torch.Tensor) -> torch.Tensor:
        k = top_embeddings.size(0)
        e1 = top_embeddings.unsqueeze(1).expand(k, k, -1)
        e2 = top_embeddings.unsqueeze(0).expand(k, k, -1)
        return torch.cat([e1, e2, e1 * e2], dim=-1)

    @staticmethod
    def _between_context(word_embeddings: torch.Tensor,
                          top_spans: List[Tuple[int, int]]) -> torch.Tensor:
        """Mean of the raw word embeddings strictly between two spans, one block per (i, j) pair
        -- vectorized over the whole K x K grid via a prefix sum (no per-pair loop; K^2 pairs
        can run into the tens of thousands per doc). `low`/`high` don't depend on pair order (i,
        j) vs (j, i), so the between-region -- and this block -- is symmetric by construction.
        Overlapping/adjacent/nested spans (count <= 0) get a zero vector."""
        device, dtype = word_embeddings.device, word_embeddings.dtype
        starts = torch.tensor([s[0] for s in top_spans], dtype=torch.long, device=device)
        ends = torch.tensor([s[1] for s in top_spans], dtype=torch.long, device=device)

        zero_row = word_embeddings.new_zeros(1, word_embeddings.size(-1))
        prefix = torch.cat([zero_row, word_embeddings.cumsum(dim=0)], dim=0)  # (num_words + 1, dim)

        low = torch.minimum(ends.unsqueeze(1), ends.unsqueeze(0))       # end of whichever span is first
        high = torch.maximum(starts.unsqueeze(1), starts.unsqueeze(0))  # start of whichever span is second
        count = (high - low - 1).clamp(min=0)                          # words strictly between

        seg_sum = prefix[high] - prefix[low + 1]                       # (K, K, dim)
        mean = seg_sum / count.clamp(min=1).unsqueeze(-1).to(dtype)
        return torch.where((count > 0).unsqueeze(-1), mean, seg_sum.new_zeros(()))

    def forward(self, dataset: str, spans: List[Tuple[int, int]], span_embeddings: torch.Tensor,
                num_words: int, word_embeddings: torch.Tensor = None,
                relation_gold: Dict[Tuple[int, int], int] = None,
                sentence: Sentence = None) -> dict:
        ns = f"{dataset}__relation_labels"
        num_to_keep = math.ceil(num_words * self.spans_per_word)

        top_embeddings, top_indices, top_scores = self.pruners[ns](span_embeddings, num_to_keep)
        top_indices_list = top_indices.tolist()
        top_spans = [spans[i] for i in top_indices_list]

        if self.use_transformer:
            starts = torch.tensor([s[0] for s in top_spans], dtype=torch.long,
                                   device=span_embeddings.device)
            top_embeddings = self.span_transformers[ns](top_embeddings, starts)

        pair_embeddings = self._pair_embeddings(top_embeddings)
        if self.use_between_context:
            between = self._between_context(word_embeddings, top_spans)
            pair_embeddings = torch.cat([pair_embeddings, between], dim=-1)

        k = top_embeddings.size(0)
        if self.use_transformer:
            scores = self.relation_scorers[ns](pair_embeddings.view(k * k, -1)).view(k, k, -1)
        else:
            projected = self.relation_feedforwards[ns](pair_embeddings.view(k * k, -1))
            scores = self.relation_scorers[ns](projected).view(k, k, -1)
        scores = scores + top_scores.view(k, 1, 1) + top_scores.view(1, k, 1)
        dummy = scores.new_zeros(k, k, 1)
        scores = torch.cat([dummy, scores], dim=-1)  # (K, K, n_labels), null at index 0

        predicted_dict, predictions = self.decode(ns, scores.detach(), top_spans, sentence)
        output = {"namespace": ns, "scores": scores, "top_spans": top_spans, "predictions": predictions}

        if relation_gold is not None:
            pos_of = {orig: pos for pos, orig in enumerate(top_indices_list)}
            gold_matrix = scores.new_zeros(k, k, dtype=torch.long)
            for (oi, oj), label_id in relation_gold.items():
                pi, pj = pos_of.get(oi), pos_of.get(oj)
                if pi is not None and pj is not None:
                    gold_matrix[pi, pj] = label_id

            output["loss"] = self._loss(scores.view(k * k, -1), gold_matrix.view(-1))
            self.metrics[ns](predicted_dict, sentence.relation_dict or {})

        return output

    def decode(self, namespace: str, scores: torch.Tensor, top_spans: List[Tuple[int, int]],
               sentence: Sentence):
        raw, predicted = scores.max(dim=-1)
        soft, _ = torch.softmax(scores, dim=-1).max(dim=-1)

        res_dict, predictions = {}, []
        for i, j in (predicted != 0).nonzero(as_tuple=False).tolist():
            label = self.vocab.get_token_from_index(int(predicted[i, j]), namespace)
            span1, span2 = top_spans[i], top_spans[j]
            res_dict[(span1, span2)] = label
            entry = [span1[0], span1[1], span2[0], span2[1], label, float(raw[i, j]), float(soft[i, j])]
            predictions.append(PredictedRelation(entry, sentence, sentence_offsets=True))
        return res_dict, predictions

    def get_metrics(self, reset: bool = False) -> Dict[str, float]:
        res = {}
        for ns, metric in self.metrics.items():
            prefix = ns.replace("_labels", "")
            per_class = metric.get_per_class()  # {relation label -> F1}, read BEFORE the reset below
            for label, class_f1 in per_class.items():
                res[f"relation_perclass/{label}"] = class_f1
            res["relation_macro_f1"] = sum(per_class.values()) / len(per_class) if per_class else 0.0
            precision, recall, f1 = metric.get_metric(reset)
            res[f"{prefix}_precision"] = precision
            res[f"{prefix}_recall"] = recall
            res[f"{prefix}_f1"] = f1
        for name in ("precision", "recall", "f1"):
            values = [v for k, v in res.items()
                      if k.endswith(f"_{name}") and "perclass/" not in k and "_macro_" not in k]
            res[f"MEAN__relation_{name}"] = sum(values) / len(values) if values else 0.0
        return res