# Copyright 2019 Google LLC, Andrew Hines
#
# 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.

"""Fits a 3rd order polynomial to map Visqol NSIM to MOS.

  Each invocation will produce a different set of coefficients due to random
  sampling of the
  dataset (unless test_set_ratio is 1.0)

  The output coefficients and scale can be used in visqol source.
  The original 2012 v238 coefficients are included for reference.

  Takes as input 2 files:
    - Both of these files are in the visqol/model/ folder, see the example
    below.
    - Subjective mean MOS score with degraded filename
    - Visqol NSIM results in csv that map degraded filenames to NSIM (Optional)
      - If NSIM methods change, this can be regenerated by modifying
        visqol to output nsim and running visqol in batch mode with csv output,
        e.g.:
          /google/src/head/depot/google3/devtools/blaze/scripts/blaze-run.sh
          //third_party/visqol:main -- --batch_input_csv
          /path/to/tcd-voip/TCD-VOIP/test_set.csv
          --results_csv /tmp/visqol_results_nan.csv --use_speech_mos_mapping
          --verbose
      - The TCD-VOIP test set can be downloaded from TCD directly.  A simple
        script called make_tcd_voip_batch_csv.sh can be used to generate
        test_set.csv is included in the same directory as this script.

  Example invocation:
    blaze run third_party/visqol/scripts:merge_vnsim_mos --
    --mos_csv=$PWD/third_party/visqol/model/tcd_voip_testset_subjective_mos.csv
    --nsim_csv=$PWD/third_party/visqol/model/tcd_voip_testset_speech_nsim.csv
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import csv
import random

from absl import app
from absl import flags
from absl import logging

import numpy as np

FLAGS = flags.FLAGS

flags.DEFINE_string('mos_csv', None,
                    'Path to a csv with degraded, subjective MOS')
flags.DEFINE_string('nsim_csv', None,
                    'Path to a csv with reference, degraded, nsim')
flags.DEFINE_float(
    'train_set_ratio', 0.7,
    'The ratio of data to use for training set.  The remainder is for eval.'
)


def LoadVNSIM(nsim_csv):
  """Returns dictionary with degraded file key and mean nsim value.

  The CSV should have three values: reference path, degraded path, nsim value

  Args:
    nsim_csv: Path to CSV file with NSIM values, format described above.

  Returns:
    Dictionary with degraded file key and NSIM value.
  """
  nsim_dict = {}

  with open(nsim_csv, 'rb') as csvfile:
    nsim_reader = csv.reader(csvfile)
    for row in nsim_reader:
      # Skip header
      if row[2] == 'vnsim':
        continue
      # Keep the degraded file without the directory info as key.
      # This will match what the mos dictionary has.
      deg_file = row[1].split('/')[-1]
      nsim_dict[deg_file] = row[2]
  return nsim_dict


def LoadMOS(mos_csv):
  """Load a csv file withMOS scores.

  Args:
    mos_csv: Path to a csv file that has degraded and MOS.

  Returns:
    Dictionary with filename keys and MOS values
  """
  mos_dict = {}
  with open(mos_csv, 'rb') as csvfile:
    mos_reader = csv.reader(csvfile)
    for row in mos_reader:
      # Skip header
      if row[1] == 'MOS':
        continue
      mos_dict[row[0]] = row[1]
  return mos_dict


def MergeVNSIMAndMOS(mos_dict, nsim_dict):
  """Combines NSIM and MOS CSVs into a dictionary.

  Args:
    mos_dict: Dictionary returned by LoadMOS() function.
    nsim_dict: Dictionary returned by LoadVNSIM() function.

  Returns:
    Dictionary with degraded filename keys and (NSIM, MOS) tuple values.
  """

  # NSIM keys are full paths, wheras MOS keys are filenames
  merged_dict = {}

  for nsim_key, nsim_value in nsim_dict.iteritems():
    # Some nsim values are NaN, skip those.
    if nsim_value == 'nan':
      logging.warn('skipping nan nsim key %s', nsim_key)
      continue

    if nsim_key not in mos_dict:
      logging.warn('skipping missing mos key %s', nsim_key)
      continue

    mos_value = mos_dict[nsim_key]
    merged_dict[nsim_key] = [nsim_value, mos_value]

  return merged_dict


def PolyFitTrain(train_list):
  """Fits a third order polynomial to the input.

  Uses np.polyfit for fitting, which is a least squares fitting.
  Args:
    train_list: A list of tuples with tags and a nested (nsim, mos) tuple.

  Returns:
    Optimal coeficients that fit the input.
  """

  mean_nsim = []
  subj_mos = []

  for _, score in train_list:
    mean_nsim.append(float(score[0]))
    subj_mos.append(float(score[1]))

  return np.polyfit(mean_nsim, subj_mos, 3)


def PolyFitError(test_list, coefs):
  """Calculate mean squared error given 3rd order polynomial coefficients.

  Args:
    test_list: list of tuples with (tag, MOS score).
    coefs: 3rd order polynomial coefficients with np.polyval order.

  Returns:
    The mean of the squared error over all examples.
  """
  errors = []

  for _, scores in test_list:
    subjective_mos = float(scores[1])
    nsim = float(scores[0])
    estimated_mos = np.polyval(coefs, nsim)

    # Additional clamping to MOS range as applied in ViSQOL.
    estimated_mos = max(1., min(5., estimated_mos))

    errors.append((subjective_mos - estimated_mos)**2)
    logging.debug('nsim %d estimatemos %d subjmos %d', nsim, estimated_mos,
                  subjective_mos)
  return np.mean(errors)


def ScaleForPerfectRef(coefs):
  """Calculate a scalar so that maximum NSIM of 1.0 maps to maximum MOS of 5.0.

  Applying this scaling is the default behavior but is optionally disabled via
  a flag in ViSQOL.

  Args:
    coefs: List of coefficients for np.polyval.

  Returns:
    A scalar, which when multiplied by the estimated MOS, scales it to a perfect
    score of 5.0 for a perfect NSIM of 1.0.
  """

  perfect_mos = np.polyval(coefs, 1.0)
  return 1.0 if perfect_mos > 5.0 else 5.0 / perfect_mos


def main(argv):
  del argv  # Unused.

  nsim_dict = LoadVNSIM(FLAGS.nsim_csv)
  mos_dict = LoadMOS(FLAGS.mos_csv)
  merged_dict = MergeVNSIMAndMOS(mos_dict, nsim_dict)

  logging.debug('merged dict: %s', merged_dict)

  merged_list = merged_dict.items()
  train_list = random.sample(merged_list,
                             int(len(merged_list) * FLAGS.train_set_ratio))
  test_list = [x for x in merged_list if x not in train_list]

  logging.info('train size %d, test size %d, total size %d', len(train_list),
               len(test_list), len(merged_list))

  coef_dict = {
      'new_coefs': PolyFitTrain(train_list),
      'v238': [158.7423, -373.5843, 295.5249, -75.2952]
  }

  for coef_label in ['new_coefs', 'v238']:
    coefs = coef_dict[coef_label]
    # formatted for C++ so you can just paste it in to some_nsim_mapper.cc
    logging.info('[%s] 3rd order polynomial coefficients {%f, %f, %f, %f}',
                 coef_label, coefs[0], coefs[1], coefs[2], coefs[3])

    train_error = PolyFitError(train_list, coefs)
    logging.info('[%s] train error %f', coef_label, train_error)

    test_error = PolyFitError(test_list, coefs)
    logging.info('[%s] test error %f', coef_label, test_error)

    scale_ref = ScaleForPerfectRef(coefs)
    logging.info('[%s] scale for perfect ref %f', coef_label, scale_ref)


if __name__ == '__main__':
  app.run(main)
