(ns clojure-sandbox.prisoners)

(defn random-drawers []
  "Returns a list of shuffled numbers"
  (-> 100
      range
      shuffle))

(defn search-50-random-drawers [prisoner-number drawers]
  "Select 50 random drawers and return true if the prisoner's number was found"
  (->> drawers
      shuffle 
      (take 50) 
      (filter (fn [x] (= x prisoner-number))) 
      count
      (= 1))) 

(defn search-50-optimal-drawers [prisoner-number drawers]
  "Open 50 drawers according to the agreed strategy, returning true if prisoner's number was found"
  (loop [next-drawer prisoner-number 
         drawers-opened 0] 
    (if (= drawers-opened 50)
      false 
      (let [result (nth drawers next-drawer)] 
        (if (= result prisoner-number) 
          true 
          (recur result (inc drawers-opened))))))) 

(defn try-luck [drawers drawer-searching-function]
  "Returns 1 if all prisoners find their number otherwise 0"
  (loop [prisoners (range 100)] 
    (if (empty? prisoners) 
      1 
      (let [res (-> prisoners
                    first
                    (drawer-searching-function drawers))] 
        (if (false? res) 
          0 
          (recur (rest prisoners))))))) 

(defn simulate-100-prisoners []
  "Simulates all prisoners searching the same drawers by both strategies, returns map showing whether each was successful"
  (let [drawers (random-drawers)] 
    {:random (try-luck drawers search-50-random-drawers) 
    :optimal (try-luck drawers search-50-optimal-drawers)})) 

(defn simulate-n-runs [n]
  "Simulate n runs of the 100 prisoner problem and returns a success count for each search method"
  (loop [random-successes 0
         optimal-successes 0
         run-count 0]
    (if (= n run-count) 
      {:random-successes random-successes 
      :optimal-successes optimal-successes
      :run-count run-count}
      (let [next-result (simulate-100-prisoners)] 
        (recur (+ random-successes (:random next-result)) 
               (+ optimal-successes (:optimal next-result))
               (inc run-count)))))) 

(defn -main [& args]
  "For 5000 runs, print out the success frequency for both search methods"
  (let [{:keys [random-successes optimal-successes run-count]} (simulate-n-runs 5000)]
    (println (str "Probability of survival with random search: " (float (/ random-successes run-count))))
    (println (str "Probability of survival with ordered search: " (float (/ optimal-successes run-count))))))