// SPDX-License-Identifier: Apache-2.0
//
// quadriga-lib c++/MEX Utility library for radio channel modelling and simulations
// Copyright (C) 2022-2025 Stephan Jaeckel (http://quadriga-lib.org)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ------------------------------------------------------------------------

#include "quadriga_arrayant.hpp"
#include "qd_arrayant_interpolate.hpp"

#include <stdexcept>
#include <string>
#include <vector>
#include <cmath>
#include <algorithm>

#include "slerp.h"

/*!SECTION
Array antenna functions
SECTION!*/

/*!MD
# arrayant_interpolate_multi
Interpolate multi-frequency array antenna patterns at arbitrary frequencies

## Description:
- Computes interpolated complex-valued antenna field patterns at arbitrary spatial angles and
  frequencies from a multi-frequency `arrayant` vector (e.g. a loudspeaker model generated by
  `generate_speaker`).
- For each requested frequency, the function identifies the two arrayant entries with the closest
  center frequencies that bracket the target frequency, performs spatial pattern interpolation on
  both using `qd_arrayant_interpolate`, and then blends the two results in the frequency dimension.
- Frequency interpolation uses spherical interpolation (SLERP) of the complex field values with
  automatic fallback to linear interpolation when the phase difference between the two bracketing
  values exceeds a threshold. This is the same mixed-mode algorithm used by `qd_arrayant_interpolate`
  for spatial grid interpolation, ensuring smooth amplitude and phase transitions.
- If the requested frequency is below the lowest or above the highest `center_frequency` in the
  vector, the nearest entry is used without extrapolation (clamping). This is intentional: a
  subwoofer model queried at 12 kHz returns its 12 kHz response if available, or its highest
  frequency sample otherwise, reflecting the physical roll-off already encoded in the pattern data.
- The function calls `qd_arrayant_interpolate` directly (bypassing the member function wrapper) for
  maximum performance. This avoids redundant validation on each call since the arrayant vector is
  validated once upfront (when `validate_input` is true).
- A simple caching mechanism avoids redundant spatial interpolation when consecutive frequency
  requests share the same bracketing entry.
- Allowed datatypes (`dtype`): `float` or `double`

## Declaration:
```
void quadriga_lib::arrayant_interpolate_multi(
        const std::vector<arrayant<dtype>> &arrayant_vec,
        const arma::Mat<dtype> *azimuth,
        const arma::Mat<dtype> *elevation,
        const arma::Col<dtype> *frequency,
        arma::Cube<dtype> *V_re,
        arma::Cube<dtype> *V_im,
        arma::Cube<dtype> *H_re,
        arma::Cube<dtype> *H_im,
        arma::uvec i_element = {},
        const arma::Cube<dtype> *orientation = nullptr,
        const arma::Mat<dtype> *element_pos_i = nullptr,
        bool validate_input = true)
```

## Arguments:
- `const std::vector<arrayant<dtype>> &**arrayant_vec**` (input)<br>
  Multi-frequency arrayant vector. Each entry must have a valid `center_frequency`. Entries need not
  be sorted by frequency (sorting is handled internally). All entries must share the same angular
  grids and element structure.

- `const arma::Mat<dtype> ***azimuth**` (input)<br>
  Azimuth angles in [rad] for spatial interpolation. Cannot be NULL.<br>
  Size `[1, n_ang]` (same angles for all elements) or `[n_out, n_ang]` (per-element angles).

- `const arma::Mat<dtype> ***elevation**` (input)<br>
  Elevation angles in [rad] for spatial interpolation. Cannot be NULL.<br>
  Size must match `azimuth`.

- `const arma::Col<dtype> ***frequency**` (input)<br>
  Target frequencies in [Hz] at which to evaluate the pattern. Cannot be NULL or empty.<br>
  Size `[n_freq]`. Values need not be sorted.

- `arma::Cube<dtype> ***V_re**` (output)<br>
  Real part of the interpolated e-theta (vertical) field component.<br>
  Size `[n_out, n_ang, n_freq]`. Resized automatically if needed. Cannot be NULL.

- `arma::Cube<dtype> ***V_im**` (output)<br>
  Imaginary part of the interpolated e-theta (vertical) field component.<br>
  Size `[n_out, n_ang, n_freq]`. Cannot be NULL.

- `arma::Cube<dtype> ***H_re**` (output)<br>
  Real part of the interpolated e-phi (horizontal) field component.<br>
  Size `[n_out, n_ang, n_freq]`. Cannot be NULL.

- `arma::Cube<dtype> ***H_im**` (output)<br>
  Imaginary part of the interpolated e-phi (horizontal) field component.<br>
  Size `[n_out, n_ang, n_freq]`. Cannot be NULL.

- `arma::uvec **i_element** = {}` (optional input)<br>
  0-based element indices to interpolate. If empty, all elements are included (`n_out = n_elements`).

- `const arma::Cube<dtype> ***orientation** = nullptr` (optional input)<br>
  Array antenna orientation (bank, tilt, heading) in [rad]. Applied identically at all frequencies.
  Sizes: `nullptr`, `[3,1,1]`, `[3,n_out,1]`, `[3,1,n_ang]`, or `[3,n_out,n_ang]`.

- `const arma::Mat<dtype> ***element_pos_i** = nullptr` (optional input)<br>
  Alternative element positions in [m], Size `[3, n_out]`. If `nullptr`, positions from the first
  arrayant entry are used (positions are identical across all frequency entries by construction).

- `bool **validate_input** = true` (optional input)<br>
  If `true`, the arrayant vector is validated with `arrayant_is_valid_multi` before processing. Set
  to `false` to skip validation in performance-critical loops where the vector is known to be valid.

## Performance Notes:
- The function caches spatial interpolation results to avoid redundant `qd_arrayant_interpolate` calls
  when consecutive frequency requests share the same bracketing entry. For best cache utilization,
  provide the `frequency` vector sorted in ascending or descending order.
- For single-frequency queries that exactly match a `center_frequency`, no frequency interpolation is
  performed (direct copy from the spatial interpolation result).
- For maximum throughput in RT simulations, set `validate_input = false` after initial validation.

## Example:
```
// Build a 2-way speaker model
arma::vec freqs_model = {100.0, 200.0, 500.0, 1000.0, 2000.0, 5000.0, 10000.0};
auto woofer = quadriga_lib::generate_speaker<double>("piston", 0.083, 50.0, 3000.0,
    12.0, 24.0, 87.0, "hemisphere", 0.0, 0.0, 0.0, 0.20, 0.30, freqs_model, 10.0);
auto tweeter = quadriga_lib::generate_speaker<double>("piston", 0.013, 1500.0, 20000.0,
    24.0, 12.0, 90.0, "hemisphere", 0.0, 0.0, 0.0, 0.20, 0.30, freqs_model, 10.0);
auto speaker = quadriga_lib::arrayant_concat_multi(woofer, tweeter);

// Query at 4 specific directions, 3 arbitrary frequencies
arma::mat azimuth = {0.0, 1.5708, -1.5708, 3.14159};
arma::mat elevation(1, 4, arma::fill::zeros);
arma::vec query_freqs = {250.0, 1500.0, 8000.0};

arma::cube V_re, V_im, H_re, H_im;
quadriga_lib::arrayant_interpolate_multi(speaker, &azimuth, &elevation, &query_freqs,
                                          &V_re, &V_im, &H_re, &H_im);

// V_re has size [2, 4, 3]:  2 elements × 4 angles × 3 frequencies
// V_re.slice(0) = response at 250 Hz (interpolated between 200 and 500 Hz entries)
// V_re.slice(1) = response at 1500 Hz (interpolated between 1000 and 2000 Hz entries)
// V_re.slice(2) = response at 8000 Hz (interpolated between 5000 and 10000 Hz entries)
```

## See also:
- <a href="#.interpolate">arrayant.interpolate</a>
- <a href="#arrayant_concat_multi">arrayant_concat_multi</a>
- <a href="#generate_speaker">generate_speaker</a>
- <a href="#arrayant_is_valid_multi">arrayant_is_valid_multi</a>
MD!*/

template <typename dtype>
void quadriga_lib::arrayant_interpolate_multi(const std::vector<quadriga_lib::arrayant<dtype>> &arrayant_vec,
                                              const arma::Mat<dtype> *azimuth,
                                              const arma::Mat<dtype> *elevation,
                                              const arma::Col<dtype> *frequency,
                                              arma::Cube<dtype> *V_re, arma::Cube<dtype> *V_im,
                                              arma::Cube<dtype> *H_re, arma::Cube<dtype> *H_im,
                                              arma::uvec i_element,
                                              const arma::Cube<dtype> *orientation,
                                              const arma::Mat<dtype> *element_pos_i,
                                              bool validate_input)
{
    // --- Input validation ---
    if (arrayant_vec.empty())
        throw std::invalid_argument("arrayant_interpolate_multi: Input vector is empty.");
    if (azimuth == nullptr)
        throw std::invalid_argument("arrayant_interpolate_multi: Input 'azimuth' cannot be NULL.");
    if (elevation == nullptr)
        throw std::invalid_argument("arrayant_interpolate_multi: Input 'elevation' cannot be NULL.");
    if (frequency == nullptr || frequency->n_elem == 0)
        throw std::invalid_argument("arrayant_interpolate_multi: Input 'frequency' cannot be NULL or empty.");
    if (V_re == nullptr || V_im == nullptr || H_re == nullptr || H_im == nullptr)
        throw std::invalid_argument("arrayant_interpolate_multi: Outputs 'V_re', 'V_im', 'H_re', 'H_im' cannot be NULL.");

    if (validate_input)
    {
        std::string err = quadriga_lib::arrayant_is_valid_multi(arrayant_vec, false);
        if (!err.empty())
            throw std::invalid_argument("arrayant_interpolate_multi: Input validation failed: " + err);
    }

    arma::uword n_entries = (arma::uword)arrayant_vec.size();
    arma::uword n_ang = azimuth->n_cols;
    arma::uword n_freq = frequency->n_elem;

    // Determine n_elements and n_out
    arma::uword n_elements = arrayant_vec[0].e_theta_re.n_slices;
    if (i_element.is_empty())
        i_element = arma::regspace<arma::uvec>(0, n_elements - 1);
    else
    {
        if (arma::any(i_element >= n_elements))
            throw std::invalid_argument("arrayant_interpolate_multi: Element indices exceed array antenna size.");
    }

    arma::uword n_out = i_element.n_elem;

    if (azimuth->n_rows != 1 && azimuth->n_rows != n_out)
        throw std::invalid_argument("arrayant_interpolate_multi: Number of rows in 'azimuth' must be 1 or n_out.");
    if (elevation->n_rows != azimuth->n_rows || elevation->n_cols != n_ang)
        throw std::invalid_argument("arrayant_interpolate_multi: Sizes of 'azimuth' and 'elevation' do not match.");

    // Validate orientation
    if (orientation != nullptr && orientation->n_elem != 0)
    {
        if (orientation->n_rows != 3)
            throw std::invalid_argument("arrayant_interpolate_multi: 'orientation' must have 3 rows.");
        if (orientation->n_cols != 1 && orientation->n_cols != n_out)
            throw std::invalid_argument("arrayant_interpolate_multi: 'orientation' must have 1 or n_out columns.");
        if (orientation->n_slices != 1 && orientation->n_slices != n_ang)
            throw std::invalid_argument("arrayant_interpolate_multi: 'orientation' must have 1 or n_ang slices.");
    }

    // Validate alternative element positions
    if (element_pos_i != nullptr && element_pos_i->n_elem != 0)
    {
        if (element_pos_i->n_rows != 3 || element_pos_i->n_cols != n_out)
            throw std::invalid_argument("arrayant_interpolate_multi: 'element_pos_i' must have 3 rows and n_out columns.");
    }

    // --- Resize output cubes ---
    if (V_re->n_rows != n_out || V_re->n_cols != n_ang || V_re->n_slices != n_freq)
        V_re->set_size(n_out, n_ang, n_freq);
    if (V_im->n_rows != n_out || V_im->n_cols != n_ang || V_im->n_slices != n_freq)
        V_im->set_size(n_out, n_ang, n_freq);
    if (H_re->n_rows != n_out || H_re->n_cols != n_ang || H_re->n_slices != n_freq)
        H_re->set_size(n_out, n_ang, n_freq);
    if (H_im->n_rows != n_out || H_im->n_cols != n_ang || H_im->n_slices != n_freq)
        H_im->set_size(n_out, n_ang, n_freq);

    // --- Build sorted frequency index from arrayant center_frequencies ---
    arma::Col<dtype> cf(n_entries);
    for (arma::uword i = 0; i < n_entries; ++i)
        cf[i] = (dtype)arrayant_vec[i].center_frequency;

    // Build sort index (center frequencies should already be sorted, but be safe)
    arma::uvec sort_idx = arma::sort_index(cf);
    arma::Col<dtype> cf_sorted(n_entries);
    for (arma::uword i = 0; i < n_entries; ++i)
        cf_sorted[i] = cf[sort_idx[i]];

    // --- Prepare element index in 1-based unsigned format for qd_arrayant_interpolate ---
    arma::Col<unsigned> i_element_1based(n_out);
    {
        unsigned *pi = i_element_1based.memptr();
        arma::uword *pu = i_element.memptr();
        for (arma::uword i = 0; i < n_out; ++i)
            pi[i] = (unsigned)pu[i] + 1;
    }

    // --- Prepare element positions ---
    arma::Mat<dtype> element_pos_local(3, n_out);
    if (element_pos_i != nullptr && element_pos_i->n_elem != 0)
    {
        std::memcpy(element_pos_local.memptr(), element_pos_i->memptr(), 3 * n_out * sizeof(dtype));
    }
    else if (!arrayant_vec[0].element_pos.is_empty())
    {
        const dtype *ptrI = arrayant_vec[0].element_pos.memptr();
        dtype *ptrO = element_pos_local.memptr();
        for (arma::uword i = 0; i < n_out; ++i)
            std::memcpy(&ptrO[3 * i], &ptrI[3 * i_element[i]], 3 * sizeof(dtype));
    }

    // --- Prepare orientation ---
    arma::Cube<dtype> orientation_empty(3, 1, 1, arma::fill::zeros);
    const arma::Cube<dtype> *orientation_local = (orientation == nullptr || orientation->n_elem == 0)
                                                     ? &orientation_empty
                                                     : orientation;

    // --- Allocate temporary buffers for per-frequency interpolation ---
    arma::uword n_vals = n_out * n_ang;

    // We need at most 2 sets of interpolation buffers (lower and upper bracket)
    arma::Mat<dtype> V_re_lo(n_out, n_ang), V_im_lo(n_out, n_ang);
    arma::Mat<dtype> H_re_lo(n_out, n_ang), H_im_lo(n_out, n_ang);
    arma::Mat<dtype> V_re_hi(n_out, n_ang), V_im_hi(n_out, n_ang);
    arma::Mat<dtype> H_re_hi(n_out, n_ang), H_im_hi(n_out, n_ang);

    // Cache: track which arrayant index is currently in lo/hi buffers to avoid redundant interpolation
    arma::sword cached_lo = -1, cached_hi = -1;

    // --- Process each requested frequency ---
    const dtype *p_freq = frequency->memptr();

    for (arma::uword fi = 0; fi < n_freq; ++fi)
    {
        dtype f = p_freq[fi];

        // Find bracketing indices in sorted center_frequency list
        arma::uword idx_lo, idx_hi;
        dtype w; // Interpolation weight: 0.0 = use lo, 1.0 = use hi

        if (n_entries == 1)
        {
            // Single entry: no frequency interpolation
            idx_lo = sort_idx[0];
            idx_hi = idx_lo;
            w = dtype(0.0);
        }
        else if (f <= cf_sorted[0])
        {
            // Below lowest: clamp to lowest
            idx_lo = sort_idx[0];
            idx_hi = idx_lo;
            w = dtype(0.0);
        }
        else if (f >= cf_sorted[n_entries - 1])
        {
            // Above highest: clamp to highest
            idx_lo = sort_idx[n_entries - 1];
            idx_hi = idx_lo;
            w = dtype(0.0);
        }
        else
        {
            // Find the interval [cf_sorted[k], cf_sorted[k+1]] containing f
            arma::uword k = 0;
            while (k + 1 < n_entries && cf_sorted[k + 1] <= f)
                ++k;

            idx_lo = sort_idx[k];
            idx_hi = sort_idx[k + 1];
            dtype f_lo = cf_sorted[k];
            dtype f_hi = cf_sorted[k + 1];
            w = (f - f_lo) / (f_hi - f_lo);
        }

        // --- Interpolate lower bracket (if not cached) ---
        if ((arma::sword)idx_lo != cached_lo)
        {
            const auto &ant_lo = arrayant_vec[idx_lo];
            qd_arrayant_interpolate<dtype>(ant_lo.e_theta_re, ant_lo.e_theta_im,
                                           ant_lo.e_phi_re, ant_lo.e_phi_im,
                                           ant_lo.azimuth_grid, ant_lo.elevation_grid,
                                           *azimuth, *elevation,
                                           i_element_1based, *orientation_local, element_pos_local,
                                           V_re_lo, V_im_lo, H_re_lo, H_im_lo);
            cached_lo = (arma::sword)idx_lo;
        }

        if (idx_lo == idx_hi || w < std::numeric_limits<dtype>::epsilon())
        {
            // No frequency interpolation needed: copy lo directly to output slice
            std::memcpy(V_re->slice_memptr(fi), V_re_lo.memptr(), n_vals * sizeof(dtype));
            std::memcpy(V_im->slice_memptr(fi), V_im_lo.memptr(), n_vals * sizeof(dtype));
            std::memcpy(H_re->slice_memptr(fi), H_re_lo.memptr(), n_vals * sizeof(dtype));
            std::memcpy(H_im->slice_memptr(fi), H_im_lo.memptr(), n_vals * sizeof(dtype));
        }
        else
        {
            // --- Interpolate upper bracket (if not cached) ---
            if ((arma::sword)idx_hi != cached_hi)
            {
                const auto &ant_hi = arrayant_vec[idx_hi];
                qd_arrayant_interpolate<dtype>(ant_hi.e_theta_re, ant_hi.e_theta_im,
                                               ant_hi.e_phi_re, ant_hi.e_phi_im,
                                               ant_hi.azimuth_grid, ant_hi.elevation_grid,
                                               *azimuth, *elevation,
                                               i_element_1based, *orientation_local, element_pos_local,
                                               V_re_hi, V_im_hi, H_re_hi, H_im_hi);
                cached_hi = (arma::sword)idx_hi;
            }

            // --- Frequency interpolation using spherical interpolation with linear fallback ---
            const dtype *pVr_lo = V_re_lo.memptr(), *pVi_lo = V_im_lo.memptr();
            const dtype *pHr_lo = H_re_lo.memptr(), *pHi_lo = H_im_lo.memptr();
            const dtype *pVr_hi = V_re_hi.memptr(), *pVi_hi = V_im_hi.memptr();
            const dtype *pHr_hi = H_re_hi.memptr(), *pHi_hi = H_im_hi.memptr();
            dtype *pVr_out = V_re->slice_memptr(fi), *pVi_out = V_im->slice_memptr(fi);
            dtype *pHr_out = H_re->slice_memptr(fi), *pHi_out = H_im->slice_memptr(fi);

            for (arma::uword j = 0; j < n_vals; ++j)
            {
                slerp_complex_mf(pVr_lo[j], pVi_lo[j], pVr_hi[j], pVi_hi[j], w, pVr_out[j], pVi_out[j]);
                slerp_complex_mf(pHr_lo[j], pHi_lo[j], pHr_hi[j], pHi_hi[j], w, pHr_out[j], pHi_out[j]);
            }
        }
    }
}

template void quadriga_lib::arrayant_interpolate_multi(
    const std::vector<quadriga_lib::arrayant<float>> &, const arma::Mat<float> *, const arma::Mat<float> *,
    const arma::Col<float> *, arma::Cube<float> *, arma::Cube<float> *, arma::Cube<float> *, arma::Cube<float> *,
    arma::uvec, const arma::Cube<float> *, const arma::Mat<float> *, bool);

template void quadriga_lib::arrayant_interpolate_multi(
    const std::vector<quadriga_lib::arrayant<double>> &, const arma::Mat<double> *, const arma::Mat<double> *,
    const arma::Col<double> *, arma::Cube<double> *, arma::Cube<double> *, arma::Cube<double> *, arma::Cube<double> *,
    arma::uvec, const arma::Cube<double> *, const arma::Mat<double> *, bool);