# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Example Generator Pipeline.

The run method implements a Beam pipeline that generates TFRecord files for
training and evaluation by joining audio files, using the full filename, to
annotations from CSV files, reading the audio files into short clips, and
serializing the labeled clips as TensorFlow Examples.
"""

import csv
import datetime
import functools
import io
import itertools
import os
from typing import BinaryIO, Dict, Iterable, Optional, Tuple, TypeVar, Union

import apache_beam as beam
from apache_beam.io import fileio
from dataclasses import dataclass
from dateutil import parser as date_parser
import intervaltree
from multispecies_whale_detection import dataset
from multispecies_whale_detection import xwav
import numpy as np
import resampy
import soundfile
import tensorflow as tf

T = TypeVar('T')


def _only_element(iterable: Iterable[T], default_value: T = None) -> T:
  """Unwraps an iterator that is expected to have at most one element.

  Args:
    iterable: An Iterable with at most one element.
    default_value: The value to return when the Iterable has no elements.

  Returns:
    The unique element of the Iterable or the default value when the Iterable is
    empty.

  Raises:
    ValueError: if the Iterable has more than one element.
  """
  iterator = iter(iterable)
  element = next(iterator, default_value)
  try:
    _ = next(iterator)
    raise ValueError('Iterable had more than one element')
  except StopIteration:
    pass
  return element


def _parse_utc(text: str) -> datetime.datetime:
  """Leniently parses a datetime, defaulting to UTC if no zone is provided."""
  parsed = date_parser.parse(text)
  if not parsed.tzinfo:
    parsed = parsed.replace(tzinfo=datetime.timezone.utc)
  return parsed


@dataclass
class Configuration:
  """Application settings for the pipeline."""
  input_directory: str
  output_directory: str
  clip_duration_seconds: float = 10
  resample_rate: int = 16000


@dataclass
class ClipMetadata:
  """Description of a clip of audio from a longer file."""
  filename: str
  sample_rate: int
  duration: datetime.timedelta
  index_in_file: int
  start_relative_to_file: datetime.timedelta
  start_utc: Optional[datetime.datetime]


# to annotate a factory method
AnnotationType = TypeVar('AnnotationType', bound='Annotation')


@dataclass
class Annotation:
  """Value object for a labeled time interval within some audio.

  Depending on context, begin and end may be interpreted as relative to the
  start of an audio file, a clip taken from a longer file, or the UNIX epoch.
  """
  label: str

  @classmethod
  def parse_csv_row(cls, row: Dict[str, str]) -> AnnotationType:
    """Parses an annotation from a CSV row.

    Args:
      row: Dictionary mapping CSV headers to field values. The headers must
        include "label" and one of ("begin" / "end" floating-point endpoints in
        seconds relative to the start of the file) or ("begin_utc" / "end_utc"
        endpoints as absolute times as strings in any format that
        dateutil.parser will convert to an aware datetime). When both are
        present, only the endpoints relative to the start of the file are used.

    Returns:
      Annotation parsed from the CSV row.

    Raises:
      ValueError if the given row does not include all the required fields.
    """
    label = row.get('label', None)
    if not label:
      raise ValueError('label was empty or not provided')
    begin_rel_file = row.get('begin', None)
    end_rel_file = row.get('end', None)
    if begin_rel_file and end_rel_file:
      return FileAnnotation(
          label=row['label'],
          begin=datetime.timedelta(seconds=float(begin_rel_file)),
          end=datetime.timedelta(seconds=float(end_rel_file)),
      )
    begin_text = row.get('begin_utc', None)
    end_text = row.get('end_utc', None)
    if begin_text and end_text:
      return UTCAnnotation(
          label=row['label'],
          begin=_parse_utc(begin_text),
          end=_parse_utc(end_text),
      )
    raise ValueError('row should have either both "begin" and "end" fields '
                     'or both "begin_utc" and "end_utc" fields')


@dataclass
class ClipAnnotation(Annotation):
  """Annotation as time differences from the start of a clip.

  This type is used for annotations relative to an extracted clip of audio
  intended to be included in its entirety into a TensorFlow Example.

  This is in contrast to FileAnnotation (below). There, begin / end are in a
  timeline whose zero point is at the start of the file. Here, the zero point
  is at the beginning of a clip which often has been extracted from a longer
  file.

  audio_example packs these ClipAnnotations into parallel "annotation" fields
  of a TensorFlow Example.
  """
  begin: datetime.timedelta
  end: datetime.timedelta


def _restrict_to_clip(
    begin: datetime.timedelta,
    end: datetime.timedelta,
    clip_metadata: ClipMetadata,
    label: str,
) -> Optional[ClipAnnotation]:
  """Restricts an interval to the duration from ClipMetadata.

  Args:
    begin: The start of the interval, relative to the clip. May be negative.
    end: The end of the interval, relative to the clip. May be negative.
    clip_metadata: Description of the clip.
    label: Label to set in the returned ClipAnnotation.

  Returns:
    ClipAnnotation with the intersection of (begin, end) with the clip described
    by clip_metadata or None if that interstion is empty.
  """
  assert end > begin
  begin = max(begin, datetime.timedelta(seconds=0))
  end = min(end, clip_metadata.duration)
  if begin < clip_metadata.duration and end > datetime.timedelta(seconds=0):
    return ClipAnnotation(begin=begin, end=end, label=label)
  else:
    return None


@dataclass
class FileAnnotation(Annotation):
  """Annotation as time differences from the start of a file."""
  begin: datetime.timedelta
  end: datetime.timedelta

  def make_relative(self,
                    clip_metadata: ClipMetadata) -> Optional[ClipAnnotation]:
    """Expresses this annotation as an interval within a given clip.

    Args:
      clip_metadata: Description of the clip, including its position relative to
        the file it came from.

    Returns:
      An annotation relative to the given clip or None if there is no overlap.
    """
    return _restrict_to_clip(
        self.begin - clip_metadata.start_relative_to_file,
        self.end - clip_metadata.start_relative_to_file,
        clip_metadata,
        self.label,
    )


@dataclass
class UTCAnnotation(Annotation):
  """Annotation whose endpoints are absolute datetimes.

  To avoid ambiguity, the datetimes must be time zone aware.
  """
  begin: datetime.datetime
  end: datetime.datetime

  def __init__(self, label, begin, end):
    if not (begin.tzinfo and end.tzinfo):
      raise ValueError('endpoint datetimes must be time zone aware')
    self.label = label
    self.begin = begin
    self.end = end

  def make_relative(self,
                    clip_metadata: ClipMetadata) -> Optional[ClipAnnotation]:
    """Expresses this annotation as an interval within a given clip.

    Args:
      clip_metadata: Description of the clip. start_utc must be set.

    Returns:
      An annotation relative to the given clip or None if there is no overlap.

    Raises:
      ValueError if clip_metadata.start_utc is None.
    """
    return _restrict_to_clip(
        self.begin - clip_metadata.start_utc,
        self.end - clip_metadata.start_utc,
        clip_metadata,