// inflect_tts_cli — Inflect AX TTS C++ SDK example / smoke tool.
//
// Synthesize from phoneme ids (eSpeak text frontend is intentionally out of
// scope for the C++ SDK — see README):
//
//   ./inflect_tts_cli \
//       --encoder models/ax620e/encoder.axmodel \
//       --decoder models/ax620e/decoder.axmodel \
//       --tokens 81,83,16,53,65,102,53 --output out.wav
//
// Host-side self test of the pure-C++ host chain (no AX runtime needed):
//
//   ./inflect_tts_cli --selftest

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <fstream>
#include <iostream>
#include <random>
#include <stdexcept>
#include <string>
#include <vector>

#include "host_chain.h"
#include "inflect_tts.h"
#include "wav_writer.h"

namespace {

int selftest() {
    int failures = 0;
    auto check = [&](bool ok, const char* name) {
        std::cout << (ok ? "[PASS] " : "[FAIL] ") << name << "\n";
        if (!ok) ++failures;
    };

    // 1) expand_priors == explicit generate_path + matmul reference.
    {
        std::mt19937 rng(42);
        std::normal_distribution<float> norm(0.0f, 1.0f);
        const int t_total = 16, x_len = 12;
        std::vector<float> logw(t_total), m_p(inflect::kHiddenChannels * t_total),
            logs_p(inflect::kHiddenChannels * t_total);
        for (auto& v : logw) v = norm(rng);
        for (auto& v : m_p) v = norm(rng);
        for (auto& v : logs_p) v = norm(rng);
        const float ls = 1.25f;
        auto priors = inflect::expand_priors(logw.data(), m_p.data(), logs_p.data(),
                                             t_total, x_len, ls);

        // Reference: integer cumsum + explicit 0/1 attn matmul.
        std::vector<int64_t> cum(t_total);
        int64_t acc = 0;
        for (int i = 0; i < t_total; ++i) {
            const float w = (i < x_len) ? std::exp(logw[i]) * ls : 0.0f;
            acc += static_cast<int64_t>(std::ceil(w));
            cum[i] = acc;
        }
        const int t_prime = static_cast<int>(std::max<int64_t>(acc, 1));
        check(priors.t_prime == t_prime, "expand_priors: T' == cumsum sum");
        double max_diff = 0.0;
        for (int t = 0; t < t_prime; ++t) {
            for (int c = 0; c < inflect::kHiddenChannels; ++c) {
                float ref = 0.0f;
                for (int i = 0; i < t_total; ++i) {
                    const int64_t lo = (i == 0) ? 0 : cum[i - 1];
                    if (lo <= t && t < cum[i]) {
                        ref = m_p[static_cast<size_t>(c) * t_total + i];
                        break;
                    }
                }
                max_diff = std::max(max_diff,
                    static_cast<double>(std::fabs(
                        priors.m_p[static_cast<size_t>(c) * t_prime + t] - ref)));
            }
        }
        check(max_diff == 0.0, "expand_priors: gather == attn matmul (exact)");
    }

    // 2) Unity durations: identity expansion.
    {
        const int t_total = 8, x_len = 5;
        std::vector<float> logw(t_total, 0.0f);  // exp(0)=1 -> 1 frame each
        std::vector<float> m_p(inflect::kHiddenChannels * t_total),
            logs_p(inflect::kHiddenChannels * t_total);
        for (size_t i = 0; i < m_p.size(); ++i) {
            m_p[i] = static_cast<float>(i % 7) * 0.1f;
            logs_p[i] = -1.0f;
        }
        auto priors = inflect::expand_priors(logw.data(), m_p.data(), logs_p.data(),
                                             t_total, x_len, 1.0f);
        bool ok = priors.t_prime == x_len;
        for (int t = 0; ok && t < x_len; ++t) {
            for (int c = 0; c < inflect::kHiddenChannels; ++c) {
                ok = priors.m_p[static_cast<size_t>(c) * priors.t_prime + t] ==
                     m_p[static_cast<size_t>(c) * t_total + t];
            }
        }
        check(ok, "expand_priors: unity durations == identity");
    }

    // 3) Noise injection determinism.
    {
        inflect::ExpandedPriors priors;
        priors.t_prime = 4;
        priors.m_p.assign(inflect::kHiddenChannels * 4, 0.5f);
        priors.logs_p.assign(inflect::kHiddenChannels * 4, -1.0f);
        const auto a = inflect::inject_noise(priors, 0.667f, 7);
        const auto b = inflect::inject_noise(priors, 0.667f, 7);
        const auto c = inflect::inject_noise(priors, 0.667f, 8);
        check(a == b, "inject_noise: same seed reproducible");
        check(a != c, "inject_noise: different seed differs");
    }

    // 4) Decoder chunk starts: coverage + >=64-frame overlap.
    {
        const auto single = inflect::decoder_chunk_starts(51);
        check(single.size() == 1 && single[0] == 0, "chunk_starts: T'<=512 single");
        const auto multi = inflect::decoder_chunk_starts(812);
        bool ok = multi.size() >= 2 && multi.front() == 0 &&
                  multi.back() + inflect::kDecoderTp >= 812;
        for (size_t i = 1; ok && i < multi.size(); ++i) {
            ok = multi[i] < multi[i - 1] + inflect::kDecoderTp - inflect::kDecoderOverlap + 1;
        }
        check(ok, "chunk_starts: T'=812 coverage + overlap");
    }

    // 5) WAV writer round trip header.
    {
        const std::string path = "inflect_selftest_tmp.wav";
        std::vector<float> wav(2400, 0.25f);
        write_wav(path, wav, 24000);
        std::ifstream f(path, std::ios::binary);
        char riff[4];
        f.read(riff, 4);
        const bool ok = f && riff[0] == 'R' && riff[1] == 'I' && riff[2] == 'F' &&
                        riff[3] == 'F';
        f.close();
        std::remove(path.c_str());
        check(ok, "write_wav: RIFF header");
    }

    std::cout << (failures ? "SELFTEST FAILED" : "SELFTEST PASS") << "\n";
    return failures ? 1 : 0;
}

std::vector<int64_t> parse_tokens(const std::string& csv) {
    std::vector<int64_t> ids;
    size_t pos = 0;
    while (pos <= csv.size()) {
        const size_t comma = csv.find(',', pos);
        const std::string part = csv.substr(pos, comma - pos);
        if (!part.empty()) {
            ids.push_back(std::stoll(part));
        }
        if (comma == std::string::npos) break;
        pos = comma + 1;
    }
    return ids;
}

void usage(const char* argv0) {
    std::cerr
        << "usage:\n"
        << "  " << argv0 << " --selftest\n"
        << "  " << argv0 << " --encoder ENC.axmodel --decoder DEC.axmodel\n"
        << "      --tokens 81,83,16,53,65,102,53 --output out.wav\n"
        << "      [--speed 1.0] [--variation 0.667] [--seed 0]\n";
}

}  // namespace

int main(int argc, char** argv) {
    std::string encoder, decoder, tokens_csv, output;
    float speed = 1.0f, variation = 0.667f;
    uint64_t seed = 0;
    for (int i = 1; i < argc; ++i) {
        const std::string arg = argv[i];
        auto next = [&](const char* name) -> std::string {
            if (i + 1 >= argc) {
                throw std::runtime_error(std::string("missing value for ") + name);
            }
            return argv[++i];
        };
        if (arg == "--selftest") return selftest();
        if (arg == "--encoder") encoder = next("--encoder");
        else if (arg == "--decoder") decoder = next("--decoder");
        else if (arg == "--tokens") tokens_csv = next("--tokens");
        else if (arg == "--output") output = next("--output");
        else if (arg == "--speed") speed = std::stof(next("--speed"));
        else if (arg == "--variation") variation = std::stof(next("--variation"));
        else if (arg == "--seed") seed = std::stoull(next("--seed"));
        else if (arg == "--help" || arg == "-h") { usage(argv[0]); return 0; }
        else {
            std::cerr << "unknown argument: " << arg << "\n";
            usage(argv[0]);
            return 2;
        }
    }
    if (encoder.empty() || decoder.empty() || tokens_csv.empty() || output.empty()) {
        usage(argv[0]);
        return 2;
    }
    try {
        InflectTTS tts(encoder, decoder);
        const auto ids = parse_tokens(tokens_csv);
        auto wav = tts.synthesize_tokens(ids, speed, variation, seed);
        write_wav(output, wav, inflect::kSampleRate);
        std::cout << "wrote " << output << ": " << wav.size() << " samples ("
                  << static_cast<double>(wav.size()) / inflect::kSampleRate
                  << " s @ " << inflect::kSampleRate << " Hz)\n";
    } catch (const std::exception& exc) {
        std::cerr << "error: " << exc.what() << "\n";
        return 1;
    }
    return 0;
}