#include <algorithm>
#include <cstring>
#include <iostream>
#include <math.h>
#include <sstream>
#include <stdexcept>

#include <Eigen/Dense>

#include "../get_dsp.h"
#include "../registry.h"
#include "slimmable.h"
#include "model.h"

#if defined(NAM_ENABLE_A2_FAST)
  #include "a2_fast.h"
#endif

// detail::Head (WaveNet post-stack head) =====================================

nam::wavenet::detail::Head::Head(const HeadParams& params)
: _in_channels(params.in_channels)
, _out_channels(params.out_channels)
{
  if (params.kernel_sizes.empty())
    throw std::runtime_error("WaveNet Head: kernel_sizes must be non-empty");
  const size_t n = params.kernel_sizes.size();
  int cin = params.in_channels;
  for (size_t i = 0; i < n; i++)
  {
    const int cout = (i + 1 == n) ? params.out_channels : params.channels;
    const int k = params.kernel_sizes[i];
    if (k < 1)
      throw std::runtime_error("WaveNet Head: kernel_sizes entries must be >= 1");
    nam::activations::Activation::Ptr act = nam::activations::Activation::get_activation(params.activation_config);
    if (act == nullptr)
      throw std::runtime_error("WaveNet Head: unsupported activation for post-stack head");
    _activations.push_back(std::move(act));
    nam::Conv1D conv;
    conv.set_size_(cin, cout, k, true, 1, 1);
    _convs.push_back(std::move(conv));
    cin = cout;
  }
}

void nam::wavenet::detail::Head::set_weights_(std::vector<float>::iterator& weights)
{
  for (size_t i = 0; i < _convs.size(); i++)
    _convs[i].set_weights_(weights);
}

void nam::wavenet::detail::Head::SetMaxBufferSize(const int maxBufferSize)
{
  for (size_t i = 0; i < _convs.size(); i++)
    _convs[i].SetMaxBufferSize(maxBufferSize);
}

long nam::wavenet::detail::Head::receptive_field() const
{
  long rf = 1;
  for (size_t i = 0; i < _convs.size(); i++)
  {
    const long k = _convs[i].get_kernel_size();
    rf += k - 1;
  }
  return rf;
}

void nam::wavenet::detail::Head::process(Eigen::MatrixXf& work, const int num_frames)
{
  for (size_t i = 0; i < _convs.size(); i++)
  {
    const long in_ch = _convs[i].get_in_channels();
    if (i == 0)
    {
      _activations[i]->apply(work.data(), (long)(in_ch * num_frames));
      _convs[i].Process(work, num_frames);
    }
    else
    {
      auto& prev = _convs[i - 1].GetOutput();
      _activations[i]->apply(prev.data(), (long)(in_ch * num_frames));
      _convs[i].Process(prev, num_frames);
    }
  }
}

// Layer ======================================================================

void nam::wavenet::detail::Layer::SetMaxBufferSize(const int maxBufferSize)
{
  _conv.SetMaxBufferSize(maxBufferSize);
  _input_mixin.SetMaxBufferSize(maxBufferSize);
  const long z_channels = this->_conv.get_out_channels(); // This is 2*bottleneck when gated, bottleneck when not
  _z.resize(z_channels, maxBufferSize);
  if (this->_layer1x1)
  {
    this->_layer1x1->SetMaxBufferSize(maxBufferSize);
  }
  // Pre-allocate output buffers
  const long channels = this->get_channels();
  this->_output_next_layer.resize(channels, maxBufferSize);
  // _output_head stores the activated portion: bottleneck rows when no head1x1, or head1x1 out_channels when head1x1 is
  // active
  if (_head1x1)
  {
    this->_output_head.resize(_head1x1->get_out_channels(), maxBufferSize);
    this->_output_head.setZero(); // Ensure consistent initialization across platforms
    _head1x1->SetMaxBufferSize(maxBufferSize);
  }
  else
  {
    this->_output_head.resize(this->_bottleneck, maxBufferSize);
    this->_output_head.setZero(); // Ensure consistent initialization across platforms
  }
  // Set max buffer size for FiLM objects
  if (this->_conv_pre_film)
    this->_conv_pre_film->SetMaxBufferSize(maxBufferSize);
  if (this->_conv_post_film)
    this->_conv_post_film->SetMaxBufferSize(maxBufferSize);
  if (this->_input_mixin_pre_film)
    this->_input_mixin_pre_film->SetMaxBufferSize(maxBufferSize);
  if (this->_input_mixin_post_film)
    this->_input_mixin_post_film->SetMaxBufferSize(maxBufferSize);
  if (this->_activation_pre_film)
    this->_activation_pre_film->SetMaxBufferSize(maxBufferSize);
  if (this->_activation_post_film)
    this->_activation_post_film->SetMaxBufferSize(maxBufferSize);
  if (this->_layer1x1_post_film)
    this->_layer1x1_post_film->SetMaxBufferSize(maxBufferSize);
  if (this->_head1x1_post_film)
    this->_head1x1_post_film->SetMaxBufferSize(maxBufferSize);
}

void nam::wavenet::detail::Layer::set_weights_(std::vector<float>::iterator& weights)
{
  this->_conv.set_weights_(weights);
  this->_input_mixin.set_weights_(weights);
  if (this->_layer1x1)
  {
    this->_layer1x1->set_weights_(weights);
  }
  if (this->_head1x1)
  {
    this->_head1x1->set_weights_(weights);
  }
  // Set weights for FiLM objects
  if (this->_conv_pre_film)
    this->_conv_pre_film->set_weights_(weights);
  if (this->_conv_post_film)
    this->_conv_post_film->set_weights_(weights);
  if (this->_input_mixin_pre_film)
    this->_input_mixin_pre_film->set_weights_(weights);
  if (this->_input_mixin_post_film)
    this->_input_mixin_post_film->set_weights_(weights);
  if (this->_activation_pre_film)
    this->_activation_pre_film->set_weights_(weights);
  if (this->_activation_post_film)
    this->_activation_post_film->set_weights_(weights);
  if (this->_layer1x1_post_film)
    this->_layer1x1_post_film->set_weights_(weights);
  if (this->_head1x1_post_film)
    this->_head1x1_post_film->set_weights_(weights);
}

void nam::wavenet::detail::Layer::Process(const Eigen::MatrixXf& input, const Eigen::MatrixXf& condition,
                                          const int num_frames)
{
  const long bottleneck = this->_bottleneck; // Use the actual bottleneck value, not the doubled output channels

  // Step 1: input convolutions
  if (this->_conv_pre_film)
  {
    // Use Process() instead of Process_() since input is const
    this->_conv_pre_film->Process(input, condition, num_frames);
    this->_conv.Process(this->_conv_pre_film->GetOutput(), num_frames);
  }
  else
  {
    this->_conv.Process(input, num_frames);
  }
  if (this->_conv_post_film)
  {
    Eigen::MatrixXf& conv_output = this->_conv.GetOutput();
    this->_conv_post_film->Process_(conv_output, condition, num_frames);
  }

  if (this->_input_mixin_pre_film)
  {
    // Use Process() instead of Process_() since condition is const
    this->_input_mixin_pre_film->Process(condition, condition, num_frames);
    this->_input_mixin.process_(this->_input_mixin_pre_film->GetOutput(), num_frames);
  }
  else
  {
    this->_input_mixin.process_(condition, num_frames);
  }
  if (this->_input_mixin_post_film)
  {
    Eigen::MatrixXf& input_mixin_output = this->_input_mixin.GetOutput();
    this->_input_mixin_post_film->Process_(input_mixin_output, condition, num_frames);
  }
  this->_z.leftCols(num_frames).noalias() =
    _conv.GetOutput().leftCols(num_frames) + _input_mixin.GetOutput().leftCols(num_frames);

  if (this->_activation_pre_film)
  {
    this->_activation_pre_film->Process_(this->_z, condition, num_frames);
  }

  // Step 2 & 3: activation and 1x1
  //
  // A note about the gating/blending activations:
  // They take 2x dimension as input.
  // The top channels are for the "primary" activation and will be in-place modified for the final result.
  // The bottom channels are for the "secondary" activation and should not be used post-activation.
  if (this->_gating_mode == GatingMode::NONE)
  {
    this->_activation->apply(this->_z.leftCols(num_frames));
    if (this->_activation_post_film)
    {
      this->_activation_post_film->Process_(this->_z, condition, num_frames);
    }
    if (this->_layer1x1)
    {
      this->_layer1x1->process_(this->_z, num_frames);
    }
  }
  else if (this->_gating_mode == GatingMode::GATED)
  {
    // Use the GatingActivation class
    // Extract the blocks first to avoid temporary reference issues
    auto input_block = this->_z.leftCols(num_frames);
    auto output_block = this->_z.topRows(bottleneck).leftCols(num_frames);
    this->_gating_activation->apply(input_block, output_block);
    if (this->_activation_post_film)
    {
      // Use Process() for blocks and copy result back
      this->_activation_post_film->Process(this->_z.topRows(bottleneck), condition, num_frames);
      this->_z.topRows(bottleneck).leftCols(num_frames).noalias() =
        this->_activation_post_film->GetOutput().leftCols(num_frames);
    }
    if (this->_layer1x1)
    {
      this->_layer1x1->process_(this->_z.topRows(bottleneck), num_frames);
    }
  }
  else if (this->_gating_mode == GatingMode::BLENDED)
  {
    // Use the BlendingActivation class
    // Extract the blocks first to avoid temporary reference issues
    auto input_block = this->_z.leftCols(num_frames);
    auto output_block = this->_z.topRows(bottleneck).leftCols(num_frames);
    this->_blending_activation->apply(input_block, output_block);
    if (this->_activation_post_film)
    {
      // Use Process() for blocks and copy result back
      this->_activation_post_film->Process(this->_z.topRows(bottleneck), condition, num_frames);
      this->_z.topRows(bottleneck).leftCols(num_frames).noalias() =
        this->_activation_post_film->GetOutput().leftCols(num_frames);
    }
    if (this->_layer1x1)
    {
      this->_layer1x1->process_(this->_z.topRows(bottleneck), num_frames);
      if (this->_layer1x1_post_film)
      {
        Eigen::MatrixXf& layer1x1_output = this->_layer1x1->GetOutput();
        this->_layer1x1_post_film->Process_(layer1x1_output, condition, num_frames);
      }
    }
  }

  if (this->_head1x1)
  {
    if (this->_gating_mode == GatingMode::NONE)
    {
      this->_head1x1->process_(this->_z.leftCols(num_frames), num_frames);
    }
    else
    {
      this->_head1x1->process_(this->_z.topRows(bottleneck).leftCols(num_frames), num_frames);
    }
    if (this->_head1x1_post_film)
    {
      Eigen::MatrixXf& head1x1_output = this->_head1x1->GetOutput();
      this->_head1x1_post_film->Process_(head1x1_output, condition, num_frames);
    }
#ifdef NAM_USE_INLINE_GEMM
    {
      // Pure copy - use memcpy
      const int total = (int)this->_head1x1->get_out_channels() * num_frames;
      std::memcpy(this->_output_head.data(), this->_head1x1->GetOutput().data(), total * sizeof(float));
    }
#else
    this->_output_head.leftCols(num_frames).noalias() = this->_head1x1->GetOutput().leftCols(num_frames);
#endif
  }
  else // No head 1x1
  {
    // (No FiLM)
    // Store output to head (skip connection: activated conv output)
    // When _skip_head_copy is true, GetOutputHead() returns _z directly, so no copy needed.
    if (!this->_skip_head_copy)
    {
#ifdef NAM_USE_INLINE_GEMM
      if (this->_gating_mode == GatingMode::NONE)
      {
        // _z has bottleneck rows, data is contiguous - use memcpy
        const int total = (int)bottleneck * num_frames;
        std::memcpy(this->_output_head.data(), this->_z.data(), total * sizeof(float));
      }
      else
      {
        // _z has 2*bottleneck rows but we only want top bottleneck rows
        // Column-major: need to copy column by column with stride
        const int out_rows = (int)bottleneck;
        const int z_rows = (int)this->_z.rows(); // 2*bottleneck for gated
        const float* NAM_RESTRICT src = this->_z.data();
        float* NAM_RESTRICT dst = this->_output_head.data();
        for (int f = 0; f < num_frames; f++)
        {
          const float* NAM_RESTRICT src_col = src + f * z_rows;
          float* NAM_RESTRICT dst_col = dst + f * out_rows;
          for (int r = 0; r < out_rows; r++)
            dst_col[r] = src_col[r];
        }
      }
#else
      if (this->_gating_mode == GatingMode::NONE)
        this->_output_head.leftCols(num_frames).noalias() = this->_z.leftCols(num_frames);
      else
        this->_output_head.leftCols(num_frames).noalias() = this->_z.topRows(bottleneck).leftCols(num_frames);
#endif
    }
  }

  // Store output to next layer (residual connection: input + layer1x1 output, or just input if layer1x1 inactive)
  if (this->_layer1x1)
  {
#ifdef NAM_USE_INLINE_GEMM
    {
      const int channels = (int)this->get_channels();
      const int total = channels * num_frames;
      const float* NAM_RESTRICT in_ptr = input.data();
      const float* NAM_RESTRICT layer_ptr = this->_layer1x1->GetOutput().data();
      float* NAM_RESTRICT out_ptr = this->_output_next_layer.data();
      int i = 0;
      for (; i + 3 < total; i += 4)
      {
        out_ptr[i] = in_ptr[i] + layer_ptr[i];
        out_ptr[i + 1] = in_ptr[i + 1] + layer_ptr[i + 1];
        out_ptr[i + 2] = in_ptr[i + 2] + layer_ptr[i + 2];
        out_ptr[i + 3] = in_ptr[i + 3] + layer_ptr[i + 3];
      }
      for (; i < total; i++)
        out_ptr[i] = in_ptr[i] + layer_ptr[i];
    }
#else
    this->_output_next_layer.leftCols(num_frames).noalias() =
      input.leftCols(num_frames) + this->_layer1x1->GetOutput().leftCols(num_frames);
#endif
  }
  else
  {
    // If layer1x1 is inactive, residual connection is just the input (identity)
#ifdef NAM_USE_INLINE_GEMM
    {
      // Pure copy - use memcpy
      const int total = (int)this->get_channels() * num_frames;
      std::memcpy(this->_output_next_layer.data(), input.data(), total * sizeof(float));
    }
#else
    this->_output_next_layer.leftCols(num_frames).noalias() = input.leftCols(num_frames);
#endif
  }
}

// LayerArray =================================================================

nam::wavenet::detail::LayerArray::LayerArray(const LayerArrayParams& params)
: _rechannel(params.input_size, params.channels, false)
, _head_rechannel(params.head1x1_params.active ? params.head1x1_params.out_channels : params.bottleneck,
                  params.head_size, params.head_kernel_size, params.head_bias ? 1 : 0, params.head_dilation, 1)
, _head_output_size(params.head1x1_params.active ? params.head1x1_params.out_channels : params.bottleneck)
{
  const size_t num_layers = params.dilations.size();
  for (size_t i = 0; i < num_layers; i++)
  {
    const int kernel_size = params.kernel_sizes[i];
    LayerParams layer_params(
      params.condition_size, params.channels, params.bottleneck, kernel_size, params.dilations[i],
      params.activation_configs[i], params.gating_modes[i], params.groups_input, params.groups_input_mixin,
      params.layer1x1_params, params.head1x1_params, params.secondary_activation_configs[i],
      params.conv_pre_film_params, params.conv_post_film_params, params.input_mixin_pre_film_params,
      params.input_mixin_post_film_params, params.activation_pre_film_params, params.activation_post_film_params,
      params._layer1x1_post_film_params, params.head1x1_post_film_params);
    this->_layers.push_back(Layer(layer_params));
  }
}

void nam::wavenet::detail::LayerArray::SetMaxBufferSize(const int maxBufferSize)
{
  _rechannel.SetMaxBufferSize(maxBufferSize);
  _head_rechannel.SetMaxBufferSize(maxBufferSize);
  for (auto it = _layers.begin(); it != _layers.end(); ++it)
  {
    it->SetMaxBufferSize(maxBufferSize);
  }
  // Pre-allocate output buffers
  const long channels = this->_get_channels();
  this->_layer_outputs.resize(channels, maxBufferSize);
  // _head_inputs size matches actual head output: head1x1.out_channels if active, else bottleneck
  this->_head_inputs.resize(this->_head_output_size, maxBufferSize);
}


long nam::wavenet::detail::LayerArray::get_receptive_field() const
{
  long result = 0;
  for (size_t i = 0; i < this->_layers.size(); i++)
    result += this->_layers[i].get_dilation() * (this->_layers[i].get_kernel_size() - 1);
  result += this->_head_rechannel.get_dilation() * ((long)this->_head_rechannel.get_kernel_size() - 1);
  return result;
}


void nam::wavenet::detail::LayerArray::Process(const Eigen::MatrixXf& layer_inputs, const Eigen::MatrixXf& condition,
                                               const int num_frames)
{
  // Zero head inputs accumulator (first layer array)
  this->_head_inputs.setZero();
  ProcessInner(layer_inputs, condition, num_frames);
}

void nam::wavenet::detail::LayerArray::Process(const Eigen::MatrixXf& layer_inputs, const Eigen::MatrixXf& condition,
                                               const Eigen::MatrixXf& head_inputs, const int num_frames)
{
  // Copy head inputs from previous layer array - use memcpy for pure copy
#ifdef NAM_USE_INLINE_GEMM
  {
    const int total = (int)this->_head_output_size * num_frames;
    std::memcpy(this->_head_inputs.data(), head_inputs.data(), total * sizeof(float));
  }
#else
  this->_head_inputs.leftCols(num_frames).noalias() = head_inputs.leftCols(num_frames);
#endif
  ProcessInner(layer_inputs, condition, num_frames);
}

void nam::wavenet::detail::LayerArray::ProcessInner(const Eigen::MatrixXf& layer_inputs,
                                                    const Eigen::MatrixXf& condition, const int num_frames)
{
  // Process rechannel and get output
  this->_rechannel.process_(layer_inputs, num_frames);
  Eigen::MatrixXf& rechannel_output = _rechannel.GetOutput();

  // Process layers
  for (size_t i = 0; i < this->_layers.size(); i++)
  {
    // Process first layer with rechannel output, subsequent layers with previous layer output
    // Use separate branches to avoid ternary operator creating temporaries
    if (i == 0)
    {
      // First layer consumes the rechannel output buffer
      this->_layers[i].Process(rechannel_output, condition, num_frames);
    }
    else
    {
      // Subsequent layers consume the full output buffer of the previous layer
      Eigen::MatrixXf& prev_output = this->_layers[i - 1].GetOutputNextLayer();
      this->_layers[i].Process(prev_output, condition, num_frames);
    }

    // Accumulate head output from this layer
#ifdef NAM_USE_INLINE_GEMM
    {
      const int total = (int)this->_head_output_size * num_frames;
      const float* NAM_RESTRICT src = this->_layers[i].GetOutputHead().data();
      float* NAM_RESTRICT dst = this->_head_inputs.data();
      int j = 0;
      for (; j + 3 < total; j += 4)
      {
        dst[j] += src[j];
        dst[j + 1] += src[j + 1];
        dst[j + 2] += src[j + 2];
        dst[j + 3] += src[j + 3];
      }
      for (; j < total; j++)
        dst[j] += src[j];
    }
#else
    this->_head_inputs.leftCols(num_frames).noalias() += this->_layers[i].GetOutputHead().leftCols(num_frames);
#endif
  }

  // Store output from last layer - use memcpy for pure copy
  const size_t last_layer = this->_layers.size() - 1;
#ifdef NAM_USE_INLINE_GEMM
  {
    const int total = (int)this->_get_channels() * num_frames;
    std::memcpy(
      this->_layer_outputs.data(), this->_layers[last_layer].GetOutputNextLayer().data(), total * sizeof(float));
  }
#else
  this->_layer_outputs.leftCols(num_frames).noalias() =
    this->_layers[last_layer].GetOutputNextLayer().leftCols(num_frames);
#endif

  // Process head rechannel (causal Conv1D)
  _head_rechannel.Process(this->_head_inputs, num_frames);
}


Eigen::MatrixXf& nam::wavenet::detail::LayerArray::GetHeadOutputs()
{
  return this->_head_rechannel.GetOutput();
}

const Eigen::MatrixXf& nam::wavenet::detail::LayerArray::GetHeadOutputs() const
{
  return this->_head_rechannel.GetOutput();
}


void nam::wavenet::detail::LayerArray::set_weights_(std::vector<float>::iterator& weights)
{
  this->_rechannel.set_weights_(weights);
  for (size_t i = 0; i < this->_layers.size(); i++)
    this->_layers[i].set_weights_(weights);
  this->_head_rechannel.set_weights_(weights);
}

long nam::wavenet::detail::LayerArray::_get_channels() const
{
  return this->_layers.size() > 0 ? this->_layers[0].get_channels() : 0;
}

namespace
{
int wave_net_output_channels(const std::vector<nam::wavenet::LayerArrayParams>& layer_array_params,
                             const bool with_head, const std::optional<nam::wavenet::HeadParams>& head_params)
{
  if (layer_array_params.empty())
    throw std::runtime_error("WaveNet requires at least one layer array");
  if (with_head && head_params.has_value())
    return head_params->out_channels;
  return layer_array_params.back().head_size;
}
} // namespace

// WaveNet ====================================================================

nam::wavenet::WaveNet::WaveNet(const int in_channels,
                               const std::vector<nam::wavenet::LayerArrayParams>& layer_array_params,
                               const float head_scale, const bool with_head, std::optional<HeadParams> head_params,
                               std::vector<float> weights, std::unique_ptr<DSP> condition_dsp,
                               const double expected_sample_rate)
: DSP(in_channels, wave_net_output_channels(layer_array_params, with_head, head_params), expected_sample_rate)
, _condition_dsp(std::move(condition_dsp))
, _head_scale(head_scale)
{
  // Assert that if there's a condition DSP, its input is compatible with what it'll get from this WaveNet:
  if (this->_condition_dsp != nullptr)
  {
    if (this->_get_condition_dim() != this->_condition_dsp->NumInputChannels())
    {
      std::stringstream ss;
      ss << "input channels of WaveNet (" << in_channels << ") don't match input channels of condition DSP ("
         << this->_condition_dsp->NumInputChannels() << "!\n";
      throw std::runtime_error(ss.str().c_str());
    }
  }
  if (with_head)
  {
    if (!head_params.has_value())
      throw std::runtime_error("WaveNet: with_head is true but head configuration is missing");
    if (head_params->in_channels != layer_array_params.back().head_size)
    {
      std::stringstream ss;
      ss << "WaveNet head in_channels (" << head_params->in_channels << ") must match last layer array head_size ("
         << layer_array_params.back().head_size << ")";
      throw std::runtime_error(ss.str());
    }
    this->_post_stack_head = std::make_unique<detail::Head>(*head_params);
  }
  else if (head_params.has_value())
    throw std::runtime_error("WaveNet: head configuration provided but with_head is false");

  for (size_t i = 0; i < layer_array_params.size(); i++)
  {
    // Quick assert that the condition_dsp will output compatibly with this layer array
    if (this->_condition_dsp != nullptr)
    {
      if (layer_array_params[i].condition_size != this->_condition_dsp->NumOutputChannels())
      {
        std::stringstream ss;
        ss << "condition_size of layer " << i << " (" << layer_array_params[i].condition_size
           << ") doesn't match output channels of condition DSP (" << this->_condition_dsp->NumOutputChannels()
           << "!\n";
        throw std::runtime_error(ss.str().c_str());
      }
    }
    this->_layer_arrays.push_back(nam::wavenet::detail::LayerArray(layer_array_params[i]));
    if (i > 0)
      if (layer_array_params[i].channels != layer_array_params[i - 1].head_size)
      {
        std::stringstream ss;
        ss << "channels of layer " << i << " (" << layer_array_params[i].channels
           << ") doesn't match head_size of preceding layer (" << layer_array_params[i - 1].head_size << "!\n";
        throw std::runtime_error(ss.str().c_str());
      }
  }
  this->set_weights_(weights);

  // Finally, figure out how much pre-warming is needed for this model.
  mPrewarmSamples = this->_condition_dsp != nullptr ? this->_condition_dsp->GetPrewarmSamples() : 1;
  for (size_t i = 0; i < this->_layer_arrays.size(); i++)
    mPrewarmSamples += this->_layer_arrays[i].get_receptive_field();
  if (this->_post_stack_head != nullptr)
    mPrewarmSamples += this->_post_stack_head->receptive_field() - 1;
}

void nam::wavenet::WaveNet::set_weights_(std::vector<float>& weights)
{
  std::vector<float>::iterator it = weights.begin();
  // Note: condition_dsp already has its own weights from construction,
  // so we don't need to set its weights here.
  for (size_t i = 0; i < this->_layer_arrays.size(); i++)
    this->_layer_arrays[i].set_weights_(it);
  if (this->_post_stack_head != nullptr)
    this->_post_stack_head->set_weights_(it);
  this->_head_scale = *(it++); // TODO `LayerArray.absorb_head_scale()`
  if (it != weights.end())
  {
    std::stringstream ss;
    for (size_t i = 0; i < weights.size(); i++)
      if (weights[i] == *it)
      {
        ss << "Weight mismatch: assigned " << i + 1 << " weights, but " << weights.size() << " were provided.";
        throw std::runtime_error(ss.str().c_str());
      }
    ss << "Weight mismatch: provided " << weights.size() << " weights, but the model expects more.";
    throw std::runtime_error(ss.str().c_str());
  }
}

void nam::wavenet::WaveNet::SetMaxBufferSize(const int maxBufferSize)
{
  DSP::SetMaxBufferSize(maxBufferSize);
  this->_condition_input.resize(this->_get_condition_dim(), maxBufferSize);
  // Resize condition output
  if (this->_condition_dsp == nullptr)
  {
    this->_condition_output.resize(this->_get_condition_dim(), maxBufferSize);
  }
  else
  {
    this->_condition_dsp->SetMaxBufferSize(maxBufferSize);
    const int condition_output_channels = this->_condition_dsp->NumOutputChannels();
    this->_condition_output.resize(condition_output_channels, maxBufferSize);

    // Resize temporary buffers for condition DSP processing
    const int condition_dim = this->_get_condition_dim();
    this->_condition_dsp_input_buffers.resize(condition_dim);
    this->_condition_dsp_output_buffers.resize(condition_output_channels);
    this->_condition_dsp_input_ptrs.resize(condition_dim);
    this->_condition_dsp_output_ptrs.resize(condition_output_channels);

    for (int ch = 0; ch < condition_dim; ch++)
    {
      this->_condition_dsp_input_buffers[ch].resize(maxBufferSize);
      this->_condition_dsp_input_ptrs[ch] = this->_condition_dsp_input_buffers[ch].data();
    }

    for (int ch = 0; ch < condition_output_channels; ch++)
    {
      this->_condition_dsp_output_buffers[ch].resize(maxBufferSize);
      this->_condition_dsp_output_ptrs[ch] = this->_condition_dsp_output_buffers[ch].data();
    }
  }

  for (size_t i = 0; i < this->_layer_arrays.size(); i++)
    this->_layer_arrays[i].SetMaxBufferSize(maxBufferSize);

  if (this->_post_stack_head != nullptr)
  {
    this->_post_stack_head->SetMaxBufferSize(maxBufferSize);
    this->_scaled_head_scratch.resize(this->_post_stack_head->in_channels(), maxBufferSize);
  }
}

void nam::wavenet::WaveNet::SetPrewarmOnReset(const bool prewarmOnReset)
{
  DSP::SetPrewarmOnReset(prewarmOnReset);
  if (this->_condition_dsp != nullptr)
    this->_condition_dsp->SetPrewarmOnReset(prewarmOnReset);
}

void nam::wavenet::WaveNet::_process_condition(const int num_frames)
{
  if (this->_condition_dsp == nullptr)
  {
    this->_condition_output.leftCols(num_frames) = this->_condition_input.leftCols(num_frames);
  }
  else
  {
    // Copy input data from Eigen matrix to pre-allocated contiguous buffers
    // Since Eigen matrices are column-major, rows are not contiguous
    // TODO maybe use row-major here?
    const int condition_dim = this->_get_condition_dim();
    for (int ch = 0; ch < condition_dim; ch++)
    {
      for (int j = 0; j < num_frames; j++)
        this->_condition_dsp_input_buffers[ch][j] = (NAM_SAMPLE)this->_condition_input(ch, j);
    }

    // Process through condition DSP using pre-allocated buffers
    this->_condition_dsp->process(
      this->_condition_dsp_input_ptrs.data(), this->_condition_dsp_output_ptrs.data(), num_frames);

    // Copy output data back to Eigen matrix
    const int condition_output_channels = this->_condition_dsp->NumOutputChannels();
    for (int ch = 0; ch < condition_output_channels; ch++)
    {
      for (int j = 0; j < num_frames; j++)
        this->_condition_output(ch, j) = (float)this->_condition_dsp_output_buffers[ch][j];
    }
  }
}

void nam::wavenet::WaveNet::_set_condition_array(NAM_SAMPLE** input, const int num_frames)
{
  const int in_channels = NumInputChannels();
  // Fill condition array with input channels
  for (int ch = 0; ch < in_channels; ch++)
  {
    for (int j = 0; j < num_frames; j++)
    {
      this->_condition_input(ch, j) = input[ch][j];
    }
  }
}

void nam::wavenet::WaveNet::process(NAM_SAMPLE** input, NAM_SAMPLE** output, const int num_frames)
{
  assert(num_frames <= mMaxBufferSize);
  const int out_channels = NumOutputChannels();

  this->_set_condition_array(input, num_frames);
  this->_process_condition(num_frames);

  // Main layer arrays:
  // Layer-to-layer
  for (size_t i = 0; i < this->_layer_arrays.size(); i++)
  {
    if (i == 0)
    {
      // First layer array - no head input
      // layer_inputs should be the original input (before condition_dsp processing),
      // condition should be the processed condition output (after condition_dsp)
      this->_layer_arrays[i].Process(this->_condition_input, this->_condition_output, num_frames);
    }
    else
    {
      // Subsequent layer arrays - use outputs from previous layer array.
      // Pass full buffers and slice inside the callee to avoid passing Blocks
      // across API boundaries (which can cause Eigen to allocate temporaries).
      Eigen::MatrixXf& prev_layer_outputs = this->_layer_arrays[i - 1].GetLayerOutputs();
      Eigen::MatrixXf& prev_head_outputs = this->_layer_arrays[i - 1].GetHeadOutputs();
      this->_layer_arrays[i].Process(prev_layer_outputs, this->_condition_output, prev_head_outputs, num_frames);
    }
  }

  auto& final_head_outputs = this->_layer_arrays.back().GetHeadOutputs();

  if (this->_post_stack_head != nullptr)
  {
    assert(final_head_outputs.rows() == this->_post_stack_head->in_channels());
    const int head_in = this->_post_stack_head->in_channels();
    for (int ch = 0; ch < head_in; ch++)
    {
      for (int s = 0; s < num_frames; s++)
        this->_scaled_head_scratch(ch, s) = this->_head_scale * final_head_outputs(ch, s);
    }
    this->_post_stack_head->process(this->_scaled_head_scratch, num_frames);
    const Eigen::MatrixXf& head_out = this->_post_stack_head->get_last_output();
    assert(head_out.rows() == out_channels);

    if (out_channels == 1)
    {
      const float* NAM_RESTRICT src = head_out.data();
      NAM_SAMPLE* NAM_RESTRICT dst = output[0];
      for (int s = 0; s < num_frames; s++)
        dst[s] = (NAM_SAMPLE)src[s];
    }
    else
    {
      for (int ch = 0; ch < out_channels; ch++)
      {
        for (int s = 0; s < num_frames; s++)
          output[ch][s] = (NAM_SAMPLE)head_out(ch, s);
      }
    }
    return;
  }

  assert(final_head_outputs.rows() == out_channels);

  // Optimized output copy with head_scale multiplication
  if (out_channels == 1)
  {
    // Single channel: data is contiguous
    const float scale = this->_head_scale;
    const float* NAM_RESTRICT src = final_head_outputs.data();
    NAM_SAMPLE* NAM_RESTRICT dst = output[0];
    for (int s = 0; s < num_frames; s++)
    {
      dst[s] = scale * src[s];
    }
  }
  else
  {
    // Multi-channel: rows are not contiguous in column-major
    for (int ch = 0; ch < out_channels; ch++)
    {
      for (int s = 0; s < num_frames; s++)
      {
        output[ch][s] = this->_head_scale * final_head_outputs(ch, s);
      }
    }
  }
}

// Config parser - extracts all configuration from JSON without constructing the DSP
nam::wavenet::WaveNetConfig nam::wavenet::parse_config_json(const nlohmann::json& config,
                                                            const double expectedSampleRate)
{
  WaveNetConfig wc;

  // Condition DSP (eagerly built via get_dsp)
  if ((config.find("condition_dsp") != config.end()) && !config["condition_dsp"].is_null())
  {
    const nlohmann::json& condition_dsp_json = config["condition_dsp"];
    wc.condition_dsp = nam::get_dsp(condition_dsp_json);
    if (wc.condition_dsp->GetExpectedSampleRate() != expectedSampleRate)
    {
      std::stringstream ss;
      ss << "Condition DSP expected sample rate (" << wc.condition_dsp->GetExpectedSampleRate()
         << ") doesn't match WaveNet expected sample rate (" << expectedSampleRate << "!\n";
      throw std::runtime_error(ss.str().c_str());
    }
  }

  for (size_t i = 0; i < config["layers"].size(); i++)
  {
    nlohmann::json layer_config = config["layers"][i];

    const int groups = layer_config.value("groups_input", 1); // defaults to 1
    const int groups_input_mixin = layer_config.value("groups_input_mixin", 1); // defaults to 1

    const int channels = layer_config["channels"];
    const int bottleneck = layer_config.value("bottleneck", channels); // defaults to channels if not present

    // Parse layer1x1 parameters
    bool layer1x1_active = true;
    int layer1x1_groups = 1;
    if (layer_config.find("layer1x1") != layer_config.end())
    {
      const auto& layer1x1_config = layer_config["layer1x1"];
      layer1x1_active = layer1x1_config["active"];
      layer1x1_groups = layer1x1_config["groups"];
    }
    nam::wavenet::Layer1x1Params layer1x1_params(layer1x1_active, layer1x1_groups);

    const int input_size = layer_config["input_size"];
    const int condition_size = layer_config["condition_size"];

    int head_size = 0;
    int head_dilation = 1;
    int head_kernel_size = 1;
    bool head_bias = false;

    // Prefer nested "head" (matches trainer export). Legacy .nam uses head_size + head_bias (implicit kernel 1).
    if (layer_config.find("head") != layer_config.end() && !layer_config["head"].is_null())
    {
      const auto& head_json = layer_config["head"];
      if (!head_json.is_object())
      {
        throw std::runtime_error("Layer array " + std::to_string(i) + ": 'head' must be a JSON object");
      }
      head_size = head_json.at("out_channels").get<int>();

      if (head_json.contains("head_dilation"))
      {
        head_dilation = head_json.at("head_dilation").get<int>();
      }

      head_kernel_size = head_json.at("kernel_size").get<int>();
      head_bias = head_json.at("bias").get<bool>();
    }
    else if (layer_config.find("head_size") != layer_config.end())
    {
      head_size = layer_config["head_size"].get<int>();
      head_kernel_size = 1;
      head_bias = layer_config.at("head_bias").get<bool>();
    }
    else
    {
      throw std::runtime_error("Layer array " + std::to_string(i)
                               + ": expected 'head' object with out_channels, kernel_size, and bias, "
                                 "or legacy 'head_size' and 'head_bias'");
    }

    if (head_kernel_size < 1)
    {
      throw std::runtime_error("Layer array " + std::to_string(i) + ": head.kernel_size must be >= 1");
    }

    const auto dilations = layer_config["dilations"];
    const size_t num_layers = dilations.size();

    // Parse kernel sizes - support legacy single-value kernel_size or new per-layer kernel_sizes
    const bool has_kernel_size = layer_config.find("kernel_size") != layer_config.end();
    const bool has_kernel_sizes = layer_config.find("kernel_sizes") != layer_config.end();
    std::vector<int> kernel_sizes;
    if (has_kernel_size && has_kernel_sizes)
    {
      throw std::runtime_error("Layer array " + std::to_string(i)
                               + ": only one of kernel_size (int) or kernel_sizes (array) may be provided");
    }
    else if (has_kernel_sizes)
    {
      const auto& kernel_sizes_json = layer_config["kernel_sizes"];
      if (!kernel_sizes_json.is_array())
      {
        throw std::runtime_error("Layer array " + std::to_string(i) + ": kernel_sizes must be an array");
      }
      for (const auto& ks_json : kernel_sizes_json)
      {
        kernel_sizes.push_back(ks_json.get<int>());
      }
      if (kernel_sizes.size() != num_layers)
      {
        throw std::runtime_error("Layer array " + std::to_string(i) + ": kernel_sizes array size ("
                                 + std::to_string(kernel_sizes.size()) + ") must match dilations size ("
                                 + std::to_string(num_layers) + ")");
      }
    }
    else if (has_kernel_size)
    {
      const int kernel_size = layer_config["kernel_size"].get<int>();
      kernel_sizes.resize(num_layers, kernel_size);
    }
    else
    {
      throw std::runtime_error("Layer array " + std::to_string(i)
                               + ": either kernel_size (int) or kernel_sizes (array) must be provided");
    }

    // Parse activation config(s) - support both single config and array
    std::vector<activations::ActivationConfig> activation_configs;
    if (layer_config["activation"].is_array())
    {
      for (const auto& activation_json : layer_config["activation"])
      {
        activation_configs.push_back(activations::ActivationConfig::from_json(activation_json));
      }
      if (activation_configs.size() != num_layers)
      {
        throw std::runtime_error("Layer array " + std::to_string(i) + ": activation array size ("
                                 + std::to_string(activation_configs.size()) + ") must match dilations size ("
                                 + std::to_string(num_layers) + ")");
      }
    }
    else
    {
      // Single activation config - duplicate it for all layers
      const activations::ActivationConfig activation_config =
        activations::ActivationConfig::from_json(layer_config["activation"]);
      activation_configs.resize(num_layers, activation_config);
    }

    // Parse gating mode(s) - support both single value and array, and old "gated" boolean
    std::vector<GatingMode> gating_modes;
    std::vector<activations::ActivationConfig> secondary_activation_configs;

    auto parse_gating_mode_str = [](const std::string& gating_mode_str) -> GatingMode {
      if (gating_mode_str == "gated")
        return GatingMode::GATED;
      else if (gating_mode_str == "blended")
        return GatingMode::BLENDED;
      else if (gating_mode_str == "none")
        return GatingMode::NONE;
      else
        throw std::runtime_error("Invalid gating_mode: " + gating_mode_str);
    };

    if (layer_config.find("gating_mode") != layer_config.end())
    {
      if (layer_config["gating_mode"].is_array())
      {
        for (const auto& gating_mode_json : layer_config["gating_mode"])
        {
          std::string gating_mode_str = gating_mode_json.get<std::string>();
          GatingMode mode = parse_gating_mode_str(gating_mode_str);
          gating_modes.push_back(mode);

          // Parse corresponding secondary activation if gating is enabled
          if (mode != GatingMode::NONE)
          {
            if (layer_config.find("secondary_activation") != layer_config.end())
            {
              if (layer_config["secondary_activation"].is_array())
              {
                if (gating_modes.size() > layer_config["secondary_activation"].size())
                {
                  throw std::runtime_error("Layer array " + std::to_string(i)
                                           + ": secondary_activation array size must be at least "
                                           + std::to_string(gating_modes.size()));
                }
                secondary_activation_configs.push_back(activations::ActivationConfig::from_json(
                  layer_config["secondary_activation"][gating_modes.size() - 1]));
              }
              else
              {
                // Single secondary activation - use for all gated layers
                secondary_activation_configs.push_back(
                  activations::ActivationConfig::from_json(layer_config["secondary_activation"]));
              }
            }
            else
            {
              // Default to Sigmoid for backward compatibility
              secondary_activation_configs.push_back(
                activations::ActivationConfig::simple(activations::ActivationType::Sigmoid));
            }
          }
          else
          {
            secondary_activation_configs.push_back(activations::ActivationConfig{});
          }
        }
        if (gating_modes.size() != num_layers)
        {
          throw std::runtime_error("Layer array " + std::to_string(i) + ": gating_mode array size ("
                                   + std::to_string(gating_modes.size()) + ") must match dilations size ("
                                   + std::to_string(num_layers) + ")");
        }
        // Validate secondary_activation array size if it's an array
        if (layer_config.find("secondary_activation") != layer_config.end()
            && layer_config["secondary_activation"].is_array())
        {
          if (layer_config["secondary_activation"].size() != num_layers)
          {
            throw std::runtime_error("Layer array " + std::to_string(i) + ": secondary_activation array size ("
                                     + std::to_string(layer_config["secondary_activation"].size())
                                     + ") must match dilations size (" + std::to_string(num_layers) + ")");
          }
        }
      }
      else
      {
        // Single gating mode - duplicate for all layers
        std::string gating_mode_str = layer_config["gating_mode"].get<std::string>();
        GatingMode gating_mode = parse_gating_mode_str(gating_mode_str);
        gating_modes.resize(num_layers, gating_mode);

        activations::ActivationConfig secondary_activation_config;
        if (gating_mode != GatingMode::NONE)
        {
          if (layer_config.find("secondary_activation") != layer_config.end())
          {
            secondary_activation_config =
              activations::ActivationConfig::from_json(layer_config["secondary_activation"]);
          }
          else
          {
            // Default to Sigmoid for backward compatibility
            secondary_activation_config = activations::ActivationConfig::simple(activations::ActivationType::Sigmoid);
          }
        }
        secondary_activation_configs.resize(num_layers, secondary_activation_config);
      }
    }
    // Backward compatibility: convert old "gated" boolean to new enum
    else if (layer_config.find("gated") != layer_config.end())
    {
      bool gated = layer_config["gated"];
      GatingMode gating_mode = gated ? GatingMode::GATED : GatingMode::NONE;
      gating_modes.resize(num_layers, gating_mode);

      if (gated)
      {
        activations::ActivationConfig secondary_config =
          activations::ActivationConfig::simple(activations::ActivationType::Sigmoid);
        secondary_activation_configs.resize(num_layers, secondary_config);
      }
      else
      {
        secondary_activation_configs.resize(num_layers, activations::ActivationConfig{});
      }
    }
    else
    {
      // Default to NONE for all layers
      gating_modes.resize(num_layers, GatingMode::NONE);
      secondary_activation_configs.resize(num_layers, activations::ActivationConfig{});
    }

    // Parse head1x1 parameters
    bool head1x1_active = false;
    int head1x1_out_channels = channels;
    int head1x1_groups = 1;
    if (layer_config.find("head1x1") != layer_config.end())
    {
      const auto& head1x1_config = layer_config["head1x1"];
      head1x1_active = head1x1_config["active"];
      head1x1_out_channels = head1x1_config["out_channels"];
      head1x1_groups = head1x1_config["groups"];
    }
    nam::wavenet::Head1x1Params head1x1_params(head1x1_active, head1x1_out_channels, head1x1_groups);

    // Helper function to parse FiLM parameters
    auto parse_film_params = [&layer_config](const std::string& key) -> nam::wavenet::_FiLMParams {
      if (layer_config.find(key) == layer_config.end() || layer_config[key] == false)
      {
        return nam::wavenet::_FiLMParams(false, false);
      }
      const nlohmann::json& film_config = layer_config[key];
      bool active = film_config.value("active", true);
      bool shift = film_config.value("shift", true);
      int film_groups = film_config.value("groups", 1);
      return nam::wavenet::_FiLMParams(active, shift, film_groups);
    };

    // Parse FiLM parameters
    nam::wavenet::_FiLMParams conv_pre_film_params = parse_film_params("conv_pre_film");
    nam::wavenet::_FiLMParams conv_post_film_params = parse_film_params("conv_post_film");
    nam::wavenet::_FiLMParams input_mixin_pre_film_params = parse_film_params("input_mixin_pre_film");
    nam::wavenet::_FiLMParams input_mixin_post_film_params = parse_film_params("input_mixin_post_film");
    nam::wavenet::_FiLMParams activation_pre_film_params = parse_film_params("activation_pre_film");
    nam::wavenet::_FiLMParams activation_post_film_params = parse_film_params("activation_post_film");
    nam::wavenet::_FiLMParams _layer1x1_post_film_params = parse_film_params("layer1x1_post_film");
    nam::wavenet::_FiLMParams head1x1_post_film_params = parse_film_params("head1x1_post_film");

    // Validation: if layer1x1_post_film is active, layer1x1 must also be active
    if (_layer1x1_post_film_params.active && !layer1x1_active)
    {
      throw std::runtime_error("Layer array " + std::to_string(i)
                               + ": layer1x1_post_film cannot be active when layer1x1.active is false");
    }

    wc.layer_array_params.push_back(nam::wavenet::LayerArrayParams(
      input_size, condition_size, head_size, head_dilation, head_kernel_size, channels, bottleneck, std::move(kernel_sizes), dilations,
      std::move(activation_configs), std::move(gating_modes), head_bias, groups, groups_input_mixin, layer1x1_params,
      head1x1_params, std::move(secondary_activation_configs), conv_pre_film_params, conv_post_film_params,
      input_mixin_pre_film_params, input_mixin_post_film_params, activation_pre_film_params,
      activation_post_film_params, _layer1x1_post_film_params, head1x1_post_film_params));
  }

  wc.with_head = config.find("head") != config.end() && !config["head"].is_null();
  wc.head_scale = config["head_scale"];
  wc.in_channels = config.value("in_channels", 1);

  if (wc.layer_array_params.empty())
    throw std::runtime_error("WaveNet config requires at least one layer array");

  if (wc.with_head)
  {
    const nlohmann::json& hj = config["head"];
    HeadParams hp;
    const int implied_in = wc.layer_array_params.back().head_size;
    // New trainer export omits in_channels (single source: last layer head_size). Legacy .nam may include it.
    if (hj.find("in_channels") != hj.end() && !hj["in_channels"].is_null())
    {
      const int legacy_in = hj["in_channels"].get<int>();
      if (legacy_in != implied_in)
      {
        std::stringstream ss;
        ss << "WaveNet config: head.in_channels (" << legacy_in << ") must equal last layer's head_size (" << implied_in
           << ")";
        throw std::runtime_error(ss.str());
      }
    }
    hp.in_channels = implied_in;
    hp.channels = hj.at("channels").get<int>();
    hp.out_channels = hj.at("out_channels").get<int>();
    hp.kernel_sizes = hj.at("kernel_sizes").get<std::vector<int>>();
    hp.activation_config = nam::activations::ActivationConfig::from_json(hj.at("activation"));
    if (hp.kernel_sizes.empty())
      throw std::runtime_error("WaveNet config: head.kernel_sizes must be non-empty");
    wc.head_params = std::move(hp);
  }
  else
    wc.head_params = std::nullopt;

  return wc;
}

// WaveNetConfig::create()
std::unique_ptr<nam::DSP> nam::wavenet::WaveNetConfig::create(std::vector<float> weights, double sampleRate)
{
  return std::make_unique<nam::wavenet::WaveNet>(in_channels, layer_array_params, head_scale, with_head,
                                                 std::move(head_params), std::move(weights), std::move(condition_dsp),
                                                 sampleRate);
}

namespace
{
const std::string SLIMMABLE_METHOD = "slice_channels_uniform";

bool config_is_slimmable_wavenet(const nlohmann::json& config)
{
  if (config.find("layers") == config.end() || !config["layers"].is_array())
    return false;
  for (const auto& lc : config["layers"])
  {
    if (lc.find("slimmable") == lc.end() || !lc["slimmable"].is_object())
      continue;
    const std::string method = lc["slimmable"].value("method", "");
    if (method != SLIMMABLE_METHOD)
    {
      if (!method.empty())
        throw std::runtime_error("SlimmableWavenet: unsupported slimmable method '" + method + "'");
      continue;
    }
    return true;
  }
  return false;
}
} // namespace

// Config parser for ConfigParserRegistry
std::unique_ptr<nam::ModelConfig> nam::wavenet::create_config(const nlohmann::json& config, double sampleRate)
{
  if (config_is_slimmable_wavenet(config))
    return nam::slimmable_wavenet::create_config(config, sampleRate);

#if defined(NAM_ENABLE_A2_FAST)
  if (int a2_channels = 0; nam::wavenet::a2_fast::is_a2_shape(config, &a2_channels))
    return nam::wavenet::a2_fast::create_a2_fast_config(config, sampleRate);
#endif

  auto wc = std::make_unique<WaveNetConfig>();
  auto parsed = parse_config_json(config, sampleRate);
  *wc = std::move(parsed);
  return wc;
}

// Register the config parser
namespace
{
static nam::ConfigParserHelper _register_WaveNet("WaveNet", nam::wavenet::create_config);
}
