"""In-memory representation of a DyGIE-format document (JSONL).

Trimmed port of the vendored radgraph.dygie.data.dataset_readers.document module: only the
NER + relation pieces are kept (no coreference clusters, no events -- this project's configs
always set loss_weights.coref = loss_weights.events = 0, so those heads never ran in v1
either). See training/README.md for the schema.
"""
import re
import json
from typing import Any, Dict, List, Optional

import numpy as np


def _fields_to_batches(d: dict, keys_to_ignore=()):
    """Inverse of `_batches_to_fields`: {"a": [1, 2], "b": [3, 4]} -> [{"a": 1, "b": 3}, {"a": 2, "b": 4}]."""
    keys = [k for k in d.keys() if k not in keys_to_ignore]
    lengths = {k: len(d[k]) for k in keys}
    if len(set(lengths.values())) != 1:
        raise ValueError(f"For document {d.get('doc_key')}, fields have different lengths: {lengths}.")
    length = next(iter(lengths.values()))
    return [{k: d[k][i] for k in keys} for i in range(length)]


def _batches_to_fields(batches: List[dict]):
    first_keys = batches[0].keys()
    for entry in batches[1:]:
        if set(entry.keys()) != set(first_keys):
            raise ValueError("Keys do not match on all entries.")
    res = {k: [] for k in first_keys}
    for batch in batches:
        for k, v in batch.items():
            res[k].append(v)
    return res


class Span:
    """A span, tracked both sentence-relative and document-relative."""

    def __init__(self, start: int, end: int, sentence: "Sentence", sentence_offsets: bool = False):
        self.sentence = sentence
        # `sentence.text_joined` is memoized on Sentence (computed once) rather than rejoined
        # here per span: the relation head can construct O(K^2) PredictedRelation/Span objects
        # per document during decode (K = pruned span count), so re-joining a ~1000-word
        # sentence per span turns into the dominant runtime cost otherwise.
        self.sentence_text = sentence.text_joined
        self.start_sent = start if sentence_offsets else start - sentence.sentence_start
        self.end_sent = end if sentence_offsets else end - sentence.sentence_start

    @property
    def start_doc(self):
        return self.start_sent + self.sentence.sentence_start

    @property
    def end_doc(self):
        return self.end_sent + self.sentence.sentence_start

    @property
    def span_doc(self):
        return (self.start_doc, self.end_doc)

    @property
    def span_sent(self):
        return (self.start_sent, self.end_sent)

    def __repr__(self):
        return str(self.span_sent)

    def __eq__(self, other):
        return (self.span_doc == other.span_doc and self.span_sent == other.span_sent
                and self.sentence == other.sentence)

    def __hash__(self):
        return hash(self.span_sent + (self.sentence_text,))


class NER:
    def __init__(self, ner, sentence: "Sentence", sentence_offsets: bool = False):
        self.span = Span(ner[0], ner[1], sentence, sentence_offsets)
        self.label = ner[2]

    def __repr__(self):
        return f"{self.span!r}: {self.label}"

    def __eq__(self, other):
        return self.span == other.span and self.label == other.label

    def to_json(self):
        return list(self.span.span_doc) + [self.label]


def _format_float(x):
    return round(x, 4)


class PredictedNER(NER):
    def __init__(self, ner, sentence, sentence_offsets: bool = False):
        """`ner` = [span_start, span_end, label, raw_score, softmax_score]."""
        super().__init__(ner, sentence, sentence_offsets)
        self.raw_score = ner[3]
        self.softmax_score = ner[4]

    def to_json(self):
        return super().to_json() + [_format_float(self.raw_score), _format_float(self.softmax_score)]


class Relation:
    def __init__(self, relation, sentence: "Sentence", sentence_offsets: bool = False):
        start1, end1, start2, end2, label = relation
        self.pair = (Span(start1, end1, sentence, sentence_offsets),
                     Span(start2, end2, sentence, sentence_offsets))
        self.label = label

    def __repr__(self):
        return f"{self.pair[0]!r}, {self.pair[1]!r}: {self.label}"

    def __eq__(self, other):
        return self.pair == other.pair and self.label == other.label

    def to_json(self):
        return list(self.pair[0].span_doc) + list(self.pair[1].span_doc) + [self.label]


class PredictedRelation(Relation):
    def __init__(self, relation, sentence, sentence_offsets: bool = False):
        """`relation` = [start1, end1, start2, end2, label, raw_score, softmax_score]."""
        super().__init__(relation[:5], sentence, sentence_offsets)
        self.raw_score = relation[5]
        self.softmax_score = relation[6]

    def to_json(self):
        return super().to_json() + [_format_float(self.raw_score), _format_float(self.softmax_score)]


class Sentence:
    """Despite the name, this project's documents always have exactly one "sentence" spanning
    the whole report (see training/README.md); the multi-sentence machinery is kept because
    the JSONL format is naturally list-of-sentences and nothing is gained by special-casing it.
    """

    def __init__(self, entry: dict, sentence_start: int, sentence_ix: int):
        self.sentence_start = sentence_start
        self.sentence_ix = sentence_ix
        self.text = entry["sentences"]
        self.text_joined = " ".join(self.text)  # memoized once; see Span.__init__
        self.metadata = {k: v for k, v in entry.items() if k.startswith("_")}

        if "ner" in entry:
            self.ner = [NER(x, self) for x in entry["ner"]]
            self.ner_dict = {e.span.span_sent: e.label for e in self.ner}
        else:
            self.ner, self.ner_dict = None, None

        self.predicted_ner = ([PredictedNER(x, self) for x in entry["predicted_ner"]]
                               if "predicted_ner" in entry else None)

        if "relations" in entry:
            self.relations = [Relation(x, self) for x in entry["relations"]]
            self.relation_dict = {(r.pair[0].span_sent, r.pair[1].span_sent): r.label
                                   for r in self.relations}
        else:
            self.relations, self.relation_dict = None, None

        self.predicted_relations = ([PredictedRelation(x, self) for x in entry["predicted_relations"]]
                                     if "predicted_relations" in entry else None)

    def to_json(self):
        res = {"sentences": self.text}
        if self.ner is not None:
            res["ner"] = [e.to_json() for e in self.ner]
        if self.predicted_ner is not None:
            res["predicted_ner"] = [e.to_json() for e in self.predicted_ner]
        if self.relations is not None:
            res["relations"] = [r.to_json() for r in self.relations]
        if self.predicted_relations is not None:
            res["predicted_relations"] = [r.to_json() for r in self.predicted_relations]
        res.update(self.metadata)
        return res

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

    def __repr__(self):
        return " ".join(self.text)


class Document:
    _ALLOWED_FIELDS = re.compile(r"doc_key|dataset|sentences|weight|.*ner$|.*relations$|^_.*")

    def __init__(self, doc_key, dataset, sentences: List[Sentence], weight: Optional[float] = None):
        self.doc_key = doc_key
        self.dataset = dataset
        self.sentences = sentences
        self.weight = weight

    @classmethod
    def from_json(cls, js: Dict[str, Any]) -> "Document":
        unexpected = [f for f in js if not cls._ALLOWED_FIELDS.match(f)]
        if unexpected:
            raise ValueError(f"Unexpected fields (prefix with '_' if intentional): {unexpected}")

        doc_key = js["doc_key"]
        dataset = js.get("dataset")
        entries = _fields_to_batches(js, ("doc_key", "dataset", "weight"))
        sentence_lengths = [len(e["sentences"]) for e in entries]
        sentence_starts = np.roll(np.cumsum(sentence_lengths), 1)
        sentence_starts[0] = 0
        sentences = [Sentence(entry, int(start), ix)
                     for ix, (entry, start) in enumerate(zip(entries, sentence_starts.tolist()))]
        return cls(doc_key, dataset, sentences, js.get("weight"))

    def to_json(self):
        res = {"doc_key": self.doc_key, "dataset": self.dataset}
        res.update(_batches_to_fields([s.to_json() for s in self.sentences]))
        if self.weight is not None:
            res["weight"] = self.weight
        return res

    @property
    def n_tokens(self):
        return sum(len(s) for s in self.sentences)

    def __getitem__(self, ix):
        return self.sentences[ix]

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

    def __repr__(self):
        return "\n".join(f"{i}: {' '.join(s.text)}" for i, s in enumerate(self.sentences))


class Dataset:
    def __init__(self, documents: List[Document]):
        self.documents = documents

    def __getitem__(self, i):
        return self.documents[i]

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

    @classmethod
    def from_jsonl(cls, fname):
        documents = []
        with open(fname) as f:
            for line in f:
                documents.append(Document.from_json(json.loads(line)))
        return cls(documents)

    def to_jsonl(self, fname):
        with open(fname, "w") as f:
            for doc in self.documents:
                print(json.dumps(doc.to_json()), file=f)