#include "inflect_tts.h"

#include <algorithm>
#include <cstring>
#include <stdexcept>

#include "host_chain.h"

namespace {

// The compiled AXMODELs take S32 token inputs at runtime: the S64
// input_processors of the ONNX graphs are folded into the model at compile
// time (COMPILE_NOTES §3, confirmed by SIMULATE §1).
using token_t = int32_t;

// Index of tensor `name` in `names`; falls back to `fallback` when names are
// unavailable (older runtimes may not report tensor names).
size_t name_index(const std::vector<std::string>& names, const char* name,
                  size_t fallback) {
    const auto it = std::find(names.begin(), names.end(), name);
    if (it != names.end()) {
        return static_cast<size_t>(it - names.begin());
    }
    return fallback;
}

}  // namespace

InflectTTS::InflectTTS(const std::string& encoder_path,
                       const std::string& decoder_path)
    : encoder_(encoder_path), decoder_(decoder_path) {}

std::vector<float> InflectTTS::synthesize_tokens(
    const std::vector<int64_t>& token_ids, float speed, float variation,
    uint64_t seed) {
    if (speed < 0.5f || speed > 2.0f) {
        throw std::invalid_argument("speed must be between 0.5 and 2.0");
    }
    if (variation < 0.0f || variation > 1.0f) {
        throw std::invalid_argument("variation must be between 0.0 and 1.0");
    }
    // intersperse(add_blank): N phonemes -> 2N+1 tokens.
    const int x_len = static_cast<int>(token_ids.size()) * 2 + 1;
    if (x_len > inflect::kEncoderT) {
        throw std::invalid_argument(
            "token sequence exceeds encoder static T=256 (split the text)");
    }

    // ---- encoder: tokens [1,256] zero-padded + x_lengths [1] -------------
    std::vector<token_t> tokens(inflect::kEncoderT, 0);
    for (size_t i = 0; i < token_ids.size(); ++i) {
        tokens[2 * i + 1] = static_cast<token_t>(token_ids[i]);
    }
    token_t x_lengths[1] = {static_cast<token_t>(x_len)};
    const auto enc_in_names = encoder_.input_names();
    std::vector<std::pair<const void*, size_t>> enc_feeds(2);
    enc_feeds[name_index(enc_in_names, "tokens", 0)] =
        {tokens.data(), tokens.size() * sizeof(token_t)};
    enc_feeds[name_index(enc_in_names, "x_lengths", 1)] =
        {x_lengths, sizeof(x_lengths)};
    auto enc_out = encoder_.run(enc_feeds);
    if (enc_out.size() != 3) {
        throw std::runtime_error("encoder must produce 3 outputs (m_p, logs_p, logw)");
    }

    // ---- host chain: durations + generate_path + expansion + noise -------
    const auto enc_names = encoder_.output_names();
    const float* m_p = reinterpret_cast<const float*>(
        enc_out[name_index(enc_names, "m_p", 0)].data());
    const float* logs_p = reinterpret_cast<const float*>(
        enc_out[name_index(enc_names, "logs_p", 1)].data());
    const float* logw = reinterpret_cast<const float*>(
        enc_out[name_index(enc_names, "logw", 2)].data());
    auto priors = inflect::expand_priors(logw, m_p, logs_p, inflect::kEncoderT,
                                         x_len, 1.0f / speed);
    std::vector<float> z_p = inflect::inject_noise(priors, variation, seed);

    // ---- decoder: Tp=512 chunks, overlap crossfade, tail trim ------------
    const int t_prime = priors.t_prime;
    const auto starts = inflect::decoder_chunk_starts(t_prime);
    std::vector<float> out;
    int prev_end = 0;
    for (size_t k = 0; k < starts.size(); ++k) {
        const int start = starts[k];
        const int take = std::min(inflect::kDecoderTp, t_prime - start);
        std::vector<float> z_chunk(
            static_cast<size_t>(inflect::kHiddenChannels) * inflect::kDecoderTp, 0.0f);
        for (int c = 0; c < inflect::kHiddenChannels; ++c) {
            const float* src =
                z_p.data() + static_cast<size_t>(c) * t_prime + start;
            float* dst =
                z_chunk.data() + static_cast<size_t>(c) * inflect::kDecoderTp;
            std::memcpy(dst, src, static_cast<size_t>(take) * sizeof(float));
        }
        std::vector<float> wav = run_decoder_chunk(z_chunk);
        if (k == 0) {
            out = std::move(wav);
        } else {
            inflect::crossfade_append(out, wav, prev_end - start);
        }
        prev_end = start + inflect::kDecoderTp;
    }
    out.resize(static_cast<size_t>(t_prime) * inflect::kHopLength);
    inflect::edge_fade(out);
    inflect::clip_inplace(out);
    return out;
}

std::vector<float> InflectTTS::run_decoder_chunk(
    const std::vector<float>& z_chunk) {
    auto dec_out = decoder_.run(
        {{z_chunk.data(), z_chunk.size() * sizeof(float)}});
    if (dec_out.size() != 1) {
        throw std::runtime_error("decoder must produce 1 output (wav)");
    }
    const size_t n = dec_out[0].size() / sizeof(float);
    std::vector<float> wav(n);
    std::memcpy(wav.data(), dec_out[0].data(), dec_out[0].size());
    return wav;
}