#### RemiAudio
#### Copyright (C) 2022-2024 Remilia Scarlet <remilia@posteo.jp>
####
#### This program is free software: you can redistribute it and/or modify it
#### under the terms of the GNU Affero General Public License as published by
#### the Free Software Foundation, either version 3 of the License, or (at your
#### option) any later version.
####
#### This program 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 Affero General Public
#### License for more details.
####
#### You should have received a copy of the GNU Affero General Public License
#### along with this program.  If not, see <https://www.gnu.org/licenses/>.

module RemiAudio::DSP
  # A mono parametric EQ.
  #
  # The first band is always a low shelf, and the last band is always a high
  # shelf.  The others are all normal peaking bands.
  class ParaEQ
    @bands : Array(BiQuadFilter)

    # When true, this equalizer will process audio, otherwise it will pass any
    # audio through unmodified.
    property? active : Bool = true

    # The sample rate of the equalizer.
    getter sampleRate : Float64

    # This internal value is a linear value, not decibels.
    @postGain : Float64 = 1.0

    # The default cutoff frequency for low shelf, in Hertz.
    DEFAULT_LOW_SHELF_FREQ = 80.0

    # The default bandwidth for band 0, in octaves.
    DEFAULT_LOW_SHELF_WIDTH = 0.7071067811865475 # Butterworth-y, 1 / sqrt(2)

    # The default cutoff frequency for the high shelf, in Hertz.
    DEFAULT_HIGH_SHELF_FREQ = 15000.0

    # The default bandwidth for the high shelf, in octaves.
    DEFAULT_HIGH_SHELF_WIDTH = 0.7071067811865475 # Butterworth-y, 1 / sqrt(2)

    # The default cutoff frequency for the peaking bands, in Hertz.
    DEFAULT_PEAKING_FREQ = 500.0

    # The default bandwidth for the peaking bands, as a Q value.
    DEFAULT_PEAKING_Q = 1.0

    # Creates a new parametric EQ.  Bands 1 through numBands-2 are peaking
    # bands, while bands 0 and numBands-1 are the shelves.  numBands must be at
    # least 2, and includes the lowshelf and highshelf.
    def initialize(numBands, newSampleRate)
      raise RemiAudioError.new("Invalid number of parametric EQ bands") if numBands < 2
      @sampleRate = newSampleRate.to_f64!

      @bands = Array(BiQuadFilter).new(numBands) { |_| BiQuadFilter.new(@sampleRate) }

      # Low shelf
      @bands[0].setLowShelf(DEFAULT_LOW_SHELF_FREQ, 0.001, DEFAULT_LOW_SHELF_WIDTH)

      # Normal bands
      (1...(numBands - 1)).each do |i|
        @bands[i].setPeakingEQ(DEFAULT_PEAKING_FREQ, 0.001, DEFAULT_PEAKING_Q)
      end

      # High shelf
      @bands[numBands - 1].setLowShelf(DEFAULT_HIGH_SHELF_FREQ, 0.001, DEFAULT_HIGH_SHELF_WIDTH)
    end

    # The amount of gain to apply after the EQ, in decibels.
    def postGain : Float64
      RemiAudio.linearToDecibels(@postGain)
    end

    # The amount of gain to apply after the EQ, in decibels.
    def postGain=(value : Float64) : Nil
      @postGain = RemiAudio.decibelsToLinear(value)
      self.postGain
    end

    # Sets the frequency/band/width for the given band.
    #
    # If `freq` is greater than 49.9% of the sampling frequency, that band will
    # be disabled.
    def setBand(band : Int, freq : Float64, gain : Float64, width : Float64) : Nil
      # ameba:disable Style/RedundantBegin
      begin
        case band
        when 0 then @bands[band].setLowShelf(freq, gain, width)
        when @bands.size - 1 then @bands[band].setHighShelf(freq, gain, width)
        else @bands[band].setPeakingEQ(freq, gain, width)
        end
      rescue IndexError
        raise RemiAudioError.new("Invalid parametric EQ band: #{band}")
      end
    end

    # Processes a single sample of audio with the equalizer, returning a
    # processed sample.
    def process(sample : Float64) : Float64
      if @active
        @bands.each do |band|
          sample = band.process(sample)
        end
        sample *= @postGain
      end
      sample
    end

    # :ditto:
    @[AlwaysInline]
    def process(sample : Float32) : Float32
      self.process(sample.to_f64!).to_f32!
    end

    # Processes a block of audio with the equalizer.
    def process(block : Array(Float64)|Slice(Float64)) : Nil
      if @active
        @bands.each do |band|
          band.process(block) unless band.gain == 0
        end
        block.map!(&.*(@postGain))
      end
    end

    # Resets all bands in the equalizer.
    def reset : Nil
      @bands.each &.reset
    end

    # Generates a string that can be passed to a program to plot this equalizer
    # on a graph.  The `PlotType` dictates what kind of script is generated.
    def plot(pt : PlotType) : String
      case pt
      in .gnu_plot?
        String.build do |str|
          str << "# gnuplot file\n"
          str << "set title '5-band Parametric Equalizer, sample rate: #{@sampleRate}'\n"
          str << "set xlabel 'Frequency (Hz)'\n"
          str << "set ylabel 'Amplitude Response (dB)'\n"
          str << sprintf("Fs = %g\n", 44100.0)
          str << "o = 2 * pi / Fs\n"
          str << "set logscale x\n"
          str << "set samples 250\n"
          str << "set grid xtics ytics\n"
          str << "set key off\n"

          @bands.each_with_index do |band, idx|
            # Run the plot method.  We're only interested in the coefficients,
            # and we're going to adjust them.
            str << band.plot(RemiAudio::PlotType::GnuPlot, coeffsOnly: true).gsub(/([ab][012])/, "\\1_#{idx}")
            str << '\n'

            # We create a function for Gnuplot now.  Normally the #plot method does this
            # for when plotting a single filter.  But we're combining filters, so we
            # need to do it differently than #plot does.
            str << "H_#{idx}(f) = "
            str << "sqrt((b0_#{idx} * b0_#{idx} + b1_#{idx} * b1_#{idx} + b2_#{idx} * b2_#{idx} + 2.0 * " <<
              "(b0_#{idx} * b1_#{idx} + b1_#{idx} * b2_#{idx}) * cos(f * o) + 2.0 * " <<
              "(b0_#{idx} * b2_#{idx}) * cos(2.0 * f * o)) / " <<
              "(1.0 + a1_#{idx} * a1_#{idx} + a2_#{idx} * a2_#{idx} + 2.0 * (a1_#{idx} + a1_#{idx} * a2_#{idx}) * " <<
              "cos(f * o)+2.0 * a2_#{idx} * cos(2.0 * f * o)))\n"
          end


          # Now we make the main function.  Again, this isn't needed when you just plot
          # a single filter.
          str << "H(f) = "
          @bands.size.times { |x| str << "H_#{x}(f)#{x == @bands.size - 1 ? "\n" : " * "}" }

          # Finish up
          str << "plot [f=10 : Fs/2] [-35:25] 20 * log10(H(f))"
        end

      in .octave?
        raise RemiAudioError.new("Plotting equalizers to GNU Octave not yet supported")
      end
    end
  end
end
