/*
 * 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_SIMRESULTSWRITER_H
#define VISQOL_INCLUDE_SIMRESULTSWRITER_H

#include <fstream>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <string>

#include "absl/base/internal/raw_logging.h"
#include "conformance.h"
#include "file_path.h"
#include "google/protobuf/util/json_util.h"
#include "src/proto/similarity_result.pb.h"  // Generated by cc_proto_library rule

namespace Visqol {
class SimilarityResultsWriter {
 public:
  /**
   * Write the results of a single ViSQOL comparison. Results can be written to
   * console, to a JSON file or to both.
   *
   * @param verbose If true, write the results to console.
   * @param results_output_csv If this path is not empty, the basic comparison
   *    results will be written here in CSV format.
   * @param debug_output_path If this path is not empty, the comparison result
   *    will be written to this file in JSON format.
   * @param sim_res_msg The comparison result to write.
   */
  static void Write(const bool verbose, const FilePath& results_output_csv,
                    const FilePath& debug_output_path,
                    const SimilarityResultMsg& sim_res_msg,
                    const bool use_speech_mode, const bool use_lattice) {
    WriteToConsole(sim_res_msg, verbose, use_speech_mode, use_lattice);

    if (!debug_output_path.Path().empty()) {
      WriteDebugJSON(debug_output_path, sim_res_msg);
    }

    if (!results_output_csv.Path().empty()) {
      WriteResultsToCSV(results_output_csv, sim_res_msg);
    }
  }

 private:
  /**
   * Write the results of the comparison, along with some basic debug info, to
   * console.
   *
   * @param sim_res_msg The comparison result to write.
   */
  static void WriteToConsole(const SimilarityResultMsg& sim_res_msg,
                             bool verbose, const bool use_speech_mode,
                             const bool use_lattice) {
    std::cout << "ViSQOL conformance version: " << kVisqolConformanceNumber
              << std::endl;
    std::cout << (use_speech_mode ? "Speech mode, " : "Audio mode, ");
    if (use_speech_mode) {
      std::cout << (use_lattice ? "lattice model" : "exponential model")
                << std::endl;
    } else {
      std::cout << "support vector regression model" << std::endl;
    }
    if (verbose) {
      std::cout << std::endl;
      std::cout << "Reference Filepath:\t" << sim_res_msg.reference_filepath()
                << std::endl;
      std::cout << "Degraded Filepath:\t" << sim_res_msg.degraded_filepath()
                << std::endl;
    }
    std::cout << "MOS-LQO:\t\t" << sim_res_msg.moslqo() << std::endl;
    if (verbose) {
      std::cout << std::endl << FormatFVNSIM(sim_res_msg) << std::endl;
      std::cout << FormatPatchSimilarity(sim_res_msg) << std::endl;
    }
  }

  /**
   * Format the FVNSIM and center frequency band debug info.
   *
   * @param sim_res_msg The similarity result containing the relevant debug
   *    info.
   *
   * @return A string containing the formatted debug info.
   */
  static std::string FormatFVNSIM(const SimilarityResultMsg& sim_res_msg) {
    std::stringstream ss;
    assert(sim_res_msg.fvnsim_size() == sim_res_msg.center_freq_bands_size());
    ss << "---------------------------------------" << std::endl;
    ss << "|  FVNSIM  |  FVNSIM10  |  Freq Band  |" << std::endl;
    ss << "---------------------------------------" << std::endl;
    for (size_t i = 0; i < sim_res_msg.fvnsim_size(); i++) {
      ss << "| " << std::fixed << std::setprecision(6) << std::setw(8)
         << sim_res_msg.fvnsim(i) << " | " << std::fixed << std::setprecision(6)
         << std::setw(10) << sim_res_msg.fvnsim10(i) << " | " << std::fixed
         << std::setprecision(1) << std::setw(9) << std::right
         << sim_res_msg.center_freq_bands(i) << "Hz"
         << " |" << std::endl;
    }
    ss << "---------------------------------------" << std::endl << std::endl;
    return ss.str();
  }

  /**
   * Format the per-patch similarity result debug info.
   *
   * @param sim_res_msg The similarity result containing the relevant debug
   *    info.
   *
   * @return A string containing the formatted debug info.
   */
  static std::string FormatPatchSimilarity(
      const SimilarityResultMsg& sim_res_msg) {
    std::stringstream ss;
    // Format the column titles.
    ss << "--------------------------------------"
          "--------------------------------------"
       << std::endl;
    ss << "| Patch Idx | Similarity | Ref Patch: Start - End "
          "| Deg Patch: Start - End |"
       << std::endl;
    ss << "--------------------------------------"
          "--------------------------------------"
       << std::endl;

    // Format the rows.
    for (size_t i = 0; i < sim_res_msg.patch_sims_size(); i++) {
      // Format patch index.
      ss << "| " << std::setw(9) << std::right
         << i

         // PrFormatint patch similarity score.
         << " | " << std::fixed << std::setprecision(6) << std::setw(10)
         << std::right << sim_res_msg.patch_sims(i).similarity() << std::fixed
         << std::setprecision(3)

         // Format ref patch timestamps.
         << " | " << std::setw(11) << std::right
         << sim_res_msg.patch_sims(i).ref_patch_start_time() << "  - "
         << std::setw(7) << std::right
         << sim_res_msg.patch_sims(i).ref_patch_end_time()

         // Format deg patch timestamps.
         << " | " << std::setw(11) << std::right
         << sim_res_msg.patch_sims(i).deg_patch_start_time() << "  - "
         << std::setw(7) << std::right
         << sim_res_msg.patch_sims(i).deg_patch_end_time() << " |" << std::endl;
    }
    ss << "--------------------------------------"
          "--------------------------------------\n"
       << std::endl;

    return ss.str();
  }

  /**
   * Write the ViSQOL comparison result, including all debug info, to the given
   * file path. The data will be written in JSON format.
   *
   * Data will be appended to the file if the file already has contents.
   *
   * @param debug_output_path The file path to write the JSON results to.
   * @param sim_res_msg The comparison result to write.
   */
  static void WriteDebugJSON(const FilePath& debug_output_path,
                             const SimilarityResultMsg& sim_res_msg) {
    std::string debug_json;
    if (google::protobuf::util::MessageToJsonString(sim_res_msg, &debug_json)
            .ok()) {
      std::ofstream outFile;
      outFile.open(debug_output_path.Path(), std::ios_base::app);
      outFile << debug_json;
      outFile.close();
    } else {
      ABSL_RAW_LOG(ERROR, "Error writing debug JSON: %s ",
                   sim_res_msg.ShortDebugString().c_str());
    }
  }

  /**
   * Write the reference and degraded filepath, along with the resulting
   * MOS-LQO from their comparison, to a CSV file.
   *
   * @param csv_res_path The path to the output CSV file. If the file already
   *    exists, values will be appended to it.
   * @param sim_res_msg The comparison result to write.
   * @param output_moslqo If true, write a column for the mean opinion score.
   * @param If True, write a column with the average nsim value per frequency.
   */
  static void WriteResultsToCSV(const FilePath& csv_res_path,
                                const SimilarityResultMsg& sim_res_msg,
                                const bool output_moslqo = true,
                                const bool output_fvnsim = true,
                                const bool output_stddev = true,
                                const bool output_fvdegenergy = true) {
    // If this file does not already exist, we need to write the header.
    const bool write_header = !csv_res_path.Exists();
    std::ofstream out_file;
    out_file.open(csv_res_path.Path(), std::ios_base::app);

    if (write_header) {
      out_file << "reference,degraded";
      if (output_moslqo) {
        out_file << ",moslqo";
      }

      if (output_fvnsim) {
        for (size_t i = 0; i < sim_res_msg.fvnsim_size(); i++) {
          out_file << ",fvnsim" << i;
        }
      }
      if (output_fvnsim) {
        for (size_t i = 0; i < sim_res_msg.fvnsim10_size(); i++) {
          out_file << ",fvnsim10_" << i;
        }
      }
      if (output_stddev) {
        for (size_t i = 0; i < sim_res_msg.fstdnsim_size(); i++) {
          out_file << ",fstdnsim" << i;
        }
      }
      if (output_fvdegenergy) {
        for (size_t i = 0; i < sim_res_msg.fvdegenergy_size(); i++) {
          out_file << ",fvdegenergy" << i;
        }
      }
      out_file << std::endl;
    }

    out_file << sim_res_msg.reference_filepath() << ","
             << sim_res_msg.degraded_filepath();

    if (output_moslqo) {
      out_file << "," << std::setprecision(9) << sim_res_msg.moslqo();
    }

    if (output_fvnsim) {
      for (size_t i = 0; i < sim_res_msg.fvnsim_size(); i++) {
        out_file << "," << std::setprecision(9) << sim_res_msg.fvnsim(i);
      }
      for (size_t i = 0; i < sim_res_msg.fvnsim10_size(); i++) {
        out_file << "," << std::setprecision(9) << sim_res_msg.fvnsim10(i);
      }
    }
    if (output_stddev) {
      for (size_t i = 0; i < sim_res_msg.fstdnsim_size(); i++) {
        out_file << "," << std::setprecision(9) << sim_res_msg.fstdnsim(i);
      }
    }
    if (output_fvdegenergy) {
      for (size_t i = 0; i < sim_res_msg.fvdegenergy_size(); i++) {
        out_file << "," << std::setprecision(9) << sim_res_msg.fvdegenergy(i);
      }
    }

    out_file << std::endl;
    out_file.close();
  }
};
}  // namespace Visqol

#endif  // VISQOL_INCLUDE_SIMRESULTSWRITER_H
