(* Algorithme ID3 *)

(* Encodage des données *)

(* attr0 : Meteo - valeurs : soleil 0, nuageux 1, pluie 2 *)
(* attr1 : Temperature - chaud 0, moyen 1, froid 2 *)
(* attr2 : Humidité - normal 0, elevé 1 *)
(* attr3 : Vent - non 0, oui 1 *)

(* classe : NON 0, OUI 1 *)

let nb_cl = 2;; (* nombre de classes ; de 0 à nb_cl - 1 *)

let noms_cl = [| "NON"; "OUI" |];;

let noms_attr = [| "Meteo"; "Temperature"; "Humidité"; "Vent" |];;

let vals_attr = [| [|"soleil"; "nuageux"; "pluie"|]; (* valeurs de l'attribut 0 *)
                   [|"chaud"; "moyen"; "froid"|]; (* valeurs de l'attribut 1 *)
                   [|"normal"; "élevé"|]; (* valeurs de l'attribut 2 *)
                   [|"non"; "oui"|] |];; (* valeurs de l'attribut 3 *)

let nb_val = [| 3; 3; 2; 2|];; (* nombre de valeurs par attribut ; de 0 à na - 1 *)

let data = [ ([|0; 0; 1; 0|], 0); (* les 14 exemples *)
             ([|0; 0; 1; 1|], 0);
             ([|1; 0; 1; 0|], 1);
             ([|2; 1; 1; 0|], 1);
             ([|2; 2; 0; 0|], 1);
             ([|2; 2; 0; 1|], 0);
             ([|1; 2; 0; 1|], 1);
             ([|0; 1; 1; 0|], 0);
             ([|0; 2; 0; 0|], 1);
             ([|2; 1; 0; 0|], 1);
             ([|0; 1; 0; 1|], 1);
             ([|1; 1; 1; 1|], 1);
             ([|1; 0; 0; 0|], 1);
             ([|2; 1; 1; 1|], 0) ];;

type arbre = 
    F of (int * float) (* la classe et la proportion associée *)
  | N of (int * arbre array) (* l'attribut et le tableau des sous arbres *)


let arbre_test = N (0, [| F (0, 1.); F (1, 0.7); F (-1, 1.)|]);;

let affiche_arbre arb =
  let rec affiche_espace n = 
    if n > 0 then begin print_string " "; affiche_espace (n-1) end
  in
  let rec affiche_aux decalage arb = affiche_espace decalage; match arb with
    | F (c, prob) -> Printf.printf "[Classe %d (%f)]\n" c prob
    | N (attr, fils_tab) -> begin
      Printf.printf "%s \n" noms_attr.(attr);
      Array.iter (fun a -> affiche_aux (decalage+3) a) fils_tab;
  end
  in
  affiche_aux 0 arb
;;

let affiche_joli_arbre arb =
  let rec affiche_espace n = 
    if n > 0 then begin print_string " "; affiche_espace (n-1) end
  in
  let rec affiche_aux decalage prefixe arb = affiche_espace decalage;
    print_string prefixe; 
    match arb with
    | F (c, prob) -> Printf.printf "[%s (%f)]\n" (if c = -1 then "???" else noms_cl.(c)) prob
    | N (attr, fils_tab) -> begin
      Printf.printf "%s \n" noms_attr.(attr);
      Array.iteri (fun i a -> affiche_aux (decalage+3) (vals_attr.(attr).(i)^":") a) fils_tab;
    end
  in
  affiche_aux 0 "" arb;;

affiche_arbre arbre_test;;

affiche_joli_arbre arbre_test;;

let rec decide arb obj = match arb with
| F (c, prob) -> Printf.printf "Décision %d (proba:%f)\n" c prob
| N (attr, fils_tab) -> decide fils_tab.(obj.(attr)) obj
;;

decide arbre_test [|2;0;0;1|];;

let compteurs data =
  (* renvoie le tableau des compteurs de classes *)
  let compt = Array.make nb_cl 0 in
  List.iter (fun (_,c) -> compt.(c) <- compt.(c) + 1) data;
  compt;;

let somme tab =
  (* calcule la somme d'un tableau d'entiers *)
  let t = ref 0 in
  for i = 0 to Array.length tab - 1 do
    t := !t + tab.(i)
  done; 
  !t
;;

let entropie data =
  (* calcule l'entropie d'un ensemble d'exemples *)
  let compt = compteurs data in
  let ftot = float_of_int (somme compt) in
  let e = ref 0. in
  for i = 0 to Array.length compt - 1 do
    if compt.(i) > 0 then
      let pr = float_of_int compt.(i) /. ftot in
      e := !e -. pr *. log pr /. log 2.
  done;
  !e;;

let gain data attr = 
  (* calcule le gain d'un attribut *)
  let g = ref (entropie data) in 
  let n = List.length data in
  for i = 0 to nb_val.(attr) - 1 do
    let dataa = List.filter (fun (obj,c) -> obj.(attr) = i) data in
    let na = List.length dataa in 
    if (na > 0) then
      g := !g -. (float_of_int na) /. (float_of_int n) *. entropie dataa
  done;
  !g
;;

let meilleur_attr data attr_list = 
  let rec meilleur_attr_aux attr_list = match attr_list with
  (* renvoie le nom et le gain du meilleur attribut *)
  | [] -> failwith "impossible"
  | [attr] -> (attr, gain data attr)
  | attr :: ll -> begin
    let g = gain data attr in
    let (a, gmax) = meilleur_attr_aux ll in
    if g > gmax then (attr, g) else (a, gmax)
  end
  in
  let (a,g) = meilleur_attr_aux attr_list 
  in  
  a
;;

let rec construire e data attr_list = 
  (* e : utiliser l'entropie ?
     data : l'ensemble d'entrainement
     attr : les attributs non utilisés *)
  match data with 
  | [] -> F (-1, 1.0) 
  | (_, c)::dd -> 
    if List.for_all (fun (_, cl) -> cl = c) dd then F (c, 1.0)
    else if attr_list = [] then begin
      (* choisir la classe prépondérante *)
      let compt = compteurs data in
      let tot = somme compt in
      let imax = ref (-1) in (* indice du compteur maximal *)
      for i = 0 to Array.length compt - 1 do
        if (!imax = -1) || compt.(i) > compt.(!imax) then imax := i
      done;
      F (!imax, (float_of_int compt.(!imax)) /. float_of_int tot)
    end else begin
      (* choisir le meilleur attribut *)
      let attr_choisi = 
        if e then meilleur_attr data attr_list else List.hd attr_list
      in
      let attr_restants = List.filter (fun a -> a <> attr_choisi) attr_list in
      let nfils = nb_val.(attr_choisi) in
      let fils_tab = Array.make nfils (F (0,0.)) in
      for i = 0 to nfils - 1 do
        let dataa = List.filter (fun (obj,_) -> obj.(attr_choisi) = i) data in
        fils_tab.(i) <- construire e dataa attr_restants
      done;
      N (attr_choisi, fils_tab)
    end;;


Printf.printf "Arbre1 sans entropie :\n ";
let arbre_golf_1 = construire false data [2;3;1;0] in
affiche_joli_arbre arbre_golf_1;;
    
Printf.printf "Arbre avec entropie :\n ";
let arbre_golf_avec_entropie = construire true data [0;1;2;3] in
affiche_joli_arbre arbre_golf_avec_entropie;;