(*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *)

(** Arithmetic terms *)

open Ses.Var_intf
include Arithmetic_intf

module Representation
    (Var : VAR)
    (Trm : INDETERMINATE with type var := Var.t) =
struct
  module Prod = struct
    include Multiset.Make
              (Int)
              (struct
                type t = Trm.trm [@@deriving compare, equal, sexp]
              end)

    let t_of_sexp = t_of_sexp Trm.trm_of_sexp
  end

  module Mono = struct
    type t = Prod.t [@@deriving compare, equal, sexp]

    let num_den m = Prod.partition m ~f:(fun _ i -> i >= 0)

    let ppx strength ppf power_product =
      let pp_factor ppf (indet, exponent) =
        if exponent = 1 then (Trm.ppx strength) ppf indet
        else Format.fprintf ppf "%a^%i" (Trm.ppx strength) indet exponent
      in
      let pp_num ppf num =
        if Prod.is_empty num then Trace.pp_styled `Magenta "1" ppf
        else Prod.pp "@ @<2>× " pp_factor ppf num
      in
      let pp_den ppf den =
        if not (Prod.is_empty den) then
          Format.fprintf ppf "@ / %a"
            (Prod.pp "@ / " pp_factor)
            (Prod.map_counts ~f:Int.neg den)
      in
      let num, den = num_den power_product in
      if Prod.is_singleton num && Prod.is_empty den then
        Format.fprintf ppf "@[<2>%a@]" pp_num num
      else Format.fprintf ppf "@[<2>(%a%a)@]" pp_num num pp_den den

    (** [one] is the empty product Πᵢ₌₁⁰ xᵢ^pᵢ *)
    let one = Prod.empty

    let equal_one = Prod.is_empty

    (** [of_ x₁ p₁] is the singleton product Πᵢ₌₁¹ x₁^p₁ *)
    let of_ x p = Prod.of_ x p

    (** [pow (Πᵢ₌₁ⁿ xᵢ^pᵢ) p] is Πᵢ₌₁ⁿ xᵢ^(pᵢ×p) *)
    let pow mono = function
      | 0 -> Prod.empty
      | 1 -> mono
      | power -> Prod.map_counts ~f:(Int.mul power) mono

    let mul x y = Prod.union x y

    (** [get_trm m] is [Some x] iff [equal m (of_ x 1)] *)
    let get_trm mono =
      match Prod.only_elt mono with Some (trm, 1) -> Some trm | _ -> None

    (* traverse *)

    let trms mono =
      Iter.from_iter (fun f -> Prod.iter mono ~f:(fun trm _ -> f trm))

    (* query *)

    let vars p = Iter.flat_map ~f:Trm.vars (trms p)
    let fv p = Var.Set.of_iter (vars p)
  end

  module Sum = struct
    include Multiset.Make (Q) (Mono)

    let t_of_sexp = t_of_sexp Mono.t_of_sexp
  end

  module Poly = struct
    type t = Sum.t [@@deriving compare, equal, sexp]
    type trm = Trm.trm
  end

  include Poly

  module Make (Embed : EMBEDDING with type trm := Trm.trm and type t := t) =
  struct
    include Poly

    let ppx strength ppf poly =
      if Sum.is_empty poly then Trace.pp_styled `Magenta "0" ppf
      else
        let pp_coeff_mono ppf (m, c) =
          if Mono.equal_one m then Trace.pp_styled `Magenta "%a" ppf Q.pp c
          else if Q.equal Q.one c then
            Format.fprintf ppf "%a" (Mono.ppx strength) m
          else Format.fprintf ppf "%a@<1>×%a" Q.pp c (Mono.ppx strength) m
        in
        if Sum.is_singleton poly then
          Format.fprintf ppf "@[<2>%a@]" (Sum.pp "@ + " pp_coeff_mono) poly
        else
          Format.fprintf ppf "@[<2>(%a)@]"
            (Sum.pp "@ + " pp_coeff_mono)
            poly

    let pp = ppx (fun _ -> None)

    let mono_invariant mono =
      let@ () = Invariant.invariant [%here] mono [%sexp_of: Mono.t] in
      Prod.iter mono ~f:(fun base power ->
          (* powers are non-zero *)
          assert (not (Int.equal Int.zero power)) ;
          match Embed.get_arith base with
          | None -> ()
          | Some poly -> (
            match Sum.classify poly with
            | `Many -> ()
            | `Zero | `One _ ->
                (* polynomial factors are not constant or singleton, which
                   should have been flattened into the parent monomial *)
                assert false ) ) ;
      match Mono.get_trm mono with
      | None -> ()
      | Some trm -> (
        match Embed.get_arith trm with
        | None -> ()
        | Some _ ->
            (* singleton monomials are not polynomials, which should have
               been flattened into the parent polynomial *)
            assert false )

    let invariant poly =
      let@ () = Invariant.invariant [%here] poly [%sexp_of: t] in
      Sum.iter poly ~f:(fun mono coeff ->
          (* coefficients are non-zero *)
          assert (not (Q.equal Q.zero coeff)) ;
          mono_invariant mono )

    (* constants *)

    let const q = Sum.of_ Mono.one q |> check invariant
    let zero = const Q.zero |> check (fun p -> assert (Sum.is_empty p))

    (* core constructors *)

    let neg poly = Sum.map_counts ~f:Q.neg poly |> check invariant
    let add p q = Sum.union p q |> check invariant
    let sub p q = add p (neg q)

    let mulc coeff poly =
      ( if Q.equal Q.one coeff then poly
      else if Q.equal Q.zero coeff then zero
      else Sum.map_counts ~f:(Q.mul coeff) poly )
      |> check invariant

    (* projections and embeddings *)

    type view = Trm of trm | Const of Q.t | Compound

    let classify poly =
      match Sum.classify poly with
      | `Zero -> Const Q.zero
      | `One (mono, coeff) -> (
        match Prod.classify mono with
        | `Zero -> Const coeff
        | `One (trm, 1) when Q.equal Q.one coeff -> Trm trm
        | _ -> Compound )
      | `Many -> Compound

    let get_const poly =
      match Sum.classify poly with
      | `Zero -> Some Q.zero
      | `One (mono, coeff) when Mono.equal_one mono -> Some coeff
      | _ -> None

    let get_mono poly =
      match Sum.only_elt poly with
      | Some (mono, coeff) when Q.equal Q.one coeff -> Some mono
      | _ -> None

    (** Terms of a polynomial: product of a coefficient and a monomial *)
    module CM = struct
      type t = Q.t * Prod.t

      let one = (Q.one, Mono.one)
      let mul (c1, m1) (c2, m2) = (Q.mul c1 c2, Mono.mul m1 m2)

      (** Monomials [Mono.t] have [trm] indeterminates, which include, via
          [get_arith], polynomials [t] over monomials themselves. To avoid
          redundant representations, singleton polynomials are flattened. *)
      let of_trm : ?power:int -> trm -> t =
       fun ?(power = 1) base ->
        match Embed.get_arith base with
        | Some poly -> (
          match Sum.classify poly with
          (* 0 ^ p₁ ==> 0 × 1 *)
          | `Zero -> (Q.zero, Mono.one)
          (* (Σᵢ₌₁¹ cᵢ × Xᵢ) ^ p₁ ==> cᵢ^p₁ × Πⱼ₌₁¹ Xⱼ^pⱼ *)
          | `One (mono, coeff) -> (Q.pow coeff power, Mono.pow mono power)
          (* (Σᵢ₌₁ⁿ cᵢ × Xᵢ) ^ p₁ ==> 1 × Πⱼ₌₁¹ (Σᵢ₌₁ⁿ cᵢ × Xᵢ)^pⱼ *)
          | `Many -> (Q.one, Mono.of_ base power) )
        (* X₁ ^ p₁ ==> 1 × Πⱼ₌₁¹ Xⱼ^pⱼ *)
        | None -> (Q.one, Mono.of_ base power)

      (** Polynomials [t] have [trm] indeterminates, which, via [get_arith],
          include polynomials themselves. To avoid redundant
          representations, singleton monomials are flattened. Also, constant
          multiplication is not interpreted in [Prod], so constant
          polynomials are multiplied by their coefficients directly. *)
      let to_poly : t -> Poly.t =
       fun (coeff, mono) ->
        ( match Mono.get_trm mono with
        | Some trm -> (
          match Embed.get_arith trm with
          (* c × (Σᵢ₌₁ⁿ cᵢ × Xᵢ) ==> Σᵢ₌₁ⁿ c×cᵢ × Xᵢ *)
          | Some poly -> mulc coeff poly
          (* c₁ × X₁ ==> Σᵢ₌₁¹ cᵢ × Xᵢ *)
          | None -> Sum.of_ mono coeff )
        (* c₁ × (Πⱼ₌₁ᵐ X₁ⱼ^p₁ⱼ) ==> Σᵢ₌₁¹ cᵢ × (Πⱼ₌₁ᵐ Xᵢⱼ^pᵢⱼ) *)
        | None -> Sum.of_ mono coeff )
        |> check invariant
    end

    (** Embed a term into a polynomial, by projecting a polynomial out of
        the term if possible *)
    let trm trm =
      ( match Embed.get_arith trm with
      | Some poly -> poly
      | None -> Sum.of_ (Mono.of_ trm 1) Q.one )
      |> check (fun poly ->
             assert (equal poly (CM.to_poly (CM.of_trm trm))) )

    (** Project out the term embedded into a polynomial, if possible *)
    let get_trm poly =
      match get_mono poly with
      | Some mono -> Mono.get_trm mono
      | None -> None

    (* constructors over indeterminates *)

    let mul e1 e2 = CM.to_poly (CM.mul (CM.of_trm e1) (CM.of_trm e2))

    let div n d =
      CM.to_poly (CM.mul (CM.of_trm n) (CM.of_trm d ~power:(-1)))

    let pow base power = CM.to_poly (CM.of_trm base ~power)

    (* transform *)

    let split_const poly =
      match Sum.find_and_remove Mono.one poly with
      | Some (c, p_c) -> (p_c, c)
      | None -> (poly, Q.zero)

    let partition_sign poly =
      Sum.partition_map poly ~f:(fun _ coeff ->
          if Q.sign coeff >= 0 then Left coeff else Right (Q.neg coeff) )

    let map poly ~f =
      [%trace]
        ~call:(fun {pf} -> pf "%a" pp poly)
        ~retn:(fun {pf} -> pf "%a" pp)
      @@ fun () ->
      let p, p' = (poly, Sum.empty) in
      let p, p' =
        Sum.fold poly (p, p') ~f:(fun mono coeff (p, p') ->
            let m, cm' = (mono, CM.one) in
            let m, cm' =
              Prod.fold mono (m, cm') ~f:(fun trm power (m, cm') ->
                  let trm' = f trm in
                  if trm == trm' then (m, cm')
                  else
                    (Prod.remove trm m, CM.mul cm' (CM.of_trm trm' ~power)) )
            in
            ( Sum.remove mono p
            , Sum.union p' (CM.to_poly (CM.mul (coeff, m) cm')) ) )
      in
      Sum.union p p' |> check invariant

    (* traverse *)

    let monos poly =
      Iter.from_iter (fun f -> Sum.iter poly ~f:(fun mono _ -> f mono))

    let trms poly = Iter.flat_map ~f:Mono.trms (monos poly)

    type product = Prod.t

    let fold_factors = Prod.fold
    let fold_monomials = Sum.fold

    (* query *)

    let vars p = Iter.flat_map ~f:Trm.vars (trms p)

    (* solve *)

    let exists_fv_in vs poly =
      Iter.exists ~f:(fun v -> Var.Set.mem v vs) (vars poly)

    (** [solve_for_mono r c m p] solves [0 = r + (c×m) + p] as [m = q]
        ([Some (m, q)]) such that [r + (c×m) + p = m - q] *)
    let solve_for_mono rejected_poly coeff mono poly =
      if Mono.equal_one mono || exists_fv_in (Mono.fv mono) poly then None
      else
        Some
          ( Sum.of_ mono Q.one
          , mulc (Q.inv (Q.neg coeff)) (Sum.union rejected_poly poly) )

    (** [solve_poly r p] solves [0 = r + p] as [m = q] ([Some (m, q)]) such
        that [r + p = m - q] *)
    let rec solve_poly rejected poly =
      [%trace]
        ~call:(fun {pf} -> pf "0 = (%a) + (%a)" pp rejected pp poly)
        ~retn:(fun {pf} s ->
          pf "%a"
            (Option.pp "%a" (fun fs (v, q) ->
                 Format.fprintf fs "%a ↦ %a" pp v pp q ))
            s )
      @@ fun () ->
      let* mono, coeff, poly = Sum.pop_min_elt poly in
      match solve_for_mono rejected coeff mono poly with
      | Some _ as soln -> soln
      | None -> solve_poly (Sum.add mono coeff rejected) poly

    (* solve [0 = e] *)
    let solve_zero_eq ?for_ e =
      [%trace]
        ~call:(fun {pf} ->
          pf "0 = %a%a" Trm.pp e (Option.pp " for %a" Trm.pp) for_ )
        ~retn:(fun {pf} s ->
          pf "%a"
            (Option.pp "%a" (fun fs (c, r) ->
                 Format.fprintf fs "%a ↦ %a" pp c pp r ))
            s ;
          match (for_, s) with
          | Some f, Some (c, _) -> assert (equal (trm f) c)
          | _ -> () )
      @@ fun () ->
      let* a = Embed.get_arith e in
      match for_ with
      | None -> solve_poly Sum.empty a
      | Some for_ -> (
          let* for_poly = Embed.get_arith for_ in
          match get_mono for_poly with
          | Some m ->
              let* c, p = Sum.find_and_remove m a in
              solve_for_mono Sum.empty c m p
          | _ -> None )
  end
end