open ExtLib;; 
let (|>) x f = f x;; 
let (>>) f g x = g (f x);;  
let identity x = x;;

type position = Left | Right | Out;; (* при взвешивании монета или на левой чашке весов, или на правой, или вне весов *)
type outcome = LeftLighter | LeftHeavier | Equal;; (* результат взвешивания *)
type weighting = position list;; (* одно взвешивание - список позиций монет *)
let cNormal = (false,true,false);;
let cAny = (true, true, true);;

(* ментальная модель набора монет, оно же состояние st, - список троек булевых значений, 
обозначающих может ли данная монета быть легкой, нормальной или тяжелой (light, normal, heavy) *)

let sum (cl,cn,ch) (l,n,h) = (cl || l, cn && n, ch || h);;
let summarize st w = 
  List.fold_left2 (fun (l,r) c p -> match p with Left -> sum l c, r | Right -> l, sum r c | Out -> l,r) 
    (cNormal, cNormal) st w;;
  
let weigh st w = (* определение результата взвешивания *)
  let (ll, ln, lh), (rl, rn, rh) = summarize st w in 
  let lst1 = if ln && rn then [Equal] else [] in
  let lst2 = if ll || rh then LeftLighter::lst1 else lst1 in
  if lh || rl then LeftHeavier::lst2 else lst2;;

let not_heavy (l,n,h) = (l,n,false);;
let not_light (l,n,h) = (false,n,h);;
let is_normal (l,n,h) = (false,true,false);;

let eval st w outcome = (* определение нового состояния из заданных состояния, взвешивания и результата *)
  let f = match outcome with
    | LeftLighter -> (function Left -> not_heavy | Right -> not_light | Out -> is_normal)
    | LeftHeavier -> (function Left -> not_light | Right -> not_heavy | Out -> is_normal)
    | Equal       -> (function Left  | Right -> is_normal | Out -> identity) in
  List.map2 f w st;;
      
let weightings f m k = (* перебирает все различные взвешивания из m монет, когда на каждой чашке k монет, и передает их в f *)
  let rec loop lst m kl kr =
    if m < 1 then (if kl = 0 && kr = 0 then f lst else ()) else
    (if kl > 0 then loop (Left::lst) (m-1) (kl-1) kr;
    if kr > 0 && kl > 0 then loop (Right::lst) (m-1) kl (kr-1);
    if kl+kr < m then loop (Out::lst) (m-1) kl kr) in
  loop [] m k k;;        
        
let good_state st = (* состояние хорошее, если все монеты кроме одной точно нормальные *)
  1 = List.fold_left (fun k c -> if c <> cNormal then k+1 else k) 0 st;;
        
let find_fake st = (* найти номер поддельной монеты *)
  List.findi (fun i c -> c <> cNormal) st |> fst;;
        
let repeat x k = (* создание списка из k элементов х *)
  let rec loop lst = function 0 -> lst | n -> loop (x::lst) (n-1) in loop [] k;;

(* описание последнего шага алгоритма - по результату взешивания говорит номер фальшивой монеты*)
type last_step = weighting * (outcome * int) list;; 
(* шаг алгоритма - взвешивание и действия для каждого возможного результата *)
type algorithm = LastStep of last_step | Step of weighting * (outcome * algorithm) list | Found of int;; 
exception AlgEx of algorithm;;

let rec check_weighting m nw st w = (* поиск решения для данного взвешивания *)
  let outcomes = weigh st w in
  if List.length outcomes < 2 then () else (* если у взв. детерминированный результат, оно неинформативно и неинтересно *)
  let newstates = List.map (eval st w) outcomes in
  if List.for_all good_state newstates then (* если при любом результате знаем где фальшивка, найден последний шаг *)
    raise (AlgEx(LastStep(w, List.map2 (fun oc s-> oc, find_fake s) outcomes newstates)))
  else (* иначе ищем решения для возможных результатов данного взвешивания *)
    if nw < 2 then () else check_states m nw w outcomes newstates

and solvable m nw st = (* поиск решений для данного состояния st, числа оставшихся взвешиваний nw и числа монет m *)
  if good_state st then Some(Found(find_fake st)) else (* если точно знаем где фальшивка, дело сделано *)
  try for k=m/2 downto 1 do (* ищем по всем вариантам взвешивания *)
        weightings (check_weighting m nw st) m k (* если решение найдется, цикл прервется исключением *)
      done;
      None (*  решение не найдено  *)
  with AlgEx a -> Some a (* решение найдено *)
  
and check_states m nw w outcomes newstates = (* поиск решений для данного набора состояний *)
  let svb = List.map (fun s-> lazy (solvable m (nw-1) s )) newstates in  (* ленивый поиск решений для каждого состояния *)
  if List.for_all (fun lo -> Lazy.force lo <> None) svb then (* если для каждого результата взвешивания знаем как найти фальшивку, то *)
    let algs = List.map (Lazy.force >> Option.get) svb in (* получить список алгоритмов по поиску фальшивки в каждом случае *)
    let a = Step(w, List.combine outcomes algs) in
    raise (AlgEx a);; (* выдать ответ с шагом алгоритма поиска *)

let solve st0 nw = (* решить задачу для исходного состояния монет st0 и числа взвешиваний nw *)
  let m = List.length st0 in  
  try
    for k=m/2 downto 1 do
      let w1 = (repeat Left k) @ (repeat Right k) @ (repeat Out (m-k*2)) in (* для каждого варианта первого взвешивания *)
      let outcomes1 = weigh st0 w1 in (* получаем результаты взешивания *)
      let states1 = List.map (eval st0 w1) outcomes1 in (* и соответствующие им состояния *)
      check_states m nw w1 outcomes1 states1 (* ищем решение при таком взвешивании *)
    done;
    None (* ничего не нашли *)
  with AlgEx alg -> Some alg;; (* решение найдено *)

let show_items sep f lst = List.map f lst |> String.concat sep;;          
let show_outcome = function LeftLighter -> "left lighter" | LeftHeavier -> "left heavier" | Equal -> "equal";;
let show_weighting ps =
  let select pos =
    List.mapi (fun i p -> i+1,p) ps |> List.find_all (fun (i,p)->p=pos) |> List.map fst in  
  Printf.sprintf "%s and %s" (show_items ", " string_of_int (select Left)) (show_items ", " string_of_int (select Right));;

let tab n = String.make n ' ';;

let rec show_alg n = function (* структурированный вывод решения-алгоритма *)
  | LastStep(w, lst) -> Printf.sprintf "\n%slast weighting [%s]:\n%s" (tab n) (show_weighting w) 
      (show_items "\n" (fun (oc, k) -> Printf.sprintf "%sif %s then fake is %d" (tab (n+1)) (show_outcome oc) (k+1)) lst)
  | Step(w, lst) -> Printf.sprintf "\n%sWeigh [%s]:\n%s" (tab n) (show_weighting w)
      (show_items "\n" (fun (oc, alg) -> 
        Printf.sprintf "%sif %s then %s" (tab (n+1)) (show_outcome oc) (show_alg (n+2) alg)) lst)
  | Found k -> Printf.sprintf "fake is %d" (k+1);;      

let num_weightings = if Array.length Sys.argv > 2 then int_of_string Sys.argv.(2) else 3 in
let num_coins = if Array.length Sys.argv > 1 then int_of_string Sys.argv.(1) else 8 in
match solve (repeat cAny num_coins) num_weightings with (* запуск поиска *)
  | None -> print_endline "no solution found"
  | Some alg -> print_endline (show_alg 0 alg);;