open Expr
open Diff
open Eval

type condition = expr
type piece = condition * expr
type piecewise = piece list * expr

let create pieces default =
  (pieces, default)

let eval_piecewise env (pieces, default) =
  let rec find_first = function
    | [] -> eval env default
    | (cond, value) :: rest ->
        let cond_val = eval env cond in
        if cond_val <> 0.0 then eval env value
        else find_first rest
  in
  find_first pieces

let diff_piecewise var (pieces, default) =
  let diff_pieces = List.map (fun (cond, expr) ->
    (cond, diff var expr)
  ) pieces in
  (diff_pieces, diff var default)

let to_expr (_pieces, default) =
  default

let heaviside x =
  create [(Add (x, Const 0.0), Const 1.0)] (Const 0.0)

let signum x =
  create [
    (Add (x, Const 0.0), Const 1.0);
    (Sub (Const 0.0, x), Const (-1.0))
  ] (Const 0.0)

let abs_as_piecewise x =
  create [
    (Add (x, Const 0.0), x)
  ] (Neg x)

let max_expr e1 e2 =
  create [(Sub (e1, e2), e1)] e2

let min_expr e1 e2 =
  create [(Sub (e2, e1), e1)] e2

let continuous_extension expr _var _point =
  expr

let piecewise_to_string (pieces, default) =
  let piece_strs = List.map (fun (cond, value) ->
    Printf.sprintf "  %s if %s" (to_string value) (to_string cond)
  ) pieces in
  String.concat "\n" piece_strs ^ "\n  " ^ to_string default ^ " otherwise"