/*
 * 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.
 */

#ifndef VISQOL_INCLUDE_VISQOLCOMMANDLINE_H
#define VISQOL_INCLUDE_VISQOLCOMMANDLINE_H

#include <memory>
#include <utility>
#include <vector>

#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "comparison_patches_selector.h"
#include "file_path.h"
#include "gammatone_spectrogram_builder.h"
#include "image_patch_creator.h"
#include "similarity_result.h"
#include "src/proto/similarity_result.pb.h"  // Generated by cc_proto_library rule
#include "svr_similarity_to_quality_mapper.h"

namespace Visqol {

/**
 * This class is used for executing ViSQOL comparisons. It is initialised with
 * a similarity to quality mapper which is used for generating quality scores
 * when comparing a given degraded signal to a given reference signal.
 */
class VisqolManager {
 public:
  /**
   * The number of frames in a patch.
   */
  static const size_t kPatchSize;

  /**
   * The number of frames in a patch for speech.
   */
  static const size_t kPatchSizeSpeech;

  /**
   * Number of bands used in the gammatone filter bank for ViSQOL Audio.
   */
  static const size_t kNumBandsAudio;

  /**
   * Number of bands used in the gammatone filter bank for ViSQOL Speech.
   */
  static const size_t kNumBandsSpeech;

  /**
   * Minimum frequency used in the gammatone filter bank.
   */
  static const double kMinimumFreq;

  /**
   * The overlap to be used in the Hamming window e.g 0.5 equals 50% overlap.
   */
  static const double kOverlap;

  /**
   * The maximum delta in duration allowed (measured in seconds) between the
   * reference and degraded signals before a warning will be logged alerting on
   * the mismatch.
   */
  static const double kDurationMismatchTolerance;

  /**
   * Initializes an instance for use with the given similarity to quality
   * mapping model. Must be called before running comparisons.
   *
   * @param similarity_to_quality_mapper_model The filepath to the similarity to
   *    quality mapping model.
   * @param use_speech_mode True if the input signals should be processed as
   *    speech audio. Else, false.
   * @param use_unscaled_speech True if perfect NSIM scores of 1.0 should not
   *    be scaled to a MOS-LQO of 5.0, but instead scaled to ~4.x.
   * @param search_window The search_window parameter determines how far the
   *    comparison algorithm will search to discover the most optimal match for
   *    a given reference patch.
   * @param use_lattice_model If true, use a lattice model for mapping
   *    similarity to quality.
   *
   * @return An 'OK' status if initialised successfully, else an error status.
   */
  absl::Status Init(const FilePath& similarity_to_quality_mapper_model,
                    bool use_speech_mode, bool use_unscaled_speech,
                    int search_window, bool use_lattice_model = true);

  /**
   * Initializes an instance for use with the given similarity to quality
   * mapping model. Must be called before running comparisons.
   *
   * @param similarity_to_quality_mapper_model_string A path string of the
   *    similarity to quality mapping model.
   * @param use_speech_mode True if the input signals should be processed as
   *    speech audio. Else, false.
   * @param use_unscaled_speech True if perfect NSIM scores of 1.0 should not
   *    be scaled to a MOS-LQO of 5.0, but instead scaled to ~4.x.
   * @param search_window The search_window parameter determines how far the
   *    comparison algorithm will search to discover the most optimal match for
   *    a given reference patch.
   * @param use_lattice_model If true, use a lattice model for mapping
   *    similarity to quality.
   *
   * @return An 'OK' status if initialised successfully, else an error status.
   */
  absl::Status Init(absl::string_view similarity_to_quality_mapper_model_string,
                    bool use_speech_mode, bool use_unscaled_speech,
                    int search_window, bool use_lattice_model = true);

  /**
   * Perform a comparison on a single reference/degraded audio file pair.
   *
   * @param ref_signal_path The path to the reference audio file.
   * @param deg_signal_path The path to the degraded audio file.
   *
   * @return A StatusOr object that will contain a SimilarityResultMsg if the
   *    comparison was successful, else it will contain the error Status.
   */
  absl::StatusOr<SimilarityResultMsg> Run(const FilePath& ref_signal_path,
                                          const FilePath& deg_signal_path);

  /**
   * Perform a comparison on a single reference/degraded audio signal pair.
   *
   * @param ref_signal The reference audio signal.
   * @param deg_signal The degraded audio signal.
   *
   * @return A StatusOr object that will contain a SimilarityResultMsg if the
   *    comparison was successful, else it will contain the error Status.
   */
  absl::StatusOr<SimilarityResultMsg> Run(const AudioSignal& ref_signal,
                                          AudioSignal& deg_signal);

 private:
  /**
   * True if the input signals should be processed as speech audio.
   */
  bool use_speech_mode_ = false;

  /**
   * True if similarity to quality mapper should use a deep lattice network
   * model.
   */
  bool use_lattice_model_ = true;

  /**
   * True if perfect NSIM scores of 1.0 should not be scaled to a MOS-LQO of
   * 5.0, but instead scaled to ~4.x.
   */
  bool use_unscaled_speech_mos_mapping_ = false;

  /**
   * True if the object was successfully initialized, else false.
   */
  bool is_initialized_ = false;

  /**
   * This parameter is used to determine how far the algorithm will search in
   * order to find the most optimal match.
   */
  int search_window_ = 60;

  /**
   * Used for creating the patches from both the reference and degraded signals
   * for comparison.
   */
  std::unique_ptr<ImagePatchCreator> patch_creator_;

  /**
   * Used for comparing the patches from the reference signal to the degraded
   * signals.
   */
  std::unique_ptr<ComparisonPatchesSelector> patch_selector_;

  /**
   * Used for builidng spectrogram representations of the reference and
   * degraded signals.
   */
  std::unique_ptr<SpectrogramBuilder> spectrogram_builder_;

  /**
   * Used for generating a quality score for the degraded signal based on the
   * outcome of the spectrogram comparisons.
   */
  std::unique_ptr<SimilarityToQualityMapper> sim_to_qual_;

  /**
   * Initialises the patch creator.
   */
  void InitPatchCreator();

  /**
   * Initialises the comparison patch selector.
   */
  void InitPatchSelector();

  /**
   * Initialises the spectrogram builder.
   */
  void InitSpectrogramBuilder();

  /**
   * Initialises the similiary to quality mapper using the provided model file.
   */
  absl::Status InitSimilarityToQualityMapper(
      FilePath similarity_to_quality_mapper_model);

  /**
   * This helper function can be called to check if the object was initialized
   * correctly.
   *
   * @return An error status if the object was not initialized correctly. Else,
   * an 'ok' status is returned.
   */
  absl::Status ErrorIfNotInitialized();

  /**
   * For a given ViSQOL similarity result, populate a similarity result
   * protobuf message for return.
   *
   * Note that this function will not set the filepaths for the signals that
   * were compared. This is because filepaths are not always supplied for the
   * comparison e.g. if called through the API. This function only sets values
   * which will always be present for a given comparison.
   *
   * @param sim_result The ViSQOL similarity comparison result.
   *
   * @return The similarity result in protobuf format.
   */
  SimilarityResultMsg PopulateSimResultMsg(const SimilarityResult& sim_result);

  /**
   * Validate that the input audio signal meet the necessary requirements.
   *
   * @param ref_signal The reference signal to validate.
   * @param deg_signal The degraded signal to validate.
   *
   * @return An 'OK' status if input signals are valid, else an error status.
   */
  absl::Status ValidateInputAudio(const AudioSignal& ref_signal,
                                  const AudioSignal& deg_signal);
};
}  // namespace Visqol

#endif  // VISQOL_INCLUDE_VISQOLCOMMANDLINE_H
