(*********************************************)
(* NEERC'07 quarterfinals                    *)
(* Problem K: K'ak'-u-pakal and Mayan script *)
(* OCaml solution: for testing purposes      *)
(*********************************************)

let pf = Printf.printf;;

type expr = Var of int * string
          | Expr of expr * int * int * int;; (* expr*$1/$2+$3 *)

let vars = ref 0;;
let add_var id = let v = !vars in vars := v + 1; Var (v, " "^id^" ");;

type block_type = Small | Main;;
type text = Glyph of expr 
          | HBlock of (text*block_type) list 
          | VBlock of (text*block_type) list;;

(* Parsing source file *)

let inp = open_in "kak-u-pakal.in";;
let s = input_line inp ^ ";";;
close_in inp;;

let n = ref 0;;
let error s = failwith (s ^ " at pos " ^ string_of_int (!n+1));;
let curr() = s.[!n];;
let next() = n := !n + 1;;
let str_of_chr c = String.make 1 c;;
let eat c = (if curr() != c then error (str_of_chr c ^ " expected")); next();;
let after a c = let _ = eat c in a;;

let rec parse_digit s = 
  match curr() with
    '0'..'9' -> let s = s^str_of_chr (curr ()) in next (); parse_digit s
  | _        -> s;;

let rec parse_hieroglyph () =
  match curr() with
    '(' -> next(); after (parse_block_body()) ')'
  | '0'..'9' -> Glyph (add_var (parse_digit ""))
  | _   -> error "Unexpected symbol"
  
and parse_block_body() =
  let b = parse_block() in
  match curr() with
    ':' -> VBlock (parse_vertical [b])
  | '.' -> HBlock (parse_horizontal [b])
  | _   -> error ". or : expected"

and parse_vertical b =
  if curr() = ':' then (next(); parse_vertical (b@[parse_block()])) else b

and parse_horizontal b =
  if curr() = '.' then (next(); parse_horizontal (b@[parse_block()])) else b

and parse_block () =
  if curr() = '[' then (next(); after (parse_hieroglyph(), Main) ']')
  else (parse_hieroglyph(), Small);;

let parse () =
  n := 0;
  parse_hieroglyph ();;

let rec print_tree t =
  match t with
    Glyph (Var (n,id)) -> pf "%d <%s>" n id
  | Glyph (_) -> failwith "Impossible Glyph"
  | VBlock l -> print_block l ":"
  | HBlock l -> print_block l "."

and print_block l c =
  pf "(";
  List.iter (fun (t,m) -> 
    print_string (match m with Small -> "s " | Main -> "M ");
    print_tree t;
    print_string c) l;
  print_string ")";;

(* Building Diophant equations *)

let rec print_equation eq =
  match eq with
    Var (n,s) -> pf "v%d <%s>" n s
  | Expr (e,m,d,c) -> pf "("; print_equation e; pf ")*%d/%d+%d" m d c;;

let rec make_equations t =
  match t with
    Glyph (Var _ as v) -> ([v],[v])
  | Glyph _ -> failwith "Impossible Glyph"
  | VBlock l -> make_equations_block l (fun (x,y) -> (x,y))
  | HBlock h -> make_equations_block h (fun (x,y) -> (y,x))

and make_equations_block l swap =
  let n = List.length l in
  let has_main = List.filter (fun (_,size) -> size = Main) l <> [] in
  let c = if has_main then 1 else 0 in

  swap (List.fold_left (fun (dl,el) (x, mx) -> 
    let (ds,es) = swap (make_equations x) in
    let div = match mx with Small->1|Main->2 in
    (List.map (fun e -> Expr (e,n+c,div,0 (*n-1 for old statement variant*))) ds @ dl, 
     es @ el)
  ) ([],[]) l);;

(* Solving Diophant equations *)

exception Contradiction;;
let try_size eq_set size is_horizontal =
  let a = Array.init !vars (fun _ -> 0) in
  try
    let rec inner sz expr = 
      match expr with
        Var (id,name) -> 
          let min_sz = if is_horizontal then String.length name + 1 
                       (* new statement variant, was without +1 *)
                       else 2 (*new statement variant, was 1*) in
          if sz < min_sz then raise Contradiction;
          if a.(id) = 0 then a.(id) <- sz
          else if a.(id) <> sz then raise Contradiction
      | Expr (e,m,d,c) -> 
        let e_sz' = (sz - c)*d in
        if e_sz' mod m <> 0 then raise Contradiction
        else inner (e_sz' / m) e
    in
    List.iter (inner size) eq_set; Some a
  with
    Contradiction -> None;;

exception Found of int * int array;;
let solve l is_horizontal =
  try
    for i = 1 to 40000 do
      match try_size l i is_horizontal with
        Some a -> raise (Found (i,a))
      | None -> ()
    done;
    None
  with
    Found (s,a) -> Some a;;

let print_solution s =
  match s with 
    Some a -> pf "Solution:\n"; Array.iteri (pf "v%d = %d\n") a; pf "\n"
  | None -> pf "No solution\n";;

(* Draw solution *)

let place text (av,ah) =
  let borders = Array.init (!vars) (fun _ -> (-1,-1,-1,-1,None)) in
  let rec pl text (pv,ph) =
    match text with
      Glyph (Var (id,_) as v) -> 
        borders.(id) <- (pv,ph,pv + av.(id) - 1,ph + ah.(id) - 1,Some v); 
	(* -1 added for new statement variant *)
        (pv + av.(id) (* + 1 for old statement variant *), 
         ph + ah.(id) (* + 1 for old statement variant *))
    | Glyph _ -> failwith "Impossible glyph"
    | HBlock l -> List.fold_left (fun (pv',ph') (t,_) -> pl t (pv,ph')) (pv,ph) l
    | VBlock l -> List.fold_left (fun (pv',ph') (t,_) -> pl t (pv',ph)) (pv,ph) l
  in
    (borders, pl text (0,0));;

let print_place (borders,(szy,szx)) =
  pf "%d*%d:\n" szy szx;
  Array.iter (fun (yl,xl,yr,xr,_) -> pf "(%d,%d)-(%d,%d)\n" yl xl yr xr) borders;;

let draw (borders,(szy,szx)) =
  let img = Array.make_matrix (szy+1) (szx+1) ' ' in

  let put y x c =
    match (img.(y).(x),c) with
      (' ',_)               -> img.(y).(x) <- c
    | (_,'+') | ('+',_)     -> img.(y).(x) <- '+'
    | ('-','-') | ('|','|') -> ()
    | _ -> img.(y).(x) <- '?' (*failwith ("???: " ^ str_of_chr c ^ " over " ^ str_of_chr img.(y).(x))*) in

  Array.iter (fun (yl,xl,yr,xr,v) ->
    match v with 
      Some (Var (_,name)) ->
        for y = yl+1 to yr do put y xl '|'; put y (xr+1) '|' done;
        for x = xl+1 to xr do put yl x '-'; put (yr+1) x '-' done;
        put yl xl '+'; put (yr+1) xl '+'; 
        put yl (xr+1) '+'; put (yr+1) (xr+1) '+';
        for i = 0 to String.length name - 1 do
          put (yl+1) (xl+i+1) name.[i]
        done
    | None -> failwith "none???"
    | _ -> failwith "some???"
  ) borders; img;;

(* Main program *)

let x = parse ();;
(*print_tree x;;*)
let (v,h) = make_equations x;;

print_string "Vertical:";; List.iter (fun e -> print_equation e; pf "\n") v;;
print_string "\nHorizontal:";; List.iter (fun e -> print_equation e; pf "\n") h;;
(*pf "Vertical:\n";;*)

let vs = solve v false;;
(*
print_solution vs;;
pf "Horizontal:\n";;
*)
let hs = solve h true;;
(*print_solution hs;;*)

let out = open_out "kak-u-pakal.out";;
match (hs,vs) with
  (Some hs,Some vs) -> 
    let b = place x (vs,hs) in
(*  print_place b;*)
    let img = draw b in
    for i = 0 to Array.length img - 1 do
      for j = 0 to Array.length (img.(i)) - 1 do
        Printf.fprintf out "%c" img.(i).(j)
      done;
      Printf.fprintf out "\n"
    done
| _ -> ();;
close_out out;;
