#### 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/>.
require "../common"

module RemiAudio::DSP
  # A record that is used to hold dithering state for TPDF dithering.  This
  # structure can be used to convert samples from one type to another.
  #
  # The `Ditherer` struct also contains methods for converting sample formats
  # _without_ dithering.
  struct Ditherer
    class DitherError < RemiAudioError
    end

    @rnd : Random = Random.new(Time.local.to_unix_ms)
    @rand1 : Int32 = 0
    @rand2 : Int32 = 0
    @feedback1 : Float64 = 0.0
    @feedback2 : Float64 = 0.0

    # :nodoc:
    MAX_RAND = 2147483647i32

    # Creates a new `Ditherer` instance.
    def initialize
    end

    # Converts `sample` to a signed integer sample of `targetBitDepth` bits
    # while applying TPDF dithering when needed.  The format of `sample` must be
    # supplied via `srcFormat`.  The returned value is guaranteed to fit in
    # `targetBitDepth` bits.
    #
    # If the sample is already an integer sample, and `targetBitDepth` is
    # already greater than or equal to the current bit depth, then the sample is
    # simply coerced to the appropriate integer type without dithering.
    #
    # Unsupported bit depths will cause this to raise an `InvalidBitDepthError`.
    def dither(sample : RemiAudio::Sample, srcFormat : SampleFormat, targetBitDepth : Int,
               *, useNoiseShaping = false) : Sample
      if targetBitDepth < 8 || targetBitDepth > 64
        raise InvalidBitDepthError.new("Unsupported bit depth: #{targetBitDepth}")
      end

      # Pre-convert the sample to a float64.  Or possibly return if the sample
      # is already an integer format and the target bit depth is >= the current
      # sample bit depth.
      sampleF64 : Float64 =
        case srcFormat
        in SampleFormat::F64
          raise DitherError.new("Cannot dither a Float64 sample to an Int64 sample") unless targetBitDepth < 64
          sample.to_f64!

        in SampleFormat::F32
          raise DitherError.new("Cannot dither a Float32 sample to a higher bit depth") unless targetBitDepth < 32
          sample.to_f32!.to_f64!

        in SampleFormat::I64
          case targetBitDepth
          when 64 then return sample.to_i64!
          else int64toFloat64(sample.to_i64!)
          end

        in SampleFormat::I32
          case targetBitDepth
          when 64 then return int32toInt64(sample.to_i32!)
          when 32 then return sample.to_i32!
          else int32toFloat64(sample.to_i32!)
          end

        in SampleFormat::I24
          case targetBitDepth
          when 64 then return int24toInt64(sample.to_i32!)
          when 32 then return int24toInt32(sample.to_i32!)
          when 24 then return sample.to_i32! & 0x80FFFFFF
          else int24toFloat64(sample.to_i32! & 0x80FFFFFF)
          end

        in SampleFormat::I16
          case targetBitDepth
          when 64 then return int16toInt64(sample.to_i16!)
          when 32 then return int16toInt32(sample.to_i16!)
          when 24 then return int16toInt24(sample.to_i16!)
          when 16 then return sample.to_i16!
          else int16toFloat64(sample.to_i16!)
          end

        # in SampleFormat::I12
        #   case targetBitDepth
        #   when 64 then return int12toInt64(sample.to_i16!)
        #   when 32 then return int12toInt32(sample.to_i16!)
        #   when 24 then return int12toInt24(sample.to_i16!)
        #   when 16 then return int12toInt16(sample.to_i16!)
        #   when 12 then return sample.to_i16! & 0x8FFF
        #   else int12toFloat64(sample.to_i16!)
        #   end

        in SampleFormat::I8
          case targetBitDepth
          when 64 then return int8toInt64(sample.to_i8!)
          when 32 then return int8toInt32(sample.to_i8!)
          when 24 then return int8toInt24(sample.to_i8!)
          when 16 then return int8toInt16(sample.to_i8!)
          when 12 then return int8toInt12(sample.to_i8!)
          when 8  then return sample.to_i8!
          else raise "Unexpected bit depth"
          end

        in SampleFormat::U8
          case targetBitDepth
          when 64 then return int8toInt64((sample.to_u8!.to_i16! - 128).to_i8!)
          when 32 then return int8toInt32((sample.to_u8!.to_i16! - 128).to_i8!)
          when 24 then return int8toInt24((sample.to_u8!.to_i16! - 128).to_i8!)
          when 16 then return int8toInt16((sample.to_u8!.to_i16! - 128).to_i8!)
          when 12 then return int8toInt12((sample.to_u8!.to_i16! - 128).to_i8!)
          when 8  then return (sample.to_u8!.to_i16! - 128).to_i8!
          else raise "Unexpected bit depth"
          end
        end

      ##
      ## If we've reached here, then we're good to go to apply dithering to the
      ## sample, which is now a 64-bit float.
      ##

      s : Float64 = useNoiseShaping ? 0.0 : 0.5
      w : Float64 = (1i64 << (targetBitDepth - 1)).to_f64!
      wi : Float64 = 1.0 / w
      d : Float64 = wi / MAX_RAND
      o : Float64 = wi * 0.5
      in2 : Float64 = 0.0
      tmp : Float64 = 0.0

      @rand2 = @rand1
      @rand1 = @rnd.rand(MAX_RAND)
      in2 = sampleF64 + (s * ((@feedback1 + @feedback2) - @feedback2))
      tmp = in2 + o + (d * (@rand1 &- @rand2))
      ret = case targetBitDepth
            when 8 then  (w * tmp).to_i8! & 0xFF
            when 12 then (w * tmp).to_i16!.asBitSize(12).to_i16!
            when 16 then (w * tmp).to_i16!
            when 24 then (w * tmp).to_i32!.asBitSize(24).to_i32!
            when 32 then (w * tmp).to_i32!
            when 64 then (w * tmp).to_i64!
            else raise InvalidBitDepthError.new("Unsupported bit depth: #{targetBitDepth}")
            end
      ret = ret &- 1 if tmp < 0.0

      @feedback2 = @feedback1
      @feedback1 = in2 - (wi * ret.to_i64!)

      ret
    end

    # For each sample in `from`, this converts the sample using a `Ditherer` to
    # the destination format, then yields the converted sample and its index.
    def self.withDitheredSamples(from : RemiAudio::SampleData, srcFormat : SampleFormat, targetBitDepth : Int,
                                 *, useNoiseShaping : Bool = false, & : RemiAudio::Sample, Int32 ->)
      dith = Ditherer.new
      from.each_with_index do |samp, i|
        yield dith.dither(samp, srcFormat, targetBitDepth, useNoiseShaping: useNoiseShaping), i
      end
    end

    # Converts `sample` from `srcFormat` to `destFormat`, applying dithering
    # using `#dither` as-needed.
    #
    # Unlike the `#dither` method, this knows how to convert between float
    # formats.
    def ditherOrConvert(sample : RemiAudio::Sample, srcFormat : SampleFormat, destFormat : SampleFormat,
                        *, useNoiseShaping : Bool = false) : RemiAudio::Sample
      # Check for simple conversions that don't require dithering.
      if srcFormat == SampleFormat::F32
        # Float -> Float
        case destFormat
        when SampleFormat::F32 then return sample.to_f32!
        when SampleFormat::F64 then return sample.to_f64!
        end
      elsif srcFormat == SampleFormat::F64
        # Float -> Float
        case destFormat
        when SampleFormat::F32 then return sample.to_f32!
        when SampleFormat::F64 then return sample.to_f64!
        end
      else
        # Source is an integer format
        case destFormat
        # Integer -> Float
        when SampleFormat::F64
          case srcFormat
          when SampleFormat::U8  then return int8toFloat64((sample.to_u8!.to_i16! - 128).to_i8!)
          when SampleFormat::I8  then return int8toFloat64(sample.to_i8!)
          #when SampleFormat::I12 then return int12toFloat64(sample.to_i16!)
          when SampleFormat::I16 then return int16toFloat64(sample.to_i16!)
          when SampleFormat::I24 then return int24toFloat64(sample.to_i32!)
          when SampleFormat::I32 then return int32toFloat64(sample.to_i32!)
          when SampleFormat::I64 then return int64toFloat64(sample.to_i64!)
          end

        # Integer -> Float
        when SampleFormat::F32
          case srcFormat
          when SampleFormat::U8  then return int8toFloat64((sample.to_u8!.to_i16! - 128).to_i8!).to_f32!
          when SampleFormat::I8  then return int8toFloat64(sample.to_i8).to_f32!
          #when SampleFormat::I12 then return int12toFloat64(sample.to_i16).to_f32!
          when SampleFormat::I16 then return int16toFloat64(sample.to_i16).to_f32!
          when SampleFormat::I24 then return int24toFloat64(sample.to_i32).to_f32!
          when SampleFormat::I32 then return int32toFloat64(sample.to_i32).to_f32!
          when SampleFormat::I64 then return int64toFloat64(sample.to_i64).to_f32!
          end
        end
      end

      # If we've reached here, then we aren't doing a simple conversion and
      # instead need dithering.
      if destFormat.u8?
        # Unsigned 8-bit is handled slightly differently.  Silence is at 128,
        # not 0.
        (dither(sample, srcFormat, destFormat.getBits, useNoiseShaping: useNoiseShaping).to_i16! + 128).to_u8!
      else
        dither(sample, srcFormat, destFormat.getBits, useNoiseShaping: useNoiseShaping)
      end
    end

    # For each sample in `from`, this converts the sample using a `Ditherer` to
    # the destination format, then yields the converted sample and its index.
    # Unlike `Ditherer.withDitheredSamples`, this uses the `#ditherOrConvert`
    # method internally.
    def self.withConvertedSamples(from : SampleData, srcFormat : SampleFormat, destFormat : SampleFormat,
                                  *, useNoiseShaping : Bool = false, & : Sample, Int32 ->)
      dith = Ditherer.new
      from.each_with_index do |samp, i|
        yield dith.ditherOrConvert(samp, srcFormat, destFormat, useNoiseShaping: useNoiseShaping), i
      end
    end

    ############################################################################

    private macro defineIntUpConvFn(name, srcType, srcMask, srcBits, destType, destMask, convFn)
      private def {{name.id}}(sample : {{srcType}}) : {{destType}}
        {% den = (1 << (srcBits - 1)) %}
        {% if srcMask == nil %}
          {% if destMask == nil %}
            ((sample / {{den}}_f64) * {{destType}}::MAX).{{convFn.id}}
          {% else %}
            ((sample / {{den}}_f64) * {{destType}}::MAX).{{convFn.id}} & {{destMask}}
          {% end %}
        {% else %}
          {% if destMask == nil %}
            (((sample & {{srcMask}}) / {{den}}_f64) * {{destType}}::MAX).{{convFn.id}}
          {% else %}
            (((sample & {{srcMask}}) / {{den}}_f64) * {{destType}}::MAX).{{convFn.id}} & {{destMask}}
          {% end %}
        {% end %}
      end
    end

    defineIntUpConvFn(int8toInt64,  Int8,  nil,         8, Int64, nil, :to_i64!)
    defineIntUpConvFn(int12toInt64, Int16, 0x8FFF,     12, Int64, nil, :to_i64!)
    defineIntUpConvFn(int16toInt64, Int16, nil,        16, Int64, nil, :to_i64!)
    defineIntUpConvFn(int24toInt64, Int32, 0x80FFFFFF, 24, Int64, nil, :to_i64!)
    defineIntUpConvFn(int32toInt64, Int32, nil,        32, Int64, nil, :to_i64!)

    defineIntUpConvFn(int8toInt32,  Int8,  nil,         8, Int32, nil, :to_i32!)
    defineIntUpConvFn(int12toInt32, Int16, 0x8FFF,     12, Int32, nil, :to_i32!)
    defineIntUpConvFn(int16toInt32, Int16, nil,        16, Int32, nil, :to_i32!)
    defineIntUpConvFn(int24toInt32, Int32, 0x80FFFFFF, 24, Int32, nil, :to_i32!)

    defineIntUpConvFn(int8toInt24,  Int8,  nil,         8, Int32, 0x80FFFFFF, :to_i32!)
    defineIntUpConvFn(int12toInt24, Int16, 0x8FFF,     12, Int32, 0x80FFFFFF, :to_i32!)
    defineIntUpConvFn(int16toInt24, Int16, nil,        16, Int32, 0x80FFFFFF, :to_i32!)

    defineIntUpConvFn(int8toInt16,  Int8,  nil,         8, Int16, nil, :to_i16!)
    defineIntUpConvFn(int12toInt16, Int16, 0x8FFF,     12, Int16, nil, :to_i16!)

    defineIntUpConvFn(int8toInt12,  Int8,  nil,         8, Int16, 0x8FFF, :to_i16!)

    private macro defineIntToFloat64Fn(name, srcType)
      private def {{name.id}}(sample : {{srcType}}) : Float64
        sample / {{srcType}}::MAX.to_f64!
      end
    end

    defineIntToFloat64Fn(int8toFloat64, Int8)
    defineIntToFloat64Fn(int16toFloat64, Int16)
    defineIntToFloat64Fn(int32toFloat64, Int32)
    defineIntToFloat64Fn(int64toFloat64, Int64)

    private def int12toFloat64(sample : Int16) : Float64
      sample / 2048.0
    end

    private def int24toFloat64(sample : Int32) : Float64
      sample / 8388608.0
    end
  end
end
