(*************************************************************************)
(*  This file is part of Colibri2.                                       *)
(*                                                                       *)
(*  Copyright (C) 2014-2021                                              *)
(*    CEA   (Commissariat à l'énergie atomique et aux énergies           *)
(*           alternatives)                                               *)
(*                                                                       *)
(*  you can redistribute it and/or modify it under the terms of the GNU  *)
(*  Lesser General Public License as published by the Free Software      *)
(*  Foundation, version 2.1.                                             *)
(*                                                                       *)
(*  It 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 Lesser General Public License for more details.                  *)
(*                                                                       *)
(*  See the GNU Lesser General Public License version 2.1                *)
(*  for more details (enclosed in the file licenses/LGPLv2.1).           *)
(*************************************************************************)

open Colibri2_popop_lib
open Colibri2_core
open Colibri2_stdlib.Std

let debug =
  Debug.register_info_flag ~desc:"for the arithmetic theory of polynome"
    "LRA.polynome"

module T = struct
  include Polynome

  let name = "SARITH_POLY"
end

module ThE = ThTerm.Register (T)

module Table = struct
  module H = Datastructure.Hashtbl (T)

  let h = H.create Node.pp "HARITH_POLY"

  let set ~old_ ~new_ n d =
    (match old_ with
    | None -> ()
    | Some old_ -> if not (Polynome.equal old_ new_) then H.remove h d old_);
    H.change
      (function
        | None -> Some n
        | Some n' as s ->
            (* Should not have to register here *)
            Egraph.register d n;
            Egraph.register d n';
            Egraph.merge d n n';
            s)
      h d new_
end

include Pivot.Total (struct
  include Polynome

  let name = "LRA.polynome"

  type data = A.t

  let nodes p = p.poly
  let of_one_node n = monome A.one n
  let subst p n q = fst (subst p n q)

  let normalize p ~f =
    let add acc cl c = Polynome.x_p_cy acc c (f cl) in
    Polynome.fold add (Polynome.cst p.cst) p

  let set d cl ~old_ ~new_ =
    match is_one_node new_ with
    | None -> (
        match Polynome.is_cst new_ with
        | None -> Table.set ~old_ ~new_ cl d
        | Some q ->
            Egraph.set_value d cl (RealValue.nodevalue (RealValue.index q)))
    | Some cl' -> Egraph.merge d cl cl'

  let solve p1 p2 : _ Pivot.solve_total =
    let diff = sub p1 p2 in
    (* 0 = other - repr = p1 - p2 = diff *)
    Debug.dprintf2 debug "[Arith] @[solve 0=%a@]" pp diff;
    match Polynome.extract diff with
    | Zero -> Pivot.AlreadyEqual
    | Cst c ->
        (* 0 = cst <> 0 *)
        Debug.dprintf2 Debug.contradiction
          "[LRA/Poly] Found 0 = %a when merging" A.pp c;
        Contradiction
    | Var (q, x, p') ->
        assert (not (A.equal A.zero q));
        (* diff = qx + p' *)
        Subst (x, Polynome.mult_cst (A.div A.one (A.neg q)) p')
end)

let node_of_polynome d p =
  match Polynome.is_cst p with
  | None -> (
      match Table.H.find_opt Table.h d p with
      | Some n -> n
      | None ->
          let n = ThE.node (ThE.index p) in
          Egraph.register d n;
          assume_equality d n p;
          n)
  | Some q -> RealValue.node (RealValue.index q)

(** {2 Initialization} *)
let converter d (f : Ground.t) =
  let res = Ground.node f in
  let reg n = Egraph.register d n in
  match Ground.sem f with
  | { app = { builtin = Expr.Add _ }; tyargs = []; args; _ } ->
      let a, b = IArray.extract2_exn args in
      reg a;
      reg b;
      assume_equality d res (Polynome.of_list A.zero [ (a, A.one); (b, A.one) ])
  | { app = { builtin = Expr.Sub _ }; tyargs = []; args; _ } ->
      let a, b = IArray.extract2_exn args in
      reg a;
      reg b;
      assume_equality d res
        (Polynome.of_list A.zero [ (a, A.one); (b, A.minus_one) ])
  | { app = { builtin = Expr.Minus _ }; tyargs = []; args; _ } ->
      let a = IArray.extract1_exn args in
      reg a;
      assume_equality d res (Polynome.of_list A.zero [ (a, A.minus_one) ])
  | _ -> ()

let init env =
  Daemon.attach_reg_value env RealValue.key (fun d value ->
      let v = RealValue.value value in
      let cl = RealValue.node value in
      let p1 = Polynome.of_list v [] in
      assume_equality d cl p1);
  Ground.register_converter env converter;
  Disequality.register_disequal env (fun d n1 n2 ->
      match (get_repr d n1, get_repr d n2) with
      | Some p1, Some p2 -> (
          match Polynome.is_cst (Polynome.sub p1 p2) with
          | Some c when A.is_not_zero c -> true
          | _ -> false)
      | _ -> false)