#!/usr/bin/env python
  # Copyright 2014  Johns Hopkins University (Authors: Vijayaditya Peddinti).  Apache 2.0.
  
  # normalizes the wave files provided in input file list with a common scaling factor
  # the common scaling factor is computed to 1/\sqrt(1/(total_samples) * \sum_i{\sum_j x_i(j)^2}) where total_samples is sum of all samples of all wavefiles. If the data is multi-channel then each channel is treated as a seperate wave files
  from __future__ import division
  from __future__ import print_function
  import argparse, scipy.io.wavfile, warnings, numpy as np, math
  
  def get_normalization_coefficient(file_list, is_rir, additional_scaling):
    assert(len(file_list) > 0)
    sampling_rate = None
    total_energy = 0.0
    total_samples = 0.0
    prev_dtype_max_value = None
    for file in file_list:
      try:
        [rate, data] = scipy.io.wavfile.read(file)
        if not str(data.dtype) in set(['int16', 'int32', 'int64']):
          raise Exception('Cannot process {0}, only wav files of integer type are suppported'.format(file))
  
        dtype_max_value = np.iinfo(data.dtype).max
        # ensure that all the data in the current list is of the same format
        if prev_dtype_max_value is not None:
          assert(dtype_max_value == prev_dtype_max_value)
        prev_dtype_max_value = dtype_max_value
  
        if len(data.shape) == 1:
          data = data.reshape([data.shape[0], 1])
        if sampling_rate is not None:
          assert(rate == sampling_rate)
        else:
          sampling_rate = rate
        data = data/dtype_max_value
        if is_rir:
          # just count the energy of the direct impulse response
          # this is treated as energy of signal from 0.001 seconds before impulse
          # to 0.05 seconds after impulse. This is done as we do not want the 
          # recording length to influence the scaling factor
          channel_one = data[:, 0]
          max_d = max(channel_one)
          delay_impulse = [i for i, j in enumerate(channel_one) if j == max_d][0]
          before_impulse = np.floor(rate * 0.001)
          after_impulse = np.floor(rate * 0.05)
          start_index = int(max(0, delay_impulse - before_impulse))
          end_index = int(min(len(channel_one), delay_impulse + after_impulse))
        else:
          start_index = 0
          end_index = data.shape[0]
        # numpy does not check for numerical overflow in integer type
        # so we convert the data into floats
        data = data.astype(np.float64)
        total_energy += np.sum(data[start_index:end_index, :] ** 2)
        data_shape = list(data.shape)
        data_shape[0] = end_index-start_index
        total_samples += np.prod(data_shape)
      except IOError:
        warnings.warn("Did not find the file {0}.".format(file))
    assert(total_samples > 0)
    scaling_coefficient = np.sqrt(total_samples/total_energy)
    print("Scaling coefficient is {0}.".format(scaling_coefficient))
    if math.isnan(scaling_coefficient):
      raise Exception(" Nan encountered while computing scaling coefficient. This is mostly due to numerical overflow")
    return scaling_coefficient
  
  if __name__ == "__main__":
    usage = """ Python script to normalize input wave file list"""
  
    parser = argparse.ArgumentParser(usage)
    parser.add_argument('--is-room-impulse-response', type=str, default = "false",  help='is the input a list of room impulse responses', choices = ['True', 'False', 'true', 'false'])
    parser.add_argument('--extra-scaling-factor', type=float, default = 1.0,  help='additional scaling factor to be multiplied with the wav files')
    parser.add_argument('input_file_list', type=str, help='list of wav files to be normalized collectively')
    parser.add_argument('output_file', type=str, help='output file to store normalization coefficient')
    params = parser.parse_args() 
    if params.is_room_impulse_response.lower() == 'true':
      params.is_room_impulse_response = True
    else:
      params.is_room_impulse_response = False
  
    file_list = []
    for line in  open(params.input_file_list).readlines():
      if len(line.strip()) > 0 :
        file_list.append(line.strip())
    norm_coefficient = get_normalization_coefficient(file_list, params.is_room_impulse_response, params.extra_scaling_factor)
    out_file = open(params.output_file, 'w')
    out_file.write('{0}'.format(norm_coefficient))
    out_file.close()