from pyannote.database import get_protocol, FileFinder
protocol = get_protocol(PROTOCOL, preprocessors={"audio": FileFinder()})

"""
Explanation: Overlapped speech detection with pyannote.audio
Overlapped speech detection (OSD) is the task of detecting regions where at least two speakers are speaking at the same time. In this notebook, we will train and evaluate an OSD pipeline on Debug database.
End of explanation
"""


first_training_file = next(protocol.train())

from pyannote.audio.utils.preview import listen
listen(first_training_file)

first_training_file['annotation']

"""
Explanation: pyannote.database protocols usually define 
* a training set: for training_file in protocol.train(): ..., 
* a validation set: for validation_file in protocol.development(): ... 
* an evaluation set for evaluation_file in protocol.test(): ...
Let's listen to the first training file and visualize its reference annotation:
End of explanation
"""


from pyannote.audio.pipelines.overlapped_speech_detection import OracleOverlappedSpeechDetection
oracle_osd = OracleOverlappedSpeechDetection()

oracle_osd(first_training_file).get_timeline()

"""
Explanation: The expected output of a perfect overlapped speech detection pipeline would look like this:
End of explanation
"""


from pyannote.audio.tasks import OverlappedSpeechDetection
osd = OverlappedSpeechDetection(protocol, duration=2., batch_size=16)

"""
Explanation: Training
We initialize an OSD task that describes how the model will be trained:

protocol indicates that we will use files available in protocol.train().
duration=2. and batch_size=16 indicates that the model will ingest batches of 16 two seconds long audio chunks.
End of explanation
"""


from pyannote.audio.models.segmentation.debug import SimpleSegmentationModel
model = SimpleSegmentationModel(task=osd)

"""
Explanation: We initialize the model: it needs to know about the task (task=osd) for which it is being trained for:
End of explanation
"""


import pytorch_lightning as pl
trainer = pl.Trainer(max_epochs=10)
trainer.fit(model, osd)

"""
Explanation: Now that everything is ready, let's train with pytorch-ligthning!
End of explanation
"""


test_file = next(protocol.test())
# here we use a test file provided by the protocol, but it could be any audio file
# e.g. test_file = "/path/to/test.wav".

"""
Explanation: Inference
Once trained, we will apply the model on a test file:
End of explanation
"""


from pyannote.audio import Inference
inference = Inference(model)
osd_probability = inference(test_file)

osd_probability

"""
Explanation: Because the model was trained on 2s audio chunks and that test files are likely to be much longer than that, we wrap the model with an Inference instance: it will take care of sliding a 2s window over the whole file and aggregate the output of the model.
End of explanation
"""


from pyannote.audio.pipelines import OverlappedSpeechDetection as OverlappedSpeechDetectionPipeline
pipeline = OverlappedSpeechDetectionPipeline(scores=inference).instantiate(
    {"onset": 0.5, "offset": 0.5, "min_duration_on": 0.1, "min_duration_off": 0.1})

"""
Explanation: Pipeline
Almost there! To obtain the final overlapped speech regions, we need to apply a detection threshold.
For that, we rely on the overlapped speech detection pipeline whose hyper-parameters are set manually:
- onset=0.5: mark region as active when probability goes above 0.5
- offset=0.5: switch back to inactive when probability goes below 0.5
- min_duration_on=0.1: remove active regions shorter than 100ms
- min_duration_off=0.1: fill inactive regions shorter than 100ms.
End of explanation
"""


pipeline(test_file).get_timeline()

"""
Explanation: Here we go:
End of explanation
"""


# to make things faster, we run the inference once and for all... 
validation_files = list(protocol.development())
for file in validation_files:
    file['osd'] = inference(file)
# ... and tell the pipeline to load OSD scores directly from files
pipeline = OverlappedSpeechDetectionPipeline(scores="osd")

from pyannote.pipeline import Optimizer
optimizer = Optimizer(pipeline)
optimizer.tune(validation_files, n_iterations=200, show_progress=False)

"""
Explanation: Optimizing pipeline hyper-parameters
W can try to optimize the hyper-parameters (that we chose manually above) on the validation set to get better performance.
End of explanation
"""


optimized_pipeline = OverlappedSpeechDetectionPipeline(scores=inference).instantiate(optimizer.best_params)
optimized_pipeline(test_file).get_timeline()

"""
Explanation: There you go: better hyper-parameters that should lead to better results!
End of explanation
"""


