# Lint as: python3
"""Utils for bluetooth audio testing."""

import logging as log
import math
import os
from typing import Optional
import numpy as np
from scipy import signal as scipy_signal
from scipy.io import wavfile
# Internal import
# Internal import

# Dict keys of the THD+N analysis result
_THDN_KEY = 'thd+n'
_START_TIME_KEY = 'start_time'
_END_TIME_KEY = 'end_time'


def generate_sine_wave_to_device(
    device,
    pushed_file_path='/sdcard/Music',
    frequency=480,
    channel=2,
    sample_rate=48000,
    sample_format=16,
    duration_sec=10):
  """Generates a fixed frequency sine wave file and push it to the device.

  Generates a sine wave to the Mobly device directory and push it to the device
  storage. The output file name format is such as the example:
    sine_480hz_2ch_48000rate_16bit_10sec.wav

  Args:
    device: AndroidDevice, Mobly Android controller class.
    pushed_file_path: string, the wave file path which is pushed to the device
        storage. E.g. /sdcard/Music
    frequency: int, fixed frequency in Hz.
    channel: int, number of channels.
    sample_rate: int, sampling rate in Hz.
    sample_format: int, sampling format in bit.
    duration_sec: int, audio duration in second.

  Returns:
    device_storage_path: string, the wave file on the device storage.
    mobly_directory_path: string, the wave file on the Mobly device directory.
  """
  file_name = 'sine_%dhz_%dch_%drate_%dbit_%dsec.wav' % (
      frequency, channel, sample_rate, sample_format, duration_sec)
  mobly_directory_path = os.path.join(device.log_path, file_name)
  os.system('%s -n -c %d -r %d -b %d %s synth %d sine %d' %
            (audio_processor.AudioProcessor.SOX, channel, sample_rate,
             sample_format, mobly_directory_path, duration_sec, frequency))
  device.adb.push([mobly_directory_path, pushed_file_path])
  device_storage_path = os.path.join(pushed_file_path, file_name)
  return device_storage_path, mobly_directory_path


def measure_audio_mos(recorded_audio_file, reference_audio_file):
  """Measures mean opinion score (MOS) of a recorded audio.

  This function uses the module of A/V Analysis Service to measure MOS:
  Internal reference

  Args:
    recorded_audio_file: string, the recorded audio file to be measured.
    reference_audio_file: string, the reference audio file for comparison.

  Returns:
    Float which is the mean opinion score of the recorded audio.
  """
  results = audio_calculator.AudioAnalyzer().Analyze(reference_audio_file,
                                                     recorded_audio_file)
  # Returns 0.0 if the results fails to be generated.
  if not results:
    log.warning('Failed to generate the audio analysis results.')
    return 0.0
  return results[0].mos


def measure_fundamental_frequency(signal, sample_rate):
  """Measures fundamental frequency of a signal.

  Args:
    signal: An 1-D array representing the signal data.
    sample_rate: int, sample rate of the signal.

  Returns:
    Float representing the fundamental frequency.
  """
  return sample_rate * (np.argmax(np.abs(np.fft.rfft(signal))) / len(signal))


def measure_rms(signal):
  """Measures Root Mean Square (RMS) of a signal.

  Args:
    signal: An 1-D array representing the signal data.

  Returns:
    Float representing the root mean square.
  """
  return np.sqrt(np.mean(np.absolute(signal)**2))


def measure_thdn(signal, sample_rate, q, frequency=None):
  """Measures Total Harmonic Distortion + Noise (THD+N) of a signal.

  Args:
    signal: An 1-D array representing the signal data.
    sample_rate: int, sample rate of the signal.
    q: float, quality factor for the notch filter.
    frequency: float, fundamental frequency of the signal. All other frequencies
        are noise. If not specified, will be calculated using FFT.

  Returns:
    Float representing THD+N ratio calculated from the ratio of RMS of pure
        harmonics and noise signal to RMS of original signal.
  """
  # Normalizes the signal.
  signal -= np.mean(signal)
  # Gets Blackman-Harris window from the signal.
  window = signal * scipy_signal.blackmanharris(len(signal))
  # Finds the fundamental frequency to remove if not specified.
  if not frequency:
    frequency = measure_fundamental_frequency(window, sample_rate)
  # Creates a notch filter to get noise from the signal.
  wo = frequency / (sample_rate / 2)
  b, a = scipy_signal.iirnotch(wo, q)
  noise = scipy_signal.lfilter(b, a, window)
  return measure_rms(noise) / measure_rms(window)


def measure_audio_thdn_per_window(
    audio_file,
    thdn_threshold,
    step_size,
    window_size,
    q=1.0,
    frequency=None):
  """Measures Total Harmonic Distortion + Noise (THD+N) of an audio file.

  This function is used to capture audio glitches from a recorded audio file,
  and the audio file shall record a fixed frequency sine wave.

  Args:
    audio_file: A .wav file to be measured.
    thdn_threshold: float, a THD+N threshold used to compare with the measured
        THD+N for every windows. If THD+N of a window is greater than the
        threshold, will record this to results.
    step_size: int, number of samples to move the window by for each analysis.
    window_size: int, number of samples to analyze each time.
    q: float, quality factor for the notch filter.
    frequency: float, fundamental frequency of the signal. All other frequencies
        are noise. If not specified, will be calculated using FFT.

  Returns:
    List containing each result of channels. Like the following structure:
        ```
          [
              [  # result of channel 1
                  {
                      "thd+n": <float>,  # THD+N of a window
                      "start_time": <float>,  # start time of a window
                      "end_time": <float>,  # end time of a window
                  },
                  ...,
              ],
              [...,]  # result of channel 2
              ...,
          ]
        ```
  """
  if step_size <= 0:
    raise ValueError('step_size shall be greater than 0.')
  if window_size <= 0:
    raise ValueError('window_size shall be greater than 0.')
  sample_rate, wave_data = wavfile.read(audio_file)
  wave_data = wave_data.astype('float64')
  if len(wave_data.shape) == 1:
    channel_signals = (wave_data,)
  else:
    channel_signals = wave_data.transpose()
  # Collects the result for each channels.
  results = []
  for signal in channel_signals:
    current_position = 0
    channel_result = []
    while current_position + window_size <= len(signal):
      window = signal[current_position:current_position + window_size]
      thdn = measure_thdn(
          signal=window,
          sample_rate=sample_rate,
          q=q,
          frequency=frequency)
      start_time = current_position / sample_rate
      end_time = (current_position + window_size) / sample_rate
      if thdn > thdn_threshold:
        channel_result.append({
            _THDN_KEY: thdn,
            _START_TIME_KEY: start_time,
            _END_TIME_KEY: end_time
        })
      current_position += step_size
    results.append(channel_result)
  return results


def get_audio_maximum_thdn(
    audio_file: str,
    step_size: int,
    window_size: int,
    q: float = 1.0,
    frequency: Optional[float] = None) -> float:
  """Gets maximum THD+N from each audio sample with specified window size.

  Args:
    audio_file: A .wav file to be measured.
    step_size: Number of samples to move the window by for each analysis.
    window_size: Number of samples to analyze each time.
    q: Quality factor for the notch filter.
    frequency: Fundamental frequency of the signal. All other frequencies
        are noise. If not specified, will be calculated using FFT.

  Returns:
    Float representing the maximum THD+N for the audio file.
  """
  # Gets all analysis results.
  total_results = measure_audio_thdn_per_window(
      audio_file=audio_file,
      thdn_threshold=-1.0,
      step_size=step_size,
      window_size=window_size,
      q=q,
      frequency=frequency)

  max_thdn_result = {_THDN_KEY: -1.0, _START_TIME_KEY: -1, _END_TIME_KEY: -1}
  for channel_results in total_results:
    for result in channel_results:
      if result[_THDN_KEY] > max_thdn_result[_THDN_KEY]:
        max_thdn_result = result

  log.info('Maximum THD+N result: %s', max_thdn_result)
  return max_thdn_result[_THDN_KEY]


def convert_thdn_percent_to_decibels(thdn_percent: float) -> float:
  """Converts THD+N percentage to decibels (dB).

  Args:
    thdn_percent: THD+N in percent. E.g. 0.001.

  Returns:
    THD+N in decibels.
  """
  return math.log(thdn_percent / 100, 10) * 20


def trim_audio(audio_file: str,
               duration_sec: float,
               start_time_sec: float = 0.0) -> str:
  """Trims an audio file with a specific start time and duration.

  Generates a output file and its name is such as below format:
    `<input file name>_<start time sec>-<duration sec>.<input file type>`

  Args:
    audio_file: string, an audio file to be trimed.
    duration_sec: float, the duration of the output file in seconds.
    start_time_sec: float, the start time of the audio file to be trimmed in
        seconds. Default value is 0.0 second if not specified.

  Returns:
    String, the output file of the same path of the origin file.
  """
  file_path, file_name = os.path.split(audio_file)
  file_name, file_ext = os.path.splitext(file_name)
  output_file_name = '%s_%s-%s%s' % (
      file_name,
      start_time_sec,
      (start_time_sec + duration_sec),
      file_ext)
  output_file = os.path.join(file_path, output_file_name)
  processor = audio_processor.AudioProcessor()
  processor.TrimAudio(
      input_file=audio_file,
      output_file=output_file,
      duration=duration_sec,
      start=start_time_sec)
  return output_file