"""Classes for representing speaker turns and interacting with RTTM files.

Taken from https://github.com/nryant/dscore
"""
from collections import namedtuple
from intervaltree import Interval, IntervalTree

from .uem import UEM
from .utils import groupby, warn, xor


Token = namedtuple("Token", ["type", "timestamp", "speaker", "hyp_index", "weight"])


class Turn(object):
    """Speaker turn class.

    A turn represents a segment of audio attributed to a single speaker.

    Parameters
    ----------
    onset : float
        Onset of turn in seconds from beginning of recording.

    offset : float, optional
        Offset of turn in seconds from beginning of recording. If None, then
        computed from ``onset`` and ``dur``.
        (Default: None)

    dur : float, optional
        Duration of turn in seconds. If None, then computed from ``onset`` and
        ``offset``.
        (Default: None)

    speaker_id : str, optional
        Speaker id.
        (Default: None)

    file_id : str, optional
        File id.
        (Default: none)
    """

    def __init__(
        self, onset, offset=None, dur=None, speaker_id=None, file_id=None, weight=None
    ):
        if not xor(offset is None, dur is None):
            raise ValueError("Exactly one of offset or dur must be given")
        if onset < 0:
            raise ValueError("Turn onset must be >= 0 seconds")
        if offset:
            dur = offset - onset
        if dur <= 0:
            raise ValueError("Turn duration must be > 0 seconds")
        if not offset:
            offset = onset + dur
        self.onset = onset
        self.offset = offset
        self.dur = dur
        self.speaker_id = speaker_id
        self.file_id = file_id
        self.weight = weight

    def __eq__(self, other):
        if isinstance(other, self.__class__):
            return self.__dict__ == other.__dict__
        return False

    def __ne__(self, other):
        return not self.__eq__(other)

    def __hash__(self):
        return hash((self.onset, self.offset, self.dur, self.file_id, self.speaker_id))

    def __str__(self):
        return "FILE: %s, SPEAKER: %s, ONSET: %f, OFFSET: %f, DUR: %f" % (
            self.file_id,
            self.speaker_id,
            self.onset,
            self.offset,
            self.dur,
        )

    def __repr__(self):
        speaker_id = "'%s'" % self.speaker_id if self.speaker_id is not None else None
        file_id = "'%s'" % self.file_id if self.file_id is not None else None
        return "Turn(%f, %f, None, %s, %s)" % (
            self.onset,
            self.offset,
            speaker_id,
            file_id,
        )


def merge_turns(turns):
    """Merge overlapping turns by same speaker within each file."""
    # Merge separately within each file and for each speaker.
    new_turns = []
    for (file_id, speaker_id), speaker_turns in groupby(
        turns, lambda x: (x.file_id, x.speaker_id)
    ):
        speaker_turns = list(speaker_turns)
        speaker_it = IntervalTree.from_tuples(
            [(turn.onset, turn.offset) for turn in speaker_turns]
        )
        n_turns_pre = len(speaker_it)
        speaker_it.merge_overlaps(strict=False)
        n_turns_post = len(speaker_it)
        if n_turns_post < n_turns_pre:
            speaker_turns = []
            for intrvl in speaker_it:
                speaker_turns.append(
                    Turn(
                        intrvl.begin, intrvl.end, speaker_id=speaker_id, file_id=file_id
                    )
                )
            speaker_turns = sorted(speaker_turns, key=lambda x: (x.onset, x.offset))
            # warn('Merging overlapping speaker turns. '
            #      'FILE: %s, SPEAKER: %s' % (file_id, speaker_id))
        new_turns.extend(speaker_turns)
    return new_turns


def chop_tree(tree, onset, offset):
    """Trim Intervals so that none overlap [``onset``, ``offset``].

    Intervals contained entirely within the chopped region are removed and
    those overlapping, but not contained are trimmed back. Differs from
    ``IntervalTree.chop`` in that it keeps track of which intervals in the
    tree were affected.

    This is an inplace operation.

    Parameters
    ----------
    tree : IntervalTree
        Interval tree.

    onset : float
        Onset of chopped region.

    offset : float
        Offset of chopped region.

    Returns
    -------
    affected_intervals : set of Interval
        Intervals from ``tree`` that overlap chopped region.
    """
    overlapped_intervals = set()  # Intervals overlapping chopped region.
    insertions = set()  # Intervals to add.

    # Identify intervals contained entirely within [onset, offset].
    overlapped_intervals.update(tree.envelop(onset, offset))

    # Identify all other intervals overlapping [onset, offset]. These belong
    # to two classes:
    # - right overlap  --  interval.begin < onset
    # - left overlap  --  offset < interval.end
    for intrvl in tree.at(onset):
        if intrvl.begin >= onset:
            continue
        overlap_dur = intrvl.end - onset
        if not overlap_dur:
            continue
        overlapped_intervals.add(intrvl)
        insertions.add(Interval(intrvl.begin, onset, intrvl.data))
    for intrvl in tree.at(offset):
        if intrvl.end <= offset:
            continue
        overlap_dur = offset - intrvl.begin
        if not overlap_dur:
            continue
        overlapped_intervals.add(intrvl)
        insertions.add(Interval(offset, intrvl.end, intrvl.data))

    # Update tree.
    for intrvl in overlapped_intervals:
        tree.discard(intrvl)
    tree.update(insertions)

    return overlapped_intervals


MAX_SESSION_DUR = 1e6  # Maximum duration (seconds) of session. Any outlandishly
# high number will do.


def trim_turns(turns, uem=None, score_onset=None, score_offset=None):
    """Trim turns to scoring regions defined in UEM.

    Parameters
    ----------
    turns : list of Turn
        Speaker turns.

    uem : UEM, optional
        Un-partitioned evaluation map.
        (Default: None)

    score_onset : float, optional
        Onset of scoring region in seconds from beginning of file. Only valid
        if ``uem=None``.
        (Default: None)

    score_offset : float, optional
        Offset of scoring region in seconds from beginning of file. Only
        valid if ``uem=None``.
        (Default: None)

    Returns
    -------
    trimmed_turns : list of Turn
        Trimmed turns.
    """
    # Validate arguments.
    if uem is not None:
        if not (score_onset is None and score_offset is None):
            raise ValueError(
                "Either uem or score_onset and score_offset must " "be specified."
            )
    else:
        if score_onset is None or score_offset is None:
            raise ValueError(
                "Either uem or score_onset and score_offset must " "be specified."
            )
        if score_onset < 0:
            raise ValueError("Scoring region onset must be >= 0 seconds")
        if score_offset <= score_onset:
            raise ValueError("Scoring region duration must be > 0 seconds")

    # If no UEM provided, set each file to have same scoring region:
    # (score_onset, score_offset).
    if uem is None:
        file_ids = set([turn.file_id for turn in turns])
        uem = UEM({fid: [(score_onset, score_offset)] for fid in file_ids})

    # Trim turns to scoring regions.
    new_turns = []
    for file_id, file_turns in groupby(turns, lambda x: x.file_id):
        if file_id not in uem:
            for turn in file_turns:
                warn("Skipping turn from file not in UEM. TURN: %s" % turn)
            continue

        # Remove overlaps with no score regions.
        noscore_tree = IntervalTree.from_tuples([(0.0, MAX_SESSION_DUR)])
        for score_onset, score_offset in uem[file_id]:
            noscore_tree.chop(score_onset, score_offset)
        turns_tree = IntervalTree.from_tuples(
            (turn.onset, turn.offset, turn) for turn in file_turns
        )
        overlapped_turns = set()  # Turns found to overlap a no score region.
        for noscore_intrvl in noscore_tree:
            overlapped_intrvls = chop_tree(
                turns_tree, noscore_intrvl.begin, noscore_intrvl.end
            )
            overlapped_turns.update([intrvl.data for intrvl in overlapped_intrvls])

        # Convert interval tree to turns.
        for intrvl in turns_tree:
            orig_turn = intrvl.data
            new_turns.append(
                Turn(
                    intrvl.begin,
                    intrvl.end,
                    speaker_id=orig_turn.speaker_id,
                    file_id=orig_turn.file_id,
                )
            )

        # Report any overlapping turns to STDERR.
        for turn in sorted(overlapped_turns, key=lambda x: (x.onset, x.offset)):
            warn("Truncating turn overlapping non-scoring region. TURN: %s" % turn)

    return new_turns
