#!/usr/bin/env python
# encoding: utf-8

# The MIT License (MIT)

# Copyright (c) 2016 CNRS

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# AUTHORS
# Hervé BREDIN - http://herve.niderb.fr


from .protocol import Protocol
from tqdm import tqdm


class SpeakerDiarizationProtocol(Protocol):
    """Speaker diarization protocol

    Parameters
    ----------
    preprocessors : dict or (key, preprocessor) iterable
        When provided, each protocol item (dictionary) are preprocessed, such
        that item[key] = preprocessor(item). In case 'preprocessor' is not
        callable, it should be a string containing placeholder for item keys
        (e.g. {'wav': '/path/to/{uri}.wav'})
    """
    def trn_iter(self):
        raise NotImplementedError(
            'Custom speaker diarization protocol should implement "trn_iter".')

    def dev_iter(self):
        raise NotImplementedError(
            'Custom speaker diarization protocol should implement "dev_iter".')

    def tst_iter(self):
        raise NotImplementedError(
            'Custom speaker diarization protocol should implement "tst_iter".')

    def train(self):
        """Iterate over the training set

This will yield dictionaries with the followings keys:

* database: str
  unique database identifier
* uri: str
  uniform (or unique) resource identifier
* annotated: pyannote.core.Timeline
  parts of the resource that were manually annotated
* annotation: pyannote.core.Annotation
  actual annotations

as well as keys coming from the provided preprocessors.

Usage
-----
>>> for item in protocol.train():
...     uri = item['uri']
...     annotated = item['annotated']
...     annotation = item['annotation']
        """

        generator = self.trn_iter()
        if self.progress:
            generator = tqdm(
                generator, desc='Training set',
                total=getattr(self.trn_iter, 'n_items', None))

        for item in generator:
            yield self.preprocess(item)

    def development(self):
        """Iterate over the development set

This will yield dictionaries with the followings keys:

* database: str
  unique database identifier
* uri: str
  uniform (or unique) resource identifier
* annotated: pyannote.core.Timeline
  parts of the resource that were manually annotated
* annotation: pyannote.core.Annotation
  actual annotations

as well as keys coming from the provided preprocessors.

Usage
-----
>>> for item in protocol.development():
...     uri = item['uri']
...     annotated = item['annotated']
...     annotation = item['annotation']
        """

        generator = self.dev_iter()
        if self.progress:
            generator = tqdm(
                generator, desc='Development set',
                total=getattr(self.dev_iter, 'n_items', None))

        for item in generator:
            yield self.preprocess(item)

    def test(self):
        """Iterate over the test set

This will yield dictionaries with the followings keys:

* database: str
  unique database identifier
* uri: str
  uniform (or unique) resource identifier
* annotated: pyannote.core.Timeline
  parts of the resource that were manually annotated
* annotation: pyannote.core.Annotation
  actual annotations

as well as keys coming from the provided preprocessors.

Usage
-----
>>> for item in protocol.test():
...     uri = item['uri']
...     annotated = item['annotated']
...     annotation = item['annotation']
        """

        generator = self.tst_iter()
        if self.progress:
            generator = tqdm(
                generator, desc='Test set',
                total=getattr(self.tst_iter, 'n_items', None))

        for item in generator:
            yield self.preprocess(item)

    def stats(self, subset):
        """Obtain global statistics on a given subset

Parameters
----------
subset : {'train', 'development', 'test'}

Returns
-------
stats : dict
    Dictionary with the followings keys:
    * annotated: float
      total duration (in seconds) of the parts that were manually annotated
    * annotation: float
      total duration (in seconds) of actual (speech) annotations
    * n_files: int
      number of files in the subset
    * labels: dict
      maps speakers with their total speech duration (in seconds)
        """

        annotated = 0.
        annotation = 0.
        n_files = 0
        labels = {}

        for item in getattr(self, subset)():
            annotated += item['annotated'].duration()
            annotation += item['annotation'].get_timeline().duration()
            for label, duration in item['annotation'].chart():
                if label not in labels:
                    labels[label] = 0.
                labels[label] += duration
            n_files += 1

        return {'annotated': annotated,
                'annotation': annotation,
                'n_files': n_files,
                'labels': labels}

@@ -5363,16 +5363,45 @@
 s = %7B%7D%0A%0A
+        lower_bound = False%0A%0A
         
@@ -5433,24 +5433,25 @@
  subset)():%0A
+%0A
             
@@ -5454,19 +5454,240 @@
     
-annotated +
+# increment 'annotation' total duration%0A            annotation += item%5B'annotation'%5D.get_timeline().duration()%0A%0A            # if 'annotated' is provided, use its duration%0A            if 'annotated' in item:%0A                duration 
 = it
@@ -5705,32 +5705,33 @@
 ed'%5D.duration()%0A
+%0A
             anno
@@ -5730,66 +5730,318 @@
     
-annotation += item%5B'annotation'%5D.get_timeline().duration()
+# otherwise, only a lower bound can be computed,%0A            # based on 'annotation' extent%0A            else:%0A                duration = item%5B'annotation'%5D.get_timeline().extent().duration%0A                lower_bound = True%0A%0A            # increment 'annotated' total duration%0A            annotated += duration%0A
 %0A   
@@ -6259,14 +6259,15 @@
     
-return
+stats =
  %7B'a
@@ -6272,18 +6272,19 @@
 'annotat
-ed
+ion
 ': annot
@@ -6285,18 +6285,19 @@
  annotat
-ed
+ion
 ,%0A      
@@ -6310,100 +6310,219 @@
     
-'annotation': annotation,%0A                'n_files': n_files,%0A                'labels': labels%7D
+ 'n_files': n_files,%0A                 'labels': labels%7D%0A%0A        if lower_bound:%0A            stats%5B'annotated_approximate'%5D = annotated%0A        else:%0A            stats%5B'annotated'%5D = annotated%0A%0A        return stats
 %0A

add support for approximate 'annotated' stat