1 files changed, 47 insertions, 32 deletions

diff --git a/motus/lib/entropy.ml b/motus/lib/entropy.ml
index aec5d86..fba14ce 100644
--- a/motus/lib/entropy.ml
+++ b/motus/lib/entropy.ml
@@ -45,50 +45,65 @@ let analyse : int -> catalog:Wordlist.t -> Wordlist.t -> t =
  fun base ~catalog words ->
   (* If we have only two elements, just pick one of them *)
   if Wordlist.list_size words <= 2
-  then (1.0, Option.get @@ Wordlist.pick words)
+  then (0.5, Option.get @@ Wordlist.pick words)
   else
+    let words_number = Float.of_int (Wordlist.list_size catalog) in
+
+    (* Each result from the game is stored as an integer, and we create an
+       array with as many elements as we have possibilities. *)
     let max_element =
       Float.to_int @@ (Validity.elements ** Float.of_int base)
     in
-    let words_number = Float.of_int (Wordlist.list_size catalog) in
-
     let arr =
       Bigarray.Array1.create Bigarray.Int Bigarray.C_layout max_element
     in
     match Wordlist.pick words with
     | None -> (0., "")
     | Some v ->
-        (* Build the array  *)
-        Seq.fold_left
-          (fun (score, word) word_ref ->
-            (* Reinitialize the array (we use the same in the successive
-               iterations *)
-            let set_ref = String.to_seq word_ref |> CharSet.of_seq in
+        let score, word, _ =
+          Seq.fold_left
+            (fun (score, word, already_picked) word_ref ->
+              (* Reinitialize the array (we use the same in the successive
+                 iterations) *)
+              Bigarray.Array1.fill arr 0;
 
-            Bigarray.Array1.fill arr 0;
+              let set_ref = String.to_seq word_ref |> CharSet.of_seq in
+              Seq.iter
+                (fun w2 ->
+                  let result = compare_words ~ref:(word_ref, set_ref) w2 in
+                  match result with
+                  | None -> ()
+                  | Some r ->
+                      let idx = Validity.index_of_result r in
 
-            Seq.iter
-              (fun w2 ->
-                let result = compare_words ~ref:(word_ref, set_ref) w2 in
-                match result with
-                | None -> ()
-                | Some r ->
-                    let idx = Validity.index_of_result r in
+                      let content = Bigarray.Array1.get arr idx in
+                      Bigarray.Array1.set arr idx (succ content) )
+                (Wordlist.words words);
 
-                    let content = Bigarray.Array1.get arr idx in
-                    Bigarray.Array1.set arr idx (succ content) )
-              (Wordlist.words words);
+              let entropy = get_entropy max_element words_number arr in
 
-            (* Now evaluate the entropy in the array *)
-            let entropy = get_entropy max_element words_number arr in
+              (* If get more information that we had, use this word. Choose this
+                 word to if it belongs to the final list (it gives a small chance
+                 to pick the right one…) *)
+              let is_better, already_picked =
+                if !entropy > score
+                then (true, false)
+                else if (not already_picked)
+                        && !entropy == score
+                        && Wordlist.mem word_ref words
+                then (true, true)
+                else (false, already_picked)
+              in
 
-            if !entropy > score
-            then (
-              Printf.printf
-                "Entropy for selecting %s : %.2f\n"
-                word_ref
-                !entropy;
-              (!entropy, word_ref) )
-            else (score, word) )
-          (-0., v)
-          (Wordlist.words catalog)
+              if is_better
+              then (
+                Printf.printf
+                  "Q. of information when selecting %s : %f\n"
+                  word_ref
+                  !entropy;
+                (!entropy, word_ref, already_picked) )
+              else (score, word, already_picked) )
+            (-0., v, false)
+            (Wordlist.words catalog)
+        in
+        (score, word)