"""JSONL -> training examples.

Port of the relevant half of `radgraph.dygie.data.dataset_readers.dygie.DyGIEReader`
(enumerate candidate spans, attach NER/relation gold labels, drop gold spans wider than
`max_span_width`) plus the coref/event fields it also reads but which this project never
weights (see training/README.md: loss_weights.coref = loss_weights.events = 0 in every config).

Simplification vs. v1: every document in this dataset is a single "sentence" spanning the
whole report (see training/README.md), so there is no batching-over-sentences dimension here
-- one `Example` = one document. `build_example` asserts this so a future multi-sentence
document fails loudly instead of silently mis-training.

Simplification vs. v1 (memory): AllenNLP's `AdjacencyField` materializes a dense
(num_spans, num_spans) gold-relation matrix per document at collate time -- for a long report
with ~2000 words and max_span_width=12 that's ~24k spans, i.e. a ~575M-cell tensor. We instead
keep gold relations as a sparse {(span_i, span_j): label_id} dict and only ever materialize the
small (num_kept, num_kept) submatrix after pruning (see model.py), which is mathematically
identical (same loss, same gradients) at a fraction of the memory.
"""
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Tuple

import torch
from torch.utils.data import Dataset as TorchDataset

from .document import Document, Sentence
from .vocab import Vocabulary


def enumerate_spans(words: List[str], max_span_width: int) -> List[Tuple[int, int]]:
    """All inclusive (start, end) spans of width 1..max_span_width, ordered by (start, end)."""
    n = len(words)
    spans = []
    for start in range(n):
        for end in range(start, min(start + max_span_width, n)):
            spans.append((start, end))
    return spans


def _too_long(span: Tuple[int, int], max_span_width: int) -> bool:
    return span[1] - span[0] + 1 > max_span_width


@dataclass
class Example:
    doc_key: object
    dataset: str
    words: List[str]
    spans: List[Tuple[int, int]]
    span_index: Dict[Tuple[int, int], int]
    ner_label_ids: torch.Tensor                       # (num_spans,) long, 0 = null
    relation_gold: Dict[Tuple[int, int], int]          # (span_i, span_j) -> label id (>=1), sparse
    weight: float
    sentence: Sentence                                  # unfiltered gold, for relation eval (see module docstring)


def build_example(doc: Document, vocab: Vocabulary, max_span_width: int) -> Example:
    if len(doc.sentences) != 1:
        raise NotImplementedError(
            f"training_v2 assumes one sentence per document; {doc.doc_key} has {len(doc.sentences)}. "
            "See training/README.md ('the whole report is a single sentence')."
        )
    sent = doc.sentences[0]
    ner_ns = f"{doc.dataset}__ner_labels"
    rel_ns = f"{doc.dataset}__relation_labels"

    spans = enumerate_spans(sent.text, max_span_width)
    span_index = {s: i for i, s in enumerate(spans)}

    ner_label_ids = torch.zeros(len(spans), dtype=torch.long)
    for span, label in (sent.ner_dict or {}).items():
        if _too_long(span, max_span_width):
            continue
        ner_label_ids[span_index[span]] = vocab.get_token_index(label, ner_ns)

    relation_gold: Dict[Tuple[int, int], int] = {}
    for (span1, span2), label in (sent.relation_dict or {}).items():
        if _too_long(span1, max_span_width) or _too_long(span2, max_span_width):
            continue
        relation_gold[(span_index[span1], span_index[span2])] = vocab.get_token_index(label, rel_ns)

    return Example(doc_key=doc.doc_key, dataset=doc.dataset, words=sent.text, spans=spans,
                    span_index=span_index, ner_label_ids=ner_label_ids, relation_gold=relation_gold,
                    weight=doc.weight if doc.weight is not None else 1.0, sentence=sent)


def read_documents(path: str) -> List[Document]:
    import json
    documents = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line:
                documents.append(Document.from_json(json.loads(line)))
    return documents


class ExampleDataset(TorchDataset):
    """One dataset split (train/dev/test), pre-built into `Example`s.

    `collate_fn` is the identity on the single element: every config in this project uses
    `data_loader.batch_size = 1` (the joint model does not support batching multiple documents,
    see model.py), so there is no padding/collation logic to write.
    """

    def __init__(self, documents: List[Document], vocab: Vocabulary, max_span_width: int):
        self.examples = [build_example(doc, vocab, max_span_width) for doc in documents]

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, ix) -> Example:
        return self.examples[ix]

    @classmethod
    def from_jsonl(cls, path: str, vocab: Vocabulary, max_span_width: int) -> "ExampleDataset":
        return cls(read_documents(path), vocab, max_span_width)


def collate_fn(batch: List[Example]) -> Example:
    assert len(batch) == 1, "training_v2 only supports batch_size=1 (see module docstring)."
    return batch[0]