#include <algorithm>
#include <array>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <fstream>
#include <limits>
#include <stdexcept>
#include <string>
#include <vector>

#include "ax_engine_api.h"
#include "ax_sys_api.h"
#include "src/engine_wrapper.hpp"
#include "src/wav_reader.hpp"

#ifndef AXERA_TARGET_NAME
#define AXERA_TARGET_NAME "AXERA"
#endif

namespace {
constexpr int kSampleRate = 16000;
constexpr int kChunkSamples = 1280;
constexpr int kHistorySamples = 480;
constexpr int kFftSize = 512;
constexpr int kSpectrumBins = 257;
constexpr int kMelBins = 32;
constexpr int kMelFrames = 8;
constexpr int kEmbeddingFrames = 76;
constexpr int kEmbeddingSize = 96;
constexpr int kFeatureFrames = 34;
using Clock = std::chrono::steady_clock;

double ElapsedSeconds(Clock::time_point begin, Clock::time_point end) {
  return std::chrono::duration<double>(end - begin).count();
}

struct Args {
  std::string models_dir = "models";
  std::string weights = "config/openwakeword_mel_weights.bin";
  std::string audio = "audio/openwakeword/alexa_test.wav";
  float threshold = 0.5f;
};

void Usage(const char *program) {
  std::printf(
      "Usage: %s [--models-dir DIR] [--mel-weights FILE] [--audio WAV] "
      "[--threshold VALUE]\n",
      program);
}

Args ParseArgs(int argc, char **argv) {
  Args args;
  for (int i = 1; i < argc; ++i) {
    const std::string key = argv[i];
    auto value = [&]() -> std::string {
      if (++i >= argc) throw std::runtime_error("Missing value for " + key);
      return argv[i];
    };
    if (key == "--models-dir") {
      args.models_dir = value();
    } else if (key == "--mel-weights") {
      args.weights = value();
    } else if (key == "--audio") {
      args.audio = value();
    } else if (key == "--threshold") {
      args.threshold = std::stof(value());
    } else if (key == "-h" || key == "--help") {
      Usage(argv[0]);
      std::exit(0);
    } else {
      throw std::runtime_error("Unknown argument: " + key);
    }
  }
  return args;
}

std::string Join(const std::string &left, const std::string &right) {
  return left.empty() || left.back() == '/' ? left + right
                                             : left + "/" + right;
}

class AxRuntime {
 public:
  AxRuntime() {
    if (AX_SYS_Init() != 0) throw std::runtime_error("AX_SYS_Init failed");
    sys_initialized_ = true;
    AX_ENGINE_NPU_ATTR_T attr{};
    if (AX_ENGINE_Init(&attr) != 0) {
      AX_SYS_Deinit();
      sys_initialized_ = false;
      throw std::runtime_error("AX_ENGINE_Init failed");
    }
    engine_initialized_ = true;
  }
  ~AxRuntime() {
    if (engine_initialized_) AX_ENGINE_Deinit();
    if (sys_initialized_) AX_SYS_Deinit();
  }

 private:
  bool sys_initialized_ = false;
  bool engine_initialized_ = false;
};

template <typename T>
T ReadScalar(std::istream &input) {
  T value{};
  input.read(reinterpret_cast<char *>(&value), sizeof(value));
  if (!input) throw std::runtime_error("Truncated mel weight file");
  return value;
}

struct MelWeights {
  std::vector<float> real;
  std::vector<float> imag;
  std::vector<float> mel;
  float floor = 0.0f;

  static MelWeights Load(const std::string &path) {
    std::ifstream input(path, std::ios::binary);
    if (!input) throw std::runtime_error("Cannot open mel weights: " + path);
    char magic[8]{};
    input.read(magic, 8);
    if (std::memcmp(magic, "OWWMEL1", 7) != 0 ||
        ReadScalar<uint32_t>(input) != 1) {
      throw std::runtime_error("Invalid openWakeWord mel weight file");
    }
    const uint32_t real_rows = ReadScalar<uint32_t>(input);
    const uint32_t real_cols = ReadScalar<uint32_t>(input);
    const uint32_t imag_rows = ReadScalar<uint32_t>(input);
    const uint32_t imag_cols = ReadScalar<uint32_t>(input);
    const uint32_t mel_rows = ReadScalar<uint32_t>(input);
    const uint32_t mel_cols = ReadScalar<uint32_t>(input);
    MelWeights result;
    result.floor = ReadScalar<float>(input);
    if (real_rows != kSpectrumBins || real_cols != kFftSize ||
        imag_rows != kSpectrumBins || imag_cols != kFftSize ||
        mel_rows != kSpectrumBins || mel_cols != kMelBins) {
      throw std::runtime_error("Unexpected openWakeWord mel weight shapes");
    }
    result.real.resize(static_cast<std::size_t>(real_rows) * real_cols);
    result.imag.resize(static_cast<std::size_t>(imag_rows) * imag_cols);
    result.mel.resize(static_cast<std::size_t>(mel_rows) * mel_cols);
    input.read(reinterpret_cast<char *>(result.real.data()),
               result.real.size() * sizeof(float));
    input.read(reinterpret_cast<char *>(result.imag.data()),
               result.imag.size() * sizeof(float));
    input.read(reinterpret_cast<char *>(result.mel.data()),
               result.mel.size() * sizeof(float));
    if (!input) throw std::runtime_error("Truncated openWakeWord mel weights");
    return result;
  }
};

std::array<float, kMelFrames * kMelBins> ComputeMel(
    const std::array<float, kHistorySamples + kChunkSamples> &samples,
    const MelWeights &weights) {
  std::array<float, kMelFrames * kMelBins> result{};
  std::array<float, kSpectrumBins> power{};
  float max_db = -std::numeric_limits<float>::infinity();
  for (int frame = 0; frame < kMelFrames; ++frame) {
    const float *frame_samples = samples.data() + frame * 160;
    for (int frequency = 0; frequency < kSpectrumBins; ++frequency) {
      const float *real = weights.real.data() + frequency * kFftSize;
      const float *imag = weights.imag.data() + frequency * kFftSize;
      float real_sum = 0.0f;
      float imag_sum = 0.0f;
      for (int n = 0; n < kFftSize; ++n) {
        real_sum += frame_samples[n] * real[n];
        imag_sum += frame_samples[n] * imag[n];
      }
      power[frequency] = real_sum * real_sum + imag_sum * imag_sum;
    }
    for (int bin = 0; bin < kMelBins; ++bin) {
      float value = 0.0f;
      for (int frequency = 0; frequency < kSpectrumBins; ++frequency) {
        value += power[frequency] *
                 weights.mel[frequency * kMelBins + bin];
      }
      value = std::max(value, weights.floor);
      const float db = std::log(value) * 10.0f / 2.3025851249694824f;
      result[frame * kMelBins + bin] = db;
      max_db = std::max(max_db, db);
    }
  }
  const float minimum = max_db - 80.0f;
  for (float &value : result) {
    value = std::max(value, minimum) / 10.0f + 2.0f;
  }
  return result;
}

struct Classifier {
  std::string name;
  int frames;
  EngineWrapper engine;
  std::vector<float> maximum;
};

void Run(const Args &args) {
  const PcmWav wav = ReadPcmWav(args.audio);
  if (wav.sample_rate != kSampleRate) {
    throw std::runtime_error("Input WAV must use 16 kHz sample rate");
  }
  const double audio_seconds =
      static_cast<double>(wav.samples.size()) / kSampleRate;
  if (audio_seconds <= 0.0) {
    throw std::runtime_error("Input WAV contains no samples");
  }
  const MelWeights weights = MelWeights::Load(args.weights);
  const auto model_load_begin = Clock::now();
  AxRuntime runtime;

  EngineWrapper embedding;
  if (embedding.Init(Join(args.models_dir,
                          "openwakeword__embedding_model.axmodel")) != 0) {
    throw std::runtime_error("Failed to load embedding model");
  }
  const std::array<std::pair<const char *, int>, 6> definitions{{
      {"alexa_v0.1", 16},
      {"hey_jarvis_v0.1", 16},
      {"hey_mycroft_v0.1", 16},
      {"hey_rhasspy_v0.1", 16},
      {"timer_v0.1", 34},
      {"weather_v0.1", 22},
  }};
  std::array<Classifier, 6> classifiers;
  for (std::size_t classifier_index = 0;
       classifier_index < definitions.size(); ++classifier_index) {
    const auto &definition = definitions[classifier_index];
    Classifier &classifier = classifiers[classifier_index];
    classifier.name = definition.first;
    classifier.frames = definition.second;
    const std::string model =
        Join(args.models_dir, "openwakeword__" + classifier.name + ".axmodel");
    if (classifier.engine.Init(model) != 0) {
      throw std::runtime_error("Failed to load classifier: " + classifier.name);
    }
    const int output_bytes = classifier.engine.GetOutputSizeByName(
        classifier.engine.OutputName(0));
    if (output_bytes <= 0 || output_bytes % sizeof(float) != 0) {
      throw std::runtime_error("Unexpected classifier output: " +
                               classifier.name);
    }
    classifier.maximum.assign(output_bytes / sizeof(float),
                              -std::numeric_limits<float>::infinity());
  }
  const double model_load_seconds =
      ElapsedSeconds(model_load_begin, Clock::now());

  std::vector<int16_t> padded = wav.samples;
  const std::size_t remainder = padded.size() % kChunkSamples;
  if (remainder != 0) padded.resize(padded.size() + kChunkSamples - remainder);
  std::array<int16_t, kHistorySamples> history{};
  std::array<float, kEmbeddingFrames * kMelBins> mel_buffer{};
  mel_buffer.fill(1.0f);
  std::array<float, kFeatureFrames * kEmbeddingSize> feature_buffer{};
  int chunks = 0;
  double feature_seconds = 0.0;
  double npu_seconds = 0.0;
  const auto inference_begin = Clock::now();

  for (std::size_t start = 0; start < padded.size(); start += kChunkSamples) {
    std::array<float, kHistorySamples + kChunkSamples> mel_input{};
    for (int i = 0; i < kHistorySamples; ++i) mel_input[i] = history[i];
    for (int i = 0; i < kChunkSamples; ++i) {
      mel_input[kHistorySamples + i] = padded[start + i];
    }
    for (int i = 0; i < kHistorySamples; ++i) {
      history[i] = padded[start + kChunkSamples - kHistorySamples + i];
    }
    const auto feature_begin = Clock::now();
    const auto mel = ComputeMel(mel_input, weights);
    feature_seconds += ElapsedSeconds(feature_begin, Clock::now());
    std::memmove(mel_buffer.data(), mel_buffer.data() + kMelFrames * kMelBins,
                 (kEmbeddingFrames - kMelFrames) * kMelBins * sizeof(float));
    std::memcpy(mel_buffer.data() +
                    (kEmbeddingFrames - kMelFrames) * kMelBins,
                mel.data(), mel.size() * sizeof(float));

    const std::string &embedding_input = embedding.InputName(0);
    if (embedding.SetInputByName(embedding_input, mel_buffer.data(),
                                 mel_buffer.size() * sizeof(float)) != 0) {
      throw std::runtime_error("Failed to set embedding input");
    }
    const auto embedding_begin = Clock::now();
    const int embedding_ret = embedding.RunSync();
    npu_seconds += ElapsedSeconds(embedding_begin, Clock::now());
    if (embedding_ret != 0) {
      throw std::runtime_error("Embedding inference failed");
    }
    std::array<float, kEmbeddingSize> feature{};
    if (embedding.GetOutputByName(embedding.OutputName(0), feature.data(),
                                  feature.size() * sizeof(float)) != 0) {
      throw std::runtime_error("Failed to read embedding output");
    }
    std::memmove(feature_buffer.data(), feature_buffer.data() + kEmbeddingSize,
                 (kFeatureFrames - 1) * kEmbeddingSize * sizeof(float));
    std::memcpy(feature_buffer.data() +
                    (kFeatureFrames - 1) * kEmbeddingSize,
                feature.data(), feature.size() * sizeof(float));

    for (Classifier &classifier : classifiers) {
      const float *input =
          feature_buffer.data() +
          (kFeatureFrames - classifier.frames) * kEmbeddingSize;
      const std::string &input_name = classifier.engine.InputName(0);
      if (classifier.engine.SetInputByName(
              input_name, input,
              classifier.frames * kEmbeddingSize * sizeof(float)) != 0) {
        throw std::runtime_error("Failed to set classifier input: " +
                                 classifier.name);
      }
      const auto classifier_begin = Clock::now();
      const int classifier_ret = classifier.engine.RunSync();
      npu_seconds += ElapsedSeconds(classifier_begin, Clock::now());
      if (classifier_ret != 0) {
        throw std::runtime_error("Classifier inference failed: " +
                                 classifier.name);
      }
      std::vector<float> output(classifier.maximum.size());
      if (classifier.engine.GetOutputByName(
              classifier.engine.OutputName(0), output.data(),
              output.size() * sizeof(float)) != 0) {
        throw std::runtime_error("Failed to read classifier output");
      }
      for (std::size_t i = 0; i < output.size(); ++i) {
        classifier.maximum[i] = std::max(classifier.maximum[i], output[i]);
      }
    }
    ++chunks;
  }
  const double inference_seconds =
      ElapsedSeconds(inference_begin, Clock::now());
  const double rtf = inference_seconds / audio_seconds;

  std::printf("\nopenWakeWord C++ inference complete\n");
  std::printf("target: %s\naudio: %s\nchunks: %d\nthreshold: %.3f\n",
              AXERA_TARGET_NAME, args.audio.c_str(), chunks, args.threshold);
  bool detected = false;
  for (const Classifier &classifier : classifiers) {
    std::printf("%-22s", classifier.name.c_str());
    float score = -std::numeric_limits<float>::infinity();
    for (std::size_t i = 0; i < classifier.maximum.size(); ++i) {
      const float value = classifier.maximum[i];
      std::printf(" %.6f", value);
      if (classifier.name != "timer_v0.1" || i != 0) {
        score = std::max(score, value);
      }
    }
    if (score >= args.threshold) {
      std::printf("  WAKEUP");
      detected = true;
    }
    std::printf("\n");
  }
  std::printf("detected: %s\n", detected ? "true" : "false");
  std::printf("audio_seconds: %.6f\n", audio_seconds);
  std::printf("feature_seconds: %.6f\n", feature_seconds);
  std::printf("npu_seconds: %.6f\n", npu_seconds);
  std::printf("model_load_seconds: %.6f\n", model_load_seconds);
  std::printf("inference_seconds: %.6f\n", inference_seconds);
  std::printf("rtf: %.6f\n", rtf);
}
}  // namespace

int main(int argc, char **argv) {
  try {
    const Args args = ParseArgs(argc, argv);
    Run(args);
    return 0;
  } catch (const std::exception &error) {
    std::fprintf(stderr, "ERROR: %s\n", error.what());
    return 1;
  }
}