(*
This file is part of the dpll(t) engine sammy
Copyright (C) 2004  The University of Iowa
Michael Decoster
                                                                                
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.
                                                                                
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., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
*)
                                                                                
open Basic_expression

exception No_duplicate_sorts_allowed

exception No_function_overloading

exception No_predicate_overloading

exception LogicNotFound of string
exception SortNotFound of string
exception SortIDNotFound of Global.type_id

exception Initialized_theory_reinitialized

exception Theory_change_made_after_init

exception Theory_change_after_initialization

let debug=Global.debug
let theory=Sys_state.cstore#theory

let pred_db=Sys_state.pred_db

let term_db=Term_db.term_db

class sort(name : string) =
object
  val sort_name : string = name
  val mutable sort_id : Global.type_id = -1
  method set_sort_id new_sort_id = sort_id <- new_sort_id
  method get_name = sort_name
  method get_id = sort_id
end

let int_sort = new sort("Int");;

let real_sort = new sort("Real");;

let real_real_array_sort = new sort(".Real_Real_Array");;

let int_int_array_sort = new sort(".Int_Int_Array");;

let int_real_array_sort = new sort(".Int_Real_Array");;

let real_int_array_sort = new sort(".Real_Int_Array");;

let function_type_id_table = 
    (Hashtbl.create 100:(string,Global.type_id*Global.term_id)Hashtbl.t)

let add_function_type_id = Hashtbl.replace function_type_id_table
let get_function_type_id x =
  try
    Hashtbl.find function_type_id_table x
  with Not_found -> failwith "theory_db: Function symbol unknown"

let get_function_id name = 
    let (_,id) = get_function_type_id name in
    id
let get_function_type name = 
    let (id,_) = get_function_type_id name in
    id


let theory_initialized = ref false;;

class theory_spec()=
let function_table = (Hashtbl.create 100: (string,int*sort list*sort)Hashtbl.t) in
let predicate_table = (Hashtbl.create 100: (string,int*sort list)Hashtbl.t) in
let sort_table = (Hashtbl.create 100: (string,sort)Hashtbl.t) in
let sort_id_table = (Hashtbl.create 100: (Global.type_id,sort)Hashtbl.t) in
let type_synonym_table = (Hashtbl.create 20: (string,string)Hashtbl.t) in
let function_synonym_table = (Hashtbl.create 20: (string,string)Hashtbl.t) in
let predicate_synonym_table = (Hashtbl.create 20: (string,string)Hashtbl.t) in
let numeral_sort = ref real_sort in
object(self)
  initializer
  begin
    Hashtbl.add sort_table "Int" int_sort;
    Hashtbl.add sort_table "Real" real_sort;
    Hashtbl.add sort_table ".Real_Real_Array" real_real_array_sort;
    Hashtbl.add sort_table ".Int_Int_Array" int_int_array_sort;
    Hashtbl.add sort_table ".Int_Real_Array" int_real_array_sort;
    Hashtbl.add sort_table ".Real_Int_Array" real_int_array_sort;
    Hashtbl.add function_table ".Real_Real_store" (-1,[real_real_array_sort; real_sort; real_sort],real_real_array_sort);
    Hashtbl.add function_table ".Real_Int_store" (-1,[real_int_array_sort; real_sort; int_sort],real_int_array_sort);
    Hashtbl.add function_table ".Int_Real_store" (-1,[int_real_array_sort; int_sort; real_sort],int_real_array_sort);
    Hashtbl.add function_table ".Int_Int_store" (-1,[int_int_array_sort; int_sort; int_sort],int_int_array_sort);
    Hashtbl.add function_table ".Real_Real_select" (-1,[real_real_array_sort; real_sort] ,real_sort);
    Hashtbl.add function_table ".Real_Int_select" (-1,[real_int_array_sort; real_sort],int_sort);
    Hashtbl.add function_table ".Int_Real_select" (-1,[int_real_array_sort; int_sort],real_sort);
    Hashtbl.add function_table ".Int_Int_select" (-1,[int_int_array_sort; int_sort],int_sort)
  end
  method print_sorts = 
      Printf.printf "printing all the sorts in the theory\n";
      Hashtbl.iter (fun sort_name sort -> Printf.printf "sort key=%s, sort name=%s, sort_id=%d\n" sort_name (sort#get_name) (sort#get_id) ) sort_table;
      Printf.printf "done printing all the sorts in the theory\n"
  
  method get_real_sort_name input_sort_name = 
    (try 
      Hashtbl.find type_synonym_table input_sort_name
     with Not_found ->
      input_sort_name
    )
  method get_real_function_name input_function_name =
    (try
      Hashtbl.find function_synonym_table input_function_name
     with Not_found ->
      input_function_name
    )
  method get_real_predicate_name input_predicate_name =
    (try
      Hashtbl.find predicate_synonym_table input_predicate_name
     with Not_found ->
      input_predicate_name
    )
  method add_sort_synonym input_sort_name real_sort_name =
    Hashtbl.replace type_synonym_table input_sort_name real_sort_name

  method add_function_synonym input_function_name real_function_name =
    Hashtbl.replace function_synonym_table input_function_name real_function_name

  method add_predicate_synonym input_predicate_name real_predicate_name =
    Hashtbl.replace predicate_synonym_table input_predicate_name real_predicate_name

  method initialize_theory =
    Hashtbl.iter (fun sort_name sort -> 
                      match sort_name with
                        "Int" -> 
                            sort#set_sort_id theory#intType;
                            Hashtbl.add sort_id_table theory#intType sort;
                            term_db#add_term_sort theory#intType
                      | "Real" ->
                            sort#set_sort_id theory#realType;
                            Hashtbl.add sort_id_table theory#realType sort;
                            term_db#add_term_sort theory#realType
                      | ".Real_Real_Array" ->
                            let real_real_array_sort_id= theory#decl_interpreted_type "ARRAY" [(string_of_int theory#realType); (string_of_int theory#realType)] in
                              sort#set_sort_id real_real_array_sort_id;
                              Hashtbl.add sort_id_table real_real_array_sort_id sort;
                              term_db#add_term_sort real_real_array_sort_id
                      | ".Real_Int_Array" ->
                            let real_int_array_sort_id= theory#decl_interpreted_type "ARRAY" [(string_of_int theory#realType); (string_of_int theory#intType)] in
                              sort#set_sort_id real_int_array_sort_id;
                              Hashtbl.add sort_id_table real_int_array_sort_id sort;
                              term_db#add_term_sort real_int_array_sort_id
                      | ".Int_Real_Array" ->
                            let int_real_array_sort_id= theory#decl_interpreted_type "ARRAY" [(string_of_int theory#intType); (string_of_int theory#realType)] in
                              sort#set_sort_id int_real_array_sort_id;
                              Hashtbl.add sort_id_table int_real_array_sort_id sort;
                              term_db#add_term_sort int_real_array_sort_id
                      | ".Int_Int_Array" ->
                            let int_int_array_sort_id= theory#decl_interpreted_type "ARRAY" [(string_of_int theory#intType); (string_of_int theory#intType)] in
                              sort#set_sort_id int_int_array_sort_id;
                              Hashtbl.add sort_id_table int_int_array_sort_id sort;
                              term_db#add_term_sort int_int_array_sort_id
                      | _ ->
                            let sort_id=theory#decType sort_name in
                              sort#set_sort_id sort_id;
                              Hashtbl.add sort_id_table sort_id sort;
                              term_db#add_term_sort sort_id;
                 ) sort_table;
    let defined_function_list=
        Hashtbl.fold (fun func_name (_,arg_sort_list,return_sort) acc->
                      let ret_sort_id=return_sort#get_id in
                        match arg_sort_list with
                          [] -> 
                                let fun_id = theory#decVariable func_name ret_sort_id in
                                  (func_name,(fun_id,[],return_sort))::acc
                        | _ ->
                                let sort_ids= 
                                      List.map (function sort_obj -> 
                                                      sort_obj#get_id
                                               ) arg_sort_list 
                                in
                                let dom_type_id=
                                      match sort_ids with
                                        arg_type_id::[]-> arg_type_id
                                      | _ -> theory#decl_interpreted_type "TUPLE" (List.map string_of_int sort_ids)
                                in
                                let type_id=theory#functionType dom_type_id ret_sort_id in
                                let fun_id=theory#decFunction func_name type_id in
                                  add_function_type_id func_name (type_id,fun_id);
                                  (func_name,(fun_id,arg_sort_list,return_sort))::acc
                 ) function_table []
    in
    let defined_predicate_list=
        Hashtbl.fold (fun pred_name (_,arg_sort_list) acc->
                          match arg_sort_list with
                            [] -> 
                              let pred_id = make_named_uninterpreted_pred pred_name in
                                (pred_name,(pred_id,[]))::acc
                          | _ ->
                            let sort_ids= 
                                  List.map (function sort_obj -> 
                                                  sort_obj#get_id
                                           ) arg_sort_list 
                            in
                            let dom_type_id=
                                  match sort_ids with
                                    arg_type_id::[]-> arg_type_id
                                  | _ -> theory#decl_interpreted_type "TUPLE" (List.map string_of_int sort_ids)
                            in
                            let type_id=theory#functionType dom_type_id theory#boolType in
                            let pred_id=theory#decFunction pred_name type_id in
                              add_function_type_id pred_name (type_id,pred_id);
                              (pred_name,(pred_id,arg_sort_list))::acc
                     ) predicate_table []
    in
    Hashtbl.clear function_table;
    Hashtbl.clear predicate_table;
    List.iter (fun (func_name,id_and_signature) -> Hashtbl.add function_table func_name id_and_signature ) defined_function_list;
    List.iter (fun (pred_name,id_and_signature) -> Hashtbl.add predicate_table pred_name id_and_signature ) defined_predicate_list;
    theory_initialized:= true
                    
    
  method get_numeral_sort = !numeral_sort
  
  method set_numeral_sort sort = 
      if !theory_initialized then raise Theory_change_made_after_init;
      numeral_sort := sort

  method is_sort sort_name = Hashtbl.mem sort_table sort_name
  method get_sort_ts sort_name = 
    try
      Hashtbl.find sort_table sort_name
    with Not_found -> (*self#print_sorts;*)raise (SortNotFound sort_name)
  method get_sort_by_id sort_id = 
    try
      Hashtbl.find sort_id_table sort_id (* don't use until initialized *)
    with Not_found -> raise (SortIDNotFound sort_id)
  method add_sort sort_name = 
                        if !theory_initialized then raise Theory_change_made_after_init;
                        if self#is_sort sort_name then
                          raise No_duplicate_sorts_allowed
                        else 
                          let new_sort=new sort(sort_name) in
                            Hashtbl.add sort_table sort_name new_sort
  method is_function func_name = Hashtbl.mem function_table func_name
  method get_function func_name = 
      Hashtbl.find function_table func_name
  method add_function func_name arg_sort_list return_sort=
                    if !theory_initialized then raise Theory_change_made_after_init;
                    try 
                      let (_,arg_sorts,ret_sort)= self#get_function func_name in
                      let rank_agreement= (List.length arg_sort_list) =(List.length arg_sorts) in
                      if not rank_agreement then
                        raise No_function_overloading;
                      let arg_sort_agreement=
                          List.for_all2 (fun sort1 sort2 -> 
                                            sort1=sort2
                                        ) arg_sorts arg_sort_list
                      in
                      if (not arg_sort_agreement) || (not (return_sort == ret_sort)) then
                        raise No_function_overloading
                      ()
                    with Not_found ->
                      match arg_sort_list with
                        [] -> 
                          Hashtbl.add function_table func_name 
                                      (-1,[],return_sort)
                      | _ ->
                          Hashtbl.add function_table func_name 
                                      (-1,arg_sort_list,return_sort)
  method is_predicate pred_name = Hashtbl.mem predicate_table pred_name
  method get_predicate pred_name = 
      Hashtbl.find predicate_table pred_name
  method add_predicate pred_name arg_sort_list=
                    if !theory_initialized then raise Theory_change_made_after_init;
                    try 
                      let (_,arg_sorts)= self#get_predicate pred_name in
                      let rank_agreement= (List.length arg_sort_list) =(List.length arg_sorts) in
                      if not rank_agreement then
                        raise No_predicate_overloading;
                      let arg_sort_agreement=
                          List.for_all2 (fun sort1 sort2 -> 
                                            sort1=sort2
                                        ) arg_sorts arg_sort_list
                      in
                      if not arg_sort_agreement then
                        raise No_predicate_overloading
                      ()
                    with Not_found ->
                      match arg_sort_list with
                        [] -> 
                            Hashtbl.add predicate_table pred_name (-1,[])
                      | _ ->
                          Hashtbl.add predicate_table pred_name 
                                      (-1,arg_sort_list)
end

let theory_table = (Hashtbl.create 5: (string,theory_spec)Hashtbl.t);;

let add_theory theory_name theory_spec = 
  Hashtbl.add theory_table theory_name theory_spec;;

add_theory "empty" (new theory_spec());;

let current_theory = ref (Hashtbl.find theory_table "empty")

let change_current_theory theory_name= 
  if !theory_initialized then raise Theory_change_after_initialization;
  current_theory:= 
    try
      Hashtbl.find theory_table theory_name
    with Not_found -> raise (LogicNotFound theory_name)

let initialize_current_theory() = 
    if !theory_initialized then raise Theory_change_after_initialization;
    (!current_theory)#initialize_theory
      
let extra_sort_table = (Hashtbl.create 100: (string,sort)Hashtbl.t)
let extra_sort_id_table = (Hashtbl.create 100: (Global.type_id,sort)Hashtbl.t)

let extra_functions_table = (Hashtbl.create 100: (string,int*sort list*sort)Hashtbl.t)
let extra_predicates_table = (Hashtbl.create 100: (string,int*sort list)Hashtbl.t)

let print_all_sorts() = 
      (!current_theory)#print_sorts;
      Printf.printf "printing all the extra sorts\n";
      Hashtbl.iter (fun sort_name sort -> Printf.printf "sort key=%s, sort name=%s, sort_id=%d\n" sort_name (sort#get_name) (sort#get_id) ) extra_sort_table;
      Printf.printf "done printing all the extra sorts\n"


let get_real_predicate_name input_pred_name =
      (!current_theory)#get_real_predicate_name input_pred_name

let get_real_function_name input_func_name =
      (!current_theory)#get_real_function_name input_func_name

let get_real_sort_name input_pred_name =
      (!current_theory)#get_real_predicate_name input_pred_name

let get_numeral_sort() = 
      (!current_theory)#get_numeral_sort 

let get_sort sort_name = 
      let real_sort_name=(!current_theory)#get_real_sort_name  sort_name in
      let out=(try 
        Hashtbl.find extra_sort_table real_sort_name
      with Not_found -> (!current_theory)#get_sort_ts real_sort_name
      )
      in 
      (*Printf.printf "\nget_sort:\n";flush stdout;
      Printf.printf "\nget_sort %s\n" (out#get_name);flush stdout;*)out

let get_sort_by_id sort_id = 
      let out=(try 
        Hashtbl.find extra_sort_id_table sort_id 
      with Not_found -> (!current_theory)#get_sort_by_id sort_id
      )
      in
      (*Printf.printf "\nget_sort %s\n" (out#get_name);flush stdout;*)out

let add_extra_sort sort_name = 
        let real_sort_name=(!current_theory)#get_real_sort_name  sort_name in
          (try 
            get_sort real_sort_name;
            raise No_duplicate_sorts_allowed
          with (SortNotFound _)->
            let sort_id=theory#decType real_sort_name in
            let new_sort=new sort(real_sort_name) in
              new_sort#set_sort_id sort_id;
              term_db#add_term_sort sort_id;
              Hashtbl.add extra_sort_table sort_name new_sort;
              Hashtbl.add extra_sort_id_table sort_id new_sort
          )


let add_extra_pred pred_name arg_sort_list=
    try 
      let (id,arg_sorts)= (!current_theory)#get_predicate pred_name in
      let rank_agreement= (List.length arg_sort_list) =(List.length arg_sorts) in
      if not rank_agreement then
        raise No_predicate_overloading;
      let arg_sort_agreement=
          List.for_all2 (fun sort1 sort2 -> 
                            sort1=sort2
                        ) arg_sorts arg_sort_list
      in
      if not arg_sort_agreement then
        raise No_predicate_overloading
      ()
    with Not_found ->
      match arg_sort_list with
        [] -> 
          let pred_id = make_named_uninterpreted_pred pred_name in
            Hashtbl.add extra_predicates_table pred_name (pred_id,[])

      | [arg_sort] ->
        let type_id=theory#functionType (arg_sort#get_id) theory#boolType in
        let pred_id=theory#decFunction pred_name type_id in
          add_function_type_id pred_name (type_id,pred_id);
          Hashtbl.add extra_predicates_table pred_name 
                      (pred_id,arg_sort_list)
      | _ ->
        let sort_ids= 
              List.map (function sort_obj -> 
                              sort_obj#get_id
                       ) arg_sort_list
        in
        let tuple_type_id=theory#decl_interpreted_type "TUPLE" (List.map string_of_int sort_ids) in
        let type_id=theory#functionType tuple_type_id theory#boolType in
        let pred_id=theory#decFunction pred_name type_id in
          add_function_type_id pred_name (type_id,pred_id);
          Hashtbl.add extra_predicates_table pred_name 
                      (pred_id,arg_sort_list)


let add_extra_func func_name arg_sort_list return_sort=
    try 
      let (id,arg_sorts,ret_sort)= (!current_theory)#get_function func_name in
      let rank_agreement= (List.length arg_sort_list) =(List.length arg_sorts) in
      if not rank_agreement then
        raise No_function_overloading;
      let arg_sort_agreement=
          List.for_all2 (fun sort1 sort2 -> 
                                sort1=sort2
                        ) arg_sorts arg_sort_list
      in
      if (not arg_sort_agreement) || (not (return_sort==ret_sort)) then
        raise No_function_overloading
      ()
    with Not_found ->
      let ret_sort_id=return_sort#get_id in
      match arg_sort_list with
        [] -> let fun_id = theory#decVariable func_name ret_sort_id in
          Hashtbl.add extra_functions_table func_name 
                      (fun_id,[],return_sort)
      | [arg_sort] ->
        let type_id=theory#functionType (arg_sort#get_id) ret_sort_id in
        let fun_id=theory#decFunction func_name type_id in
          add_function_type_id func_name (type_id,fun_id);
          Hashtbl.add extra_functions_table func_name 
                      (fun_id,arg_sort_list,return_sort)
      | _ ->
        let sort_ids= 
              List.map (function sort_obj -> 
                              sort_obj#get_id
                       ) arg_sort_list 
        in
        let tuple_type_id=theory#decl_interpreted_type "TUPLE" (List.map string_of_int sort_ids) in
        let type_id=theory#functionType tuple_type_id ret_sort_id in
        let fun_id=theory#decFunction func_name type_id in
          add_function_type_id func_name (type_id,fun_id);
          Hashtbl.add extra_functions_table func_name 
                      (fun_id,arg_sort_list,return_sort)

exception ExtraFunNotFound of string
exception ExtraPredNotFound of string

let get_fun_sorts func_name=
      try 
        (!current_theory)#get_function func_name
      with Not_found -> 
        begin 
          try
            Hashtbl.find extra_functions_table func_name
          with Not_found -> raise (ExtraFunNotFound func_name)
        end

let get_pred_sorts pred_name=
      try 
        (!current_theory)#get_predicate pred_name
      with Not_found -> 
        begin
          try
            Hashtbl.find extra_predicates_table pred_name
          with Not_found -> raise (ExtraPredNotFound pred_name)
        end
