(*****************************************************************************

  Liquidsoap, a programmable audio stream generator.
  Copyright 2003-2022 Savonet team

  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 2 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 General Public License for more details, fully stated in the COPYING
  file at the root of the liquidsoap distribution.

  You should have received a copy of the GNU General Public License
  along with this program; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301  USA

 *****************************************************************************)

open Source

type mode = Low_pass | High_pass

class filter ~kind (source : source) freq wet mode =
  let rate = float (Lazy.force Frame.audio_rate) in
  let dt = 1. /. rate in
  object (self)
    inherit operator ~name:"filter.rc" kind [source] as super
    method stype = source#stype
    method remaining = source#remaining
    method seek = source#seek
    method self_sync = source#self_sync
    method is_ready = source#is_ready
    method abort_track = source#abort_track
    val mutable prev = [||]
    val mutable prev_in = [||]

    method wake_up a =
      super#wake_up a;
      prev <- Array.make self#audio_channels 0.;
      prev_in <- Array.make self#audio_channels 0.

    method private get_frame buf =
      let offset = AFrame.position buf in
      source#get buf;
      let b = AFrame.pcm buf in
      let position = AFrame.position buf in
      let rc = 1. /. freq () in
      let alpha =
        match mode with
          | Low_pass -> dt /. (rc +. dt)
          | High_pass -> rc /. (rc +. dt)
      in
      let alpha' = 1. -. alpha in
      let wet = wet () in
      let wet' = 1. -. wet in
      match mode with
        | Low_pass ->
            let alpha = dt /. (rc +. dt) in
            for c = 0 to Array.length b - 1 do
              let b_c = b.(c) in
              for i = offset to position - 1 do
                prev.(c) <- (alpha *. b_c.(i)) +. (alpha' *. prev.(c));
                b_c.(i) <- (wet *. prev.(c)) +. (wet' *. b_c.(i))
              done
            done
        | High_pass ->
            let alpha = dt /. (rc +. dt) in
            for c = 0 to Array.length b - 1 do
              let b_c = b.(c) in
              for i = offset to position - 1 do
                prev.(c) <- alpha *. (prev.(c) +. b_c.(i) -. prev_in.(c));
                prev_in.(c) <- b_c.(i);
                b_c.(i) <- (wet *. prev.(c)) +. (wet' *. b_c.(i))
              done
            done
  end

let () =
  let kind = Lang.audio_pcm in
  let k = Lang.kind_type_of_kind_format kind in
  Lang.add_operator "filter.rc"
    [
      ("frequency", Lang.getter_t Lang.float_t, None, Some "Cutoff frequency.");
      ( "mode",
        Lang.string_t,
        None,
        Some
          "Available modes are 'low' (for low-pass filter), 'high' (for \
           high-pass filter)." );
      ( "wetness",
        Lang.getter_t Lang.float_t,
        Some (Lang.float 1.),
        Some
          "How much of the original signal should be added (1. means only \
           filtered and 0. means only original signal)." );
      ("", Lang.source_t k, None, None);
    ]
    ~return_t:k ~category:`Audio ~descr:"First-order filter (RC filter)."
    (fun p ->
      let f v = List.assoc v p in
      let freq, wet, mode, src =
        ( Lang.to_float_getter (f "frequency"),
          Lang.to_float_getter (f "wetness"),
          f "mode",
          Lang.to_source (f "") )
      in
      let mode =
        match Lang.to_string mode with
          | "low" -> Low_pass
          | "high" -> High_pass
          | _ -> raise (Error.Invalid_value (mode, "valid values are low|high"))
      in
      let kind = Source.Kind.of_kind kind in
      (new filter ~kind src freq wet mode :> Source.source))
