#### RemiAudio
#### Copyright (C) 2022-2024 Remilia Scarlet <remilia@posteo.jp>
####   Based on code from MeltySynth
####   Copyright (C) 2021 Nobuaki Tanaka
####
#### 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/>.
require "./reverb"
require "./schroeder-presets"

module RemiAudio::DSP
  # Produces a Schroeder-style plate reverberation effect.
  #
  # This reverb implementation is based on Freeverb, a public domain reverb
  # implementation by Jezar at Dreampoint.
  class Schroeder < Reverb
    # An implementation of a comb filter.
    private class CombFilter
      @buffer : Array(Float64) = [] of Float64
      @bufferIndex : Int32 = 0
      @filterStore : Float64 = 0.0f64
      property feedback : Float64 = 0.0f64
      @damp1 : Float64 = 0.0f64
      @damp2 : Float64 = 0.0f64

      def initialize(bufferSize)
        @buffer = Array.new(bufferSize, 0.0f64)
      end

      def damping : Float64
        @damp1
      end

      def damping=(val : Float64)
        @damp1 = val
        @damp2 = 1.0f64 - val
      end

      def mute : Nil
        @buffer.fill(0.0f64)
        @filterStore = 0.0f64
      end

      def process(inputGain : Float64, inputLeft : Array(Float64), inputRight : Array(Float64),
                  outputBlock : Array(Float64))
        blockidx = 0
        len = outputBlock.size
        while blockidx < len
          @bufferIndex = 0 if @bufferIndex == @buffer.size

          srcRem = @buffer.size - @bufferIndex
          dstRem = outputBlock.size - blockidx
          rem = Math.min(srcRem, dstRem)

          rem.times do |t|
            blockpos = blockidx + t
            bufferpos = @bufferIndex + t
            input = (inputLeft[blockpos] + inputRight[blockpos]) * inputGain
            output = @buffer.unsafe_fetch(bufferpos)

            # The following ifs are to avoid performance problem due to
            # denormalized number.  The original implementation uses unsafe cast
            # to detect denormalized number.  I tried to reproduce the original
            # implementation using unsafe_as, but the simple Math.Abs version was
            # faster according to some benchmarks.
            output = 0.0f64 if output.abs < 1.0e-6

            @filterStore = (output * @damp2) + (@filterStore * @damp1)
            @filterStore = 0.0f64 if @filterStore.abs < 1.0e-6

            @buffer.unsafe_put(bufferpos, input + (@filterStore * @feedback))
            outputBlock[blockpos] += output
          end

          @bufferIndex += rem
          blockidx += rem
        end
      end
    end

    # An implementation of an all-pass filter.
    private class AllPassFilter
      @buffer = [] of Float64
      @bufferIndex : Int32 = 0
      property feedback : Float64 = 0.0f64

      def initialize(bufferSize)
        @buffer = Array.new(bufferSize, 0.0f64)
      end

      def mute
        @buffer.fill(0)
      end

      def process(block : Array(Float64))
        blockidx = 0
        len = block.size
        while blockidx < len
          @bufferIndex = 0 if @bufferIndex == @buffer.size

          src_rem = @buffer.size - @bufferIndex
          dst_rem = block.size - blockidx
          rem = Math.min(src_rem, dst_rem)

          rem.times do |t|
            blockpos = blockidx + t
            bufferpos = @bufferIndex + t
            input = block[blockpos]
            output = @buffer.unsafe_fetch(bufferpos)
            output = 0.0f64 if output.abs < 1.0e-6
            block[blockpos] = output - input
            @buffer.unsafe_put(bufferpos, input + (output * @feedback))
          end

          @bufferIndex += rem
          blockidx += rem
        end
      end
    end

    FIXED_GAIN    = 0.015f64
    SCALE_WET     = 3.0f64
    SCALE_DAMPING = 0.4f64
    SCALE_ROOM    = 0.28f64
    OFFSET_ROOM   = 0.7f64

    DEFAULT_ROOM  = 0.5f64
    ROOM_MIN      = 0.01
    ROOM_MAX      = 0.99

    DEFAULT_DAMPING  = 0.5f64
    DAMPING_MIN      = 0.01
    DAMPING_MAX      = 2.0

    DEFAULT_WIDTH = 1.0f64
    WIDTH_MIN     = 0.0
    WIDTH_MAX     = 1.5

    DEFAULT_WET   = 1.0f64 / SCALE_WET
    WET_MIN       = 0.0
    WET_MAX       = 1.0

    STEREO_SPREAD = 23
    CF_TUNING_L1  = 1116
    CF_TUNING_R1  = 1116 + STEREO_SPREAD
    CF_TUNING_L2  = 1188
    CF_TUNING_R2  = 1188 + STEREO_SPREAD
    CF_TUNING_L3  = 1277
    CF_TUNING_R3  = 1277 + STEREO_SPREAD
    CF_TUNING_L4  = 1356
    CF_TUNING_R4  = 1356 + STEREO_SPREAD
    CF_TUNING_L5  = 1422
    CF_TUNING_R5  = 1422 + STEREO_SPREAD
    CF_TUNING_L6  = 1491
    CF_TUNING_R6  = 1491 + STEREO_SPREAD
    CF_TUNING_L7  = 1557
    CF_TUNING_R7  = 1557 + STEREO_SPREAD
    CF_TUNING_L8  = 1617
    CF_TUNING_R8  = 1617 + STEREO_SPREAD
    APF_TUNING_L1 =  556
    APF_TUNING_R1 =  556 + STEREO_SPREAD
    APF_TUNING_L2 =  441
    APF_TUNING_R2 =  441 + STEREO_SPREAD
    APF_TUNING_L3 =  341
    APF_TUNING_R3 =  341 + STEREO_SPREAD
    APF_TUNING_L4 =  225
    APF_TUNING_R4 =  225 + STEREO_SPREAD

    @cfsL : Array(CombFilter) = [] of CombFilter
    @cfsR : Array(CombFilter) = [] of CombFilter
    @apfsL : Array(AllPassFilter) = [] of AllPassFilter
    @apfsR : Array(AllPassFilter) = [] of AllPassFilter

    @gain : Float64 = 0.0f64
    @roomSize : Float64 = 0.0f64
    @roomSize1 : Float64 = 0.0f64
    @damping : Float64 = 0.0f64
    @damp1 : Float64 = 0.0f64
    @wet : Float64 = 0.0f64
    @wet1 : Float64 = 0.0f64
    @wet2 : Float64 = 0.0f64

    getter width : Float64 = 0.0f64

    def initialize(sampleRate, goldenRatioFeedback : Bool = false)
      @cfsL = [
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_L1)),
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_L2)),
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_L3)),
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_L4)),
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_L5)),
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_L6)),
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_L7)),
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_L8))
      ]

      @cfsR = [
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_R1)),
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_R2)),
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_R3)),
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_R4)),
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_R5)),
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_R6)),
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_R7)),
        CombFilter.new(scaleTuning(sampleRate, CF_TUNING_R8))
      ]

      @apfsL = [
        AllPassFilter.new(scaleTuning(sampleRate, APF_TUNING_L1)),
        AllPassFilter.new(scaleTuning(sampleRate, APF_TUNING_L2)),
        AllPassFilter.new(scaleTuning(sampleRate, APF_TUNING_L3)),
        AllPassFilter.new(scaleTuning(sampleRate, APF_TUNING_L4))
      ]

      @apfsR = [
        AllPassFilter.new(scaleTuning(sampleRate, APF_TUNING_R1)),
        AllPassFilter.new(scaleTuning(sampleRate, APF_TUNING_R2)),
        AllPassFilter.new(scaleTuning(sampleRate, APF_TUNING_R3)),
        AllPassFilter.new(scaleTuning(sampleRate, APF_TUNING_R4))
      ]

      feedback = goldenRatioFeedback ? 0.6180339887498949 : 0.5
      @apfsL.each(&.feedback=(feedback))
      @apfsR.each(&.feedback=(feedback))

      self.wet = DEFAULT_WET
      self.roomSize = DEFAULT_ROOM
      self.damping = DEFAULT_DAMPING
      self.width = DEFAULT_WIDTH
    end

    @[AlwaysInline]
    protected def scaleTuning(sampleRate : Int32, tuning : Int32) : Int32
      ((sampleRate.to_f64! / 44100.0) * tuning).round.to_i32!
    end

    def process(inputLeft : Array(Float64)|Slice(Float64), inputRight : Array(Float64)|Slice(Float64),
                outputLeft : Array(Float64)|Slice(Float64), outputRight : Array(Float64)|Slice(Float64)) : Nil
      outputLeft.fill(0.0f64)
      outputRight.fill(0.0f64)

      # Process comb filters in parallel
      @cfsL.size.times do |i|
        @cfsL.unsafe_fetch(i).process(@gain, inputLeft, inputRight, outputLeft)
        @cfsR.unsafe_fetch(i).process(@gain, inputLeft, inputRight, outputRight)
      end

      # Process the all-pass filters in series
      @apfsL.each &.process(outputLeft)
      @apfsR.each &.process(outputRight)

      # Mix the reverb with the input
      inputLeft.size.times do |i|
        outL = outputLeft[i]
        outR = outputRight[i]
        outputLeft[i] = (outL * @wet1) + (outR * @wet2)
        outputRight[i] = (outR * @wet1) + (outL * @wet2)
      end
    end

    def mute : Nil
      @cfsL.each &.mute
      @cfsR.each &.mute
      @apfsL.each &.mute
      @apfsR.each &.mute
    end

    def inputGain : Float64
      @gain
    end

    def roomSize : Float64
      (@roomSize - OFFSET_ROOM) / SCALE_ROOM
    end

    def roomSize=(val : Float64)
      Schroeder.checkRoomSize(val)
      @roomSize = (val * SCALE_ROOM) + OFFSET_ROOM
      updateInternalParams
    end

    def damping : Float64
      @damping / SCALE_DAMP
    end

    def damping=(val : Float64)
      Schroeder.checkDamping(val)
      @damping = val * SCALE_DAMPING
      updateInternalParams
    end

    def wet : Float64
      @wet / SCALE_WET
    end

    def wet=(val : Float64)
      Schroeder.checkWet(val)
      @wet = val * SCALE_WET
      updateInternalParams
    end

    def width=(val : Float64)
      Schroeder.checkWidth(val)
      @width = val
      updateInternalParams
    end

    def updateInternalParams(*, numFrames : Int32 = 0)
      @wet1 = @wet * ((@width / 2.0f64) + 0.5f64)
      @wet2 = @wet * ((1.0f64 - @width) / 2.0f64)

      @roomSize1 = @roomSize
      @damp1 = @damping
      @gain = FIXED_GAIN

      @cfsL.each do |cf|
        cf.feedback = @roomSize1
        cf.damping = @damp1
      end

      @cfsR.each do |cf|
        cf.feedback = @roomSize1
        cf.damping = @damp1
      end
    end

    def usePreset(preset : Reverb::Preset) : Nil
      if data = preset.as?(Schroeder::Preset)
        self.roomSize = data.roomSize
        self.damping = data.damping
        self.width = data.width
        self.wet = data.wet
      else
        raise "Bad preset type for Schroeder: #{preset.class}"
      end
    end

    @[AlwaysInline]
    def self.checkRoomSize(value) : Nil
      unless value >= ROOM_MIN && value <= ROOM_MAX
        raise ReverbError.new("Room size for Schroeder must be between #{ROOM_MIN} and #{ROOM_MAX}")
      end
    end

    @[AlwaysInline]
    def self.checkDamping(value) : Nil
      unless value >= DAMPING_MIN && value <= DAMPING_MAX
        raise ReverbError.new("Damping for Schroeder must be between #{DAMPING_MIN} and #{DAMPING_MAX}")
      end
    end

    @[AlwaysInline]
    def self.checkWidth(value) : Nil
      unless value >= WIDTH_MIN && value <= WIDTH_MAX
        raise ReverbError.new("Width for Schroeder must be between #{WIDTH_MIN} and #{WIDTH_MAX}")
      end
    end

    @[AlwaysInline]
    def self.checkWet(value) : Nil
      unless value >= WET_MIN && value <= WET_MAX
        raise ReverbError.new("Wet for Schroeder must be between #{WET_MIN} and #{WET_MAX}")
      end
    end
  end
end
