/*
  Copyright (C) 2016-2026 Dan Casarin (https://www.kfrlib.com)
  This file is part of KFR

  KFR is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 2 of the License, or
  (at your option) any later version.

  KFR is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with KFR.

  If GPL is not suitable for your project, you must purchase a commercial license to use KFR.
  Buying a commercial license is mandatory as soon as you develop commercial activities without
  disclosing the source code of your own applications.
  See https://www.kfrlib.com for details.
 */
#pragma once

#include <vector>

#include "../base.hpp"
#include "../test/assert.hpp"
#include "biquad.hpp"
#include "biquad_design.hpp"
#include "speaker.hpp"
#include "units.hpp"

KFR_PRAGMA_GNU(GCC diagnostic push)
#if KFR_HAS_WARNING("-Winaccessible-base")
KFR_PRAGMA_GNU(GCC diagnostic ignored "-Winaccessible-base")
#endif

namespace kfr
{
inline namespace KFR_ARCH_NAME
{

/**
 * @brief Convert mean-square energy to loudness in LUFS.
 *
 * Implements the EBU R128 mapping @f$ L = 10 \cdot \log_{10}(z) - 0.691 @f$,
 * where @c z is the K-weighted mean-square energy.
 * @param energy K-weighted mean-square energy (linear, not dB).
 * @return Loudness in LUFS.
 */
template <typename T>
KFR_INTRINSIC T energy_to_loudness(T energy)
{
    return T(10) * log10(energy) - T(0.691);
}

/**
 * @brief Inverse of @ref energy_to_loudness: convert LUFS back to linear energy.
 *
 * @param loudness Loudness in LUFS.
 * @return K-weighted mean-square energy (linear).
 */
template <typename T>
KFR_INTRINSIC T loudness_to_energy(T loudness)
{
    return exp10((loudness + T(0.691)) * T(0.1));
}

/**
 * @brief Buffer of momentary mean-square energies used to compute the
 *        integrated (program) loudness per EBU R128 §3560.
 *
 * Samples below -70 LUFS are discarded on push. The integrated value is
 * computed lazily and cached: it is the mean of the energies that lie above
 * a relative gate set 10 LU below the ungated mean.
 */
template <typename T>
struct integrated_vec : public univector<T>
{
private:
    /** @brief Compute and cache the integrated loudness from the stored energies. */
    void compute() const
    {
        const T z_total = mean(static_cast<const univector<T>&>(*this));
        T relative_gate = energy_to_loudness(z_total) - 10;

        T z        = 0;
        size_t num = 0;
        for (T v : *this)
        {
            T lk = energy_to_loudness(v);
            if (lk >= relative_gate)
            {
                z += v;
                num++;
            }
        }
        z /= num;
        if (num >= 1)
        {
            m_integrated = energy_to_loudness(z);
        }
        else
        {
            m_integrated = -c_infinity<T>;
        }

        m_integrated_cached = true;
    }

public:
    integrated_vec() : m_integrated(-c_infinity<T>), m_integrated_cached(false) {}
    /**
     * @brief Append a momentary mean-square energy sample.
     *
     * Samples whose loudness is below -70 LUFS (absolute gate) are rejected.
     * @param mean_square Momentary mean-square energy for one 400 ms window.
     */
    void push(T mean_square)
    {
        T lk = energy_to_loudness(mean_square);
        if (lk >= T(-70.))
        {
            this->push_back(mean_square);
            m_integrated_cached = false;
        }
    }
    /** @brief Clear the buffer and invalidate the cached result. */
    void reset()
    {
        m_integrated_cached = false;
        this->clear();
    }
    /** @brief Return the integrated loudness in LUFS, recomputing if needed. */
    T get() const
    {
        if (!m_integrated_cached)
        {
            compute();
        }
        return m_integrated;
    }

private:
    mutable T m_integrated;
    mutable bool m_integrated_cached;
};

/**
 * @brief Sorted buffer of short-term mean-square energies used to compute the
 *        Loudness Range (LRA) per EBU R128 / ITU-R BS.1770.
 *
 * Samples below -70 LUFS are rejected on push; the rest are kept in ascending
 * order so that percentile lookups are O(log n). LRA is the difference between
 * the 95 % and 10 % percentiles of the loudness distribution above a relative
 * gate set 20 LU below the ungated mean.
 */
template <typename T>
struct lra_vec : public univector<T>
{
private:
    /** @brief Compute and cache the 10 % / 95 % percentile loudness values. */
    void compute() const
    {
        m_range_high            = -70;
        m_range_low             = -70;
        static const T PRC_LOW  = T(0.10);
        static const T PRC_HIGH = T(0.95);

        const T z_total       = mean(static_cast<const univector<T>&>(*this));
        const T relative_gate = energy_to_loudness(z_total) - 20;

        if (this->size() < 2)
            return;

        const size_t start_index =
            std::upper_bound(this->begin(), this->end(), loudness_to_energy(relative_gate)) - this->begin();
        if (this->size() - start_index < 2)
            return;
        const size_t end_index = this->size() - 1;

        const size_t low_index =
            static_cast<size_t>(std::llround(start_index + (end_index - start_index) * PRC_LOW));
        const size_t high_index =
            static_cast<size_t>(std::llround(start_index + (end_index - start_index) * PRC_HIGH));
        m_range_low  = energy_to_loudness((*this)[low_index]);
        m_range_high = energy_to_loudness((*this)[high_index]);

        m_lra_cached = true;
    }

public:
    lra_vec() : m_range_low(-70), m_range_high(-70), m_lra_cached(false) {}
    /**
     * @brief Insert a short-term mean-square energy sample, keeping the buffer sorted.
     *
     * Samples whose loudness is below -70 LUFS are rejected.
     * @param mean_square Short-term mean-square energy for one 3 s window.
     */
    void push(T mean_square)
    {
        const T lk = energy_to_loudness(mean_square);
        if (lk >= -70)
        {
            auto it = std::upper_bound(this->begin(), this->end(), mean_square);
            this->insert(it, mean_square);
            m_lra_cached = false;
        }
    }
    /** @brief Clear the buffer and invalidate the cached result. */
    void reset()
    {
        m_lra_cached = false;
        this->clear();
    }
    /**
     * @brief Return the LRA percentile loudness bounds.
     * @param low  Out: 10 % percentile loudness in LUFS.
     * @param high Out: 95 % percentile loudness in LUFS.
     * @note LRA = @c high - @c low.
     */
    void get(T& low, T& high) const
    {
        if (!m_lra_cached)
            compute();
        low  = m_range_low;
        high = m_range_high;
    }

private:
    mutable T m_range_low;
    mutable T m_range_high;
    mutable bool m_lra_cached;
};

/**
 * @brief Build the EBU R128 K-weighting pre-filter as an IIR expression handle.
 *
 * The K-filter is a +4 dB high-shelf at 1681.81 Hz followed by a high-pass at
 * 38.1106678 Hz (Q = 0.5, normalized). It must be applied to each channel
 * before measuring energy.
 * @param samplerate Sample rate in Hz.
 * @return Expression handle accepting one input sample block.
 */
template <typename T>
KFR_INTRINSIC expression_handle<T, 1> make_kfilter(int samplerate)
{
    const biquad_section<T> bq[] = {
        biquad_highshelf(T(1681.81 / samplerate), T(+4.0)),
        biquad_highpass(T(38.1106678246655 / samplerate), T(0.5)).normalized_all()
    };
    return to_handle(iir(placeholder<T>(), iir_params{ bq }));
}

template <typename T>
struct ebu_r128;

/**
 * @brief Single-channel loudness measurement front-end.
 *
 * Applies the K-weighting filter to each packet, then accumulates the
 * filtered sum-of-squares into two ring buffers:
 *  - @c m_momentary_sum_of_squares covers 400 ms (momentary loudness, M).
 *  - @c m_short_sum_of_squares covers 3 s (short-term loudness, S).
 *
 * Per-channel output gain accounts for the EBU channel weighting
 * (e.g. +1.5 dB for surrounds, 0 dB for LFE).
 */
template <typename T>
struct ebu_channel
{
public:
    friend struct ebu_r128<T>;
    /**
     * @brief Construct a channel front-end.
     * @param sample_rate         Sample rate in Hz.
     * @param speaker             Speaker role (determines output weighting).
     * @param packet_size_factor  1 = 10 Hz refresh, 2 = 20 Hz, 3 = 30 Hz.
     * @param input_gain          Linear gain applied to the input before filtering.
     */
    ebu_channel(int sample_rate, speaker_type speaker, int packet_size_factor = 1, T input_gain = 1)
        : m_sample_rate(sample_rate), m_speaker(speaker), m_input_gain(input_gain),
          m_packet_size(sample_rate / 10 / packet_size_factor), m_kfilter(make_kfilter<T>(sample_rate)),
          m_short_sum_of_squares(3000 / 100 * packet_size_factor),
          m_momentary_sum_of_squares(400 / 100 * packet_size_factor), m_output_energy_gain(1.0),
          m_buffer_cursor(0), m_short_sum_of_squares_cursor(0), m_momentary_sum_of_squares_cursor(0)
    {
        switch (speaker)
        {
        case speaker_type::Lfe:
        case speaker_type::Lfe2:
            m_output_energy_gain = 0.0;
            break;
        case speaker_type::LeftSurround:
        case speaker_type::RightSurround:
            m_output_energy_gain = static_cast<T>(dB_to_power(+1.5));
            break;
        default:
            break;
        }
        reset();
    }

    /** @brief Zero the momentary and short-term sum-of-squares ring buffers. */
    void reset()
    {
        std::fill(m_short_sum_of_squares.begin(), m_short_sum_of_squares.end(), T(0));
        std::fill(m_momentary_sum_of_squares.begin(), m_momentary_sum_of_squares.end(), T(0));
    }

    /**
     * @brief Process one packet of @ref packet_size() samples.
     *
     * Filters the input through the K-filter and adds the resulting
     * sum-of-squares to both ring buffers (overwriting the oldest slot).
     * @param src Pointer to @ref packet_size() input samples.
     */
    void process_packet(const T* src)
    {
        substitute(m_kfilter, to_handle(make_univector(src, m_packet_size) * m_input_gain));
        const T filtered_sum_of_squares = sumsqr(truncate(m_kfilter, m_packet_size));

        m_short_sum_of_squares.ringbuf_write(m_short_sum_of_squares_cursor, filtered_sum_of_squares);
        m_momentary_sum_of_squares.ringbuf_write(m_momentary_sum_of_squares_cursor, filtered_sum_of_squares);
    }
    /** @brief Return the speaker role assigned to this channel. */
    speaker_type get_speaker() const { return m_speaker; }
    /** @brief Return the number of samples per packet. */
    size_t packet_size() const { return m_packet_size; }

private:
    const int m_sample_rate;
    const speaker_type m_speaker;
    const T m_input_gain;
    const size_t m_packet_size;
    expression_handle<T, 1> m_kfilter;
    univector<T> m_short_sum_of_squares;
    univector<T> m_momentary_sum_of_squares;
    T m_output_energy_gain;
    univector<T> m_buffer;
    size_t m_buffer_cursor;
    size_t m_short_sum_of_squares_cursor;
    size_t m_momentary_sum_of_squares_cursor;
};

/**
 * @brief EBU R128 loudness meter.
 *
 * Computes momentary (M, 400 ms), short-term (S, 3 s), integrated (I) and
 * loudness range (LRA) values for a multi-channel signal, per
 * EBU Tech 3341 / R128 s2 and EBU R128 s3 (LRA).
 *
 * Packets are @c sample_rate/10/ @c packet_size_factor samples long, i.e. the
 * measurement refresh rate is @c 10 * packet_size_factor Hz.
 */
template <typename T>
struct ebu_r128
{
public:
    /**
     * @brief Construct the meter.
     * @param sample_rate         Sample rate in Hz.
     * @param channels            Speaker roles for each input channel (L, R, C, Ls, Rs, LFE, ...).
     * @param packet_size_factor  Refresh multiplier: 1 = 10 Hz, 2 = 20 Hz, 3 = 30 Hz (range [1..6]).
     *
     * One @ref ebu_channel is created per entry in @c channels, in order.
     */
    ebu_r128(int sample_rate, std::span<const speaker_type> channels, int packet_size_factor = 1)
        : m_sample_rate(sample_rate), m_running(true), m_need_reset(false),
          m_packet_size(sample_rate / 10 / packet_size_factor)
    {
        KFR_LOGIC_CHECK(!channels.empty(), "channels must not be empty");
        KFR_LOGIC_CHECK(sample_rate > 0, "sample_rate must be greater than 0");
        KFR_LOGIC_CHECK(packet_size_factor >= 1 && packet_size_factor <= 6,
                        "packet_size_factor must be in range [1..6]");
        for (speaker_type sp : channels)
        {
            m_channels.emplace_back(sample_rate, sp, packet_size_factor, T(1));
        }
    }

    /** @brief Return the configured sample rate in Hz. */
    int sample_rate() const { return m_sample_rate; }

    /** @brief Return the number of samples per packet (= sample_rate / refresh_rate). */
    size_t packet_size() const { return m_packet_size; }

    /**
     * @brief Read all measured loudness values.
     * @param loudness_momentary     Out: momentary loudness M (400 ms), LUFS.
     * @param loudness_short         Out: short-term loudness S (3 s), LUFS.
     * @param loudness_intergrated   Out: integrated program loudness I, LUFS.
     * @param loudness_range_low     Out: LRA 10 % percentile, LUFS.
     * @param loudness_range_high    Out: LRA 95 % percentile, LUFS.
     * @note LRA = @c loudness_range_high - @c loudness_range_low.
     */
    void get_values(T& loudness_momentary, T& loudness_short, T& loudness_intergrated, T& loudness_range_low,
                    T& loudness_range_high)
    {
        T sum_of_mean_square_momentary = 0;
        T sum_of_mean_square_short     = 0;
        for (size_t ch = 0; ch < m_channels.size(); ch++)
        {
            sum_of_mean_square_momentary += mean(m_channels[ch].m_momentary_sum_of_squares) / m_packet_size *
                                            m_channels[ch].m_output_energy_gain;
            sum_of_mean_square_short += mean(m_channels[ch].m_short_sum_of_squares) / m_packet_size *
                                        m_channels[ch].m_output_energy_gain;
        }
        loudness_momentary   = energy_to_loudness(sum_of_mean_square_momentary);
        loudness_short       = energy_to_loudness(sum_of_mean_square_short);
        loudness_intergrated = m_integrated_buffer.get();
        m_lra_buffer.get(loudness_range_low, loudness_range_high);
    }

    /** @brief Access channel @c index (read-only). */
    const ebu_channel<T>& operator[](size_t index) const { return m_channels[index]; }
    /** @brief Return the number of channels. */
    size_t count() const { return m_channels.size(); }

    /**
     * @brief Process one packet for every channel.
     *
     * Each entry of @c source must contain exactly @ref packet_size() samples
     * and correspond, in order, to the channels passed to the constructor.
     * Updates the momentary, short-term, integrated and LRA buffers while
     * @ref running() is true.
     */
    void process_packet(const std::initializer_list<univector_dyn<T>>& source)
    {
        process_packet_impl<tag_dynamic_vector>(source);
    }
    /// @copydoc process_packet(const std::initializer_list<univector_dyn<T>>&)
    void process_packet(const std::initializer_list<univector_ref<T>>& source)
    {
        process_packet_impl<tag_array_ref>(source);
    }
    /// @copydoc process_packet(const std::initializer_list<univector_dyn<T>>&)
    void process_packet(std::span<const univector_dyn<T>> source)
    {
        process_packet_impl<tag_dynamic_vector>(source);
    }
    /// @copydoc process_packet(const std::initializer_list<univector_dyn<T>>&)
    void process_packet(std::span<const univector_ref<T>> source)
    {
        process_packet_impl<tag_array_ref>(source);
    }

    /** @brief (Re)start measurement: subsequent packets update the buffers. */
    void start() { m_running = true; }
    /** @brief Pause measurement: packets are still filtered but buffers are frozen. */
    void stop() { m_running = false; }
    /** @brief Return whether the meter is currently running (see @ref start(), @ref stop()). */
    bool running() const { return m_running; }
    /**
     * @brief Request a deferred reset of all channels and accumulators.
     *
     * The reset is applied at the start of the next @ref process_packet call,
     * so that the in-flight packet is not lost.
     */
    void reset() { m_need_reset = true; }

private:
    /**
     * @brief Per-packet implementation shared by all @ref process_packet overloads.
     *
     * Filters every channel, accumulates momentary/short-term energy, and (when
     * running) pushes those energies into the integrated and LRA buffers. Honors
     * a pending @ref reset() request before accumulating.
     * @param source One packet per channel, in constructor order.
     */
    template <univector_tag Tag>
    void process_packet_impl(std::span<const univector<T, Tag>> source)
    {
        T momentary = 0;
        T shortterm = 0;
        for (size_t ch = 0; ch < m_channels.size(); ch++)
        {
            KFR_ASSERT(source[ch].size() == m_packet_size);
            ebu_channel<T>& chan = m_channels[ch];
            chan.process_packet(source[ch].data());
            if (m_running)
            {
                momentary += mean(m_channels[ch].m_momentary_sum_of_squares) / m_packet_size *
                             m_channels[ch].m_output_energy_gain;
                shortterm += mean(m_channels[ch].m_short_sum_of_squares) / m_packet_size *
                             m_channels[ch].m_output_energy_gain;
            }
        }

        if (m_need_reset)
        {
            m_need_reset = false;
            for (size_t ch = 0; ch < m_channels.size(); ch++)
            {
                m_channels[ch].reset();
            }
            m_integrated_buffer.reset();
            m_lra_buffer.reset();
        }
        if (m_running)
        {
            m_integrated_buffer.push(momentary);
            m_lra_buffer.push(shortterm);
        }
    }
    int m_sample_rate;
    bool m_running;
    bool m_need_reset;
    size_t m_packet_size;
    std::vector<ebu_channel<T>> m_channels;
    integrated_vec<T> m_integrated_buffer;
    lra_vec<T> m_lra_buffer;
};

} // namespace KFR_ARCH_NAME
} // namespace kfr

KFR_PRAGMA_GNU(GCC diagnostic pop)
