#include "host_chain.h"

#include <algorithm>
#include <cmath>
#include <random>
#include <stdexcept>

namespace inflect {

ExpandedPriors expand_priors(const float* logw, const float* m_p,
                             const float* logs_p, int t_total, int x_len,
                             float length_scale) {
    if (x_len < 1 || x_len > t_total) {
        throw std::invalid_argument("x_len out of range [1, t_total]");
    }
    // Duration rounding: w = exp(logw) * x_mask * length_scale; w_ceil = ceil(w).
    // Pure integer cumsum from here on (generate_path).
    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]) * length_scale : 0.0f;
        acc += static_cast<int64_t>(std::ceil(w));
        cum[i] = acc;
    }
    ExpandedPriors priors;
    priors.t_prime = static_cast<int>(std::max<int64_t>(acc, 1));
    priors.m_p.resize(static_cast<size_t>(kHiddenChannels) * priors.t_prime);
    priors.logs_p.resize(static_cast<size_t>(kHiddenChannels) * priors.t_prime);

    // Alignment: output frame t copies input frame src(t), where src(t) is the
    // unique i with cum[i-1] <= t < cum[i]. This is exactly
    // attn @ m_p^T with the 0/1 generate_path alignment.
    int src = 0;
    for (int t = 0; t < priors.t_prime; ++t) {
        while (src < t_total - 1 && cum[src] <= t) {
            ++src;
        }
        if (cum[src] <= t) {
            break;  // t beyond total duration: leave zeros (matches masked path)
        }
        for (int c = 0; c < kHiddenChannels; ++c) {
            const size_t dst = static_cast<size_t>(c) * priors.t_prime + t;
            const size_t s = static_cast<size_t>(c) * t_total + src;
            priors.m_p[dst] = m_p[s];
            priors.logs_p[dst] = logs_p[s];
        }
    }
    return priors;
}

std::vector<float> inject_noise(const ExpandedPriors& priors, float variation,
                                uint64_t seed) {
    std::mt19937_64 rng(seed);
    std::normal_distribution<float> normal(0.0f, 1.0f);
    std::vector<float> z_p(priors.m_p.size());
    for (size_t i = 0; i < z_p.size(); ++i) {
        z_p[i] = priors.m_p[i] + normal(rng) * std::exp(priors.logs_p[i]) * variation;
    }
    return z_p;
}

std::vector<int> decoder_chunk_starts(int t_prime) {
    if (t_prime <= kDecoderTp) {
        return {0};
    }
    const int stride = kDecoderTp - kDecoderOverlap;
    std::vector<int> starts;
    for (int s = 0; s + kDecoderTp <= t_prime; s += stride) {
        starts.push_back(s);
    }
    const int last = t_prime - kDecoderTp;
    if (starts.back() < last) {
        starts.push_back(last);
    }
    return starts;
}

void crossfade_append(std::vector<float>& out, const std::vector<float>& wav,
                      int ov_frames) {
    const size_t ov = static_cast<size_t>(ov_frames) * kHopLength;
    if (ov == 0 || out.size() < ov || wav.size() < ov) {
        out.insert(out.end(), wav.begin(), wav.end());
        return;
    }
    const size_t tail = out.size() - ov;
    for (size_t i = 0; i < ov; ++i) {
        const float t = static_cast<float>(i + 1) / static_cast<float>(ov);
        out[tail + i] = out[tail + i] * (1.0f - t) + wav[i] * t;
    }
    out.insert(out.end(), wav.begin() + ov, wav.end());
}

void edge_fade(std::vector<float>& waveform, float milliseconds) {
    const size_t frames = std::min(
        static_cast<size_t>(kSampleRate * milliseconds / 1000.0f + 0.5f),
        waveform.size() / 2);
    for (size_t i = 0; i < frames; ++i) {
        const float t = static_cast<float>(i + 1) / static_cast<float>(frames);
        waveform[i] *= t;
        waveform[waveform.size() - 1 - i] *= t;
    }
}

void clip_inplace(std::vector<float>& waveform, float lo, float hi) {
    for (float& v : waveform) {
        v = std::min(hi, std::max(lo, v));
    }
}

}  // namespace inflect