# Copyright (c) 2017 The WebRTC project authors. All Rights Reserved.
#
# Use of this source code is governed by a BSD-style license
# that can be found in the LICENSE file in the root of the source
# tree. An additional intellectual property rights grant can be found
# in the file PATENTS.  All contributing project authors may
# be found in the AUTHORS file in the root of the source tree.

"""Evaluation score abstract class and implementations.
"""

from __future__ import division
import logging
import os
import re
import subprocess
import sys

try:
  import numpy as np
except ImportError:
  logging.critical('Cannot import the third-party Python package numpy')
  sys.exit(1)

from . import data_access
from . import exceptions
from . import signal_processing


class EvaluationScore(object):

  NAME = None
  REGISTERED_CLASSES = {}

  def __init__(self, score_filename_prefix):
    self._score_filename_prefix = score_filename_prefix
    self._input_signal_metadata = None
    self._reference_signal = None
    self._reference_signal_filepath = None
    self._tested_signal = None
    self._tested_signal_filepath = None
    self._output_filepath = None
    self._score = None

  @classmethod
  def RegisterClass(cls, class_to_register):
    """Registers an EvaluationScore implementation.

    Decorator to automatically register the classes that extend EvaluationScore.
    Example usage:

    @EvaluationScore.RegisterClass
    class AudioLevelScore(EvaluationScore):
      pass
    """
    cls.REGISTERED_CLASSES[class_to_register.NAME] = class_to_register
    return class_to_register

  @property
  def output_filepath(self):
    return self._output_filepath

  @property
  def score(self):
    return self._score

  def SetInputSignalMetadata(self, metadata):
    """Sets input signal metadata.

    Args:
      metadata: dict instance.
    """
    self._input_signal_metadata = metadata

  def SetReferenceSignalFilepath(self, filepath):
    """Sets the path to the audio track used as reference signal.

    Args:
      filepath: path to the reference audio track.
    """
    self._reference_signal_filepath = filepath

  def SetTestedSignalFilepath(self, filepath):
    """Sets the path to the audio track used as test signal.

    Args:
      filepath: path to the test audio track.
    """
    self._tested_signal_filepath = filepath

  def Run(self, output_path):
    """Extracts the score for the set test data pair.

    Args:
      output_path: path to the directory where the output is written.
    """
    self._output_filepath = os.path.join(
        output_path, self._score_filename_prefix + self.NAME + '.txt')
    try:
      # If the score has already been computed, load.
      self._LoadScore()
      logging.debug('score found and loaded')
    except IOError:
      # Compute the score.
      logging.debug('score not found, compute')
      self._Run(output_path)

  def _Run(self, output_path):
    # Abstract method.
    raise NotImplementedError()

  def _LoadReferenceSignal(self):
    assert self._reference_signal_filepath is not None
    self._reference_signal = signal_processing.SignalProcessingUtils.LoadWav(
        self._reference_signal_filepath)

  def _LoadTestedSignal(self):
    assert self._tested_signal_filepath is not None
    self._tested_signal = signal_processing.SignalProcessingUtils.LoadWav(
        self._tested_signal_filepath)


  def _LoadScore(self):
    return data_access.ScoreFile.Load(self._output_filepath)

  def _SaveScore(self):
    return data_access.ScoreFile.Save(self._output_filepath, self._score)


@EvaluationScore.RegisterClass
class AudioLevelPeakScore(EvaluationScore):
  """Peak audio level score.

  Defined as the difference between the peak audio level of the tested and
  the reference signals.

  Unit: dB
  Ideal: 0 dB
  Worst case: +/-inf dB
  """

  NAME = 'audio_level_peak'

  def __init__(self, score_filename_prefix):
    EvaluationScore.__init__(self, score_filename_prefix)

  def _Run(self, output_path):
    self._LoadReferenceSignal()
    self._LoadTestedSignal()
    self._score = self._tested_signal.dBFS - self._reference_signal.dBFS
    self._SaveScore()


@EvaluationScore.RegisterClass
class MeanAudioLevelScore(EvaluationScore):
  """Mean audio level score.

  Defined as the difference between the mean audio level of the tested and
  the reference signals.

  Unit: dB
  Ideal: 0 dB
  Worst case: +/-inf dB
  """

  NAME = 'audio_level_mean'

  def __init__(self, score_filename_prefix):
    EvaluationScore.__init__(self, score_filename_prefix)

  def _Run(self, output_path):
    self._LoadReferenceSignal()
    self._LoadTestedSignal()

    dbfs_diffs_sum = 0.0
    seconds = min(len(self._tested_signal), len(self._reference_signal)) // 1000
    for t in range(seconds):
      t0 = t * seconds
      t1 = t0 + seconds
      dbfs_diffs_sum += (
        self._tested_signal[t0:t1].dBFS - self._reference_signal[t0:t1].dBFS)
    self._score = dbfs_diffs_sum / float(seconds)
    self._SaveScore()


@EvaluationScore.RegisterClass
class PolqaScore(EvaluationScore):
  """POLQA score.

  See http://www.polqa.info/.

  Unit: MOS
  Ideal: 4.5
  Worst case: 1.0
  """

  NAME = 'polqa'

  def __init__(self, score_filename_prefix, polqa_bin_filepath):
    EvaluationScore.__init__(self, score_filename_prefix)

    # POLQA binary file path.
    self._polqa_bin_filepath = polqa_bin_filepath
    if not os.path.exists(self._polqa_bin_filepath):
      logging.error('cannot find POLQA tool binary file')
      raise exceptions.FileNotFoundError()

    # Path to the POLQA directory with binary and license files.
    self._polqa_tool_path, _ = os.path.split(self._polqa_bin_filepath)

  def _Run(self, output_path):
    polqa_out_filepath = os.path.join(output_path, 'polqa.out')
    if os.path.exists(polqa_out_filepath):
      os.unlink(polqa_out_filepath)

    args = [
        self._polqa_bin_filepath, '-t', '-q', '-Overwrite',
        '-Ref', self._reference_signal_filepath,
        '-Test', self._tested_signal_filepath,
        '-LC', 'NB',
        '-Out', polqa_out_filepath,
    ]
    logging.debug(' '.join(args))
    subprocess.call(args, cwd=self._polqa_tool_path)

    # Parse POLQA tool output and extract the score.
    polqa_output = self._ParseOutputFile(polqa_out_filepath)
    self._score = float(polqa_output['PolqaScore'])

    self._SaveScore()

  @classmethod
  def _ParseOutputFile(cls, polqa_out_filepath):
    """
    Parses the POLQA tool output formatted as a table ('-t' option).

    Args:
      polqa_out_filepath: path to the POLQA tool output file.

    Returns:
      A dict.
    """
    data = []
    with open(polqa_out_filepath) as f:
      for line in f:
        line = line.strip()
        if len(line) == 0 or line.startswith('*'):
          # Ignore comments.
          continue
        # Read fields.
        data.append(re.split(r'\t+', line))

    # Two rows expected (header and values).
    assert len(data) == 2, 'Cannot parse POLQA output'
    number_of_fields = len(data[0])
    assert number_of_fields == len(data[1])

    # Build and return a dictionary with field names (header) as keys and the
    # corresponding field values as values.
    return {data[0][index]: data[1][index] for index in range(number_of_fields)}


@EvaluationScore.RegisterClass
class TotalHarmonicDistorsionScore(EvaluationScore):
  """Total harmonic distorsion plus noise score.

  Total harmonic distorsion plus noise score.
  See "https://en.wikipedia.org/wiki/Total_harmonic_distortion#THD.2BN".

  Unit: -.
  Ideal: 0.
  Worst case: +inf
  """

  NAME = 'thd'

  def __init__(self, score_filename_prefix):
    EvaluationScore.__init__(self, score_filename_prefix)
    self._input_frequency = None

  def _Run(self, output_path):
    # TODO(aleloi): Integrate changes made locally.
    self._CheckInputSignal()

    self._LoadTestedSignal()
    if self._tested_signal.channels != 1:
      raise exceptions.EvaluationScoreException(
          'unsupported number of channels')
    samples = signal_processing.SignalProcessingUtils.AudioSegmentToRawData(
        self._tested_signal)

    # Init.
    num_samples = len(samples)
    duration = len(self._tested_signal) / 1000.0
    scaling = 2.0 / num_samples
    max_freq = self._tested_signal.frame_rate / 2
    f0_freq = float(self._input_frequency)
    t = np.linspace(0, duration, num_samples)

    # Analyze harmonics.
    b_terms = []
    n = 1
    while f0_freq * n < max_freq:
      x_n = np.sum(samples * np.sin(2.0 * np.pi * n * f0_freq * t)) * scaling
      y_n = np.sum(samples * np.cos(2.0 * np.pi * n * f0_freq * t)) * scaling
      b_terms.append(np.sqrt(x_n**2 + y_n**2))
      n += 1

    output_without_fundamental = samples - b_terms[0] * np.sin(
        2.0 * np.pi * f0_freq * t)
    distortion_and_noise = np.sqrt(np.sum(
        output_without_fundamental**2) * np.pi * scaling)

    # TODO(alessiob): Fix or remove if not needed.
    # thd = np.sqrt(np.sum(b_terms[1:]**2)) / b_terms[0]

    # TODO(alessiob): Check the range of |thd_plus_noise| and update the class
    # docstring above if accordingly.
    thd_plus_noise = distortion_and_noise / b_terms[0]

    self._score = thd_plus_noise
    self._SaveScore()

  def _CheckInputSignal(self):
    # Check input signal and get properties.
    try:
      if self._input_signal_metadata['signal'] != 'pure_tone':
        raise exceptions.EvaluationScoreException(
            'The THD score requires a pure tone as input signal')
      self._input_frequency = self._input_signal_metadata['frequency']
      if self._input_signal_metadata['test_data_gen_name'] != 'identity' or (
          self._input_signal_metadata['test_data_gen_config'] != 'default'):
        raise exceptions.EvaluationScoreException(
            'The THD score cannot be used with any test data generator other '
            'than "identity"')
    except TypeError:
      raise exceptions.EvaluationScoreException(
          'The THD score requires an input signal with associated metadata')
    except KeyError:
      raise exceptions.EvaluationScoreException(
          'Invalid input signal metadata to compute the THD score')