open Lam
open Types

exception TacticFailed of string

type hlam = (* hollow lam *)
  HFun of (id * Types.ty) * hlam
  | HApp of hlam * hlam
  | HVar of id
  | HExf of hlam * Types.ty
  | HPair of hlam * hlam
  | HLeft of hlam
  | HRight of hlam
  | Ref of hlam ref
  | Hole

type context = (id * id * hlam * Types.ty) list
type goal = hlam * Types.ty * context
type proof = goal option * goal list

type tactic =
  Exact_term of lam
  | Exact_proof of id
  | Assumption
  | Intros
  | Intro
  | Cut of ty
  | Apply of id
  | Split
  | Right
  | Left

let hyp_count = ref 0
let get_fresh_hyp () =
  let n = string_of_int !hyp_count in
  incr hyp_count;
  let hyp_id = "H" ^ n in
  let var_id = "x_" ^ n in
  (hyp_id, var_id)

(** replace ref's in a proof *)
let clean_hlam assoc (h : hlam) : hlam =
  let rec clean (h : hlam) : hlam= match h with
    HFun ((s, t), h) -> HFun ((s, t), clean h)
    | Hole -> Hole
    | HApp (h1, h2) -> HApp (clean h1, clean h2)
    | HVar s -> HVar s
    | HExf (h, t) -> HExf (clean h, t)
    | HPair (h1, h2) -> HPair (clean h1, clean h2)
    | HLeft h -> HLeft (clean h)
    | HRight h -> HRight (clean h)
    | Ref (hr) ->
      match List.assq_opt hr !assoc with
        None -> let new_h = ref (clean !hr)
          in (assoc := (hr, new_h)::!assoc;
          Ref new_h)
        | Some new_h -> Ref new_h
  in clean h
let rec clean_context assoc (c : context) : context = match c with
  [] -> []
  | (s1, s2, h, t)::q -> (s1, s2, clean_hlam assoc h, t)::(clean_context assoc q)
let clean_goal assoc ((h, t, c) : goal) : goal =
  (clean_hlam assoc h, t, clean_context assoc c)
let clean_proof ((g, gs) : proof) : (hlam ref * hlam ref) list ref * proof =
  let assoc = ref [] in
  let g' = match g with 
    | Some g -> Some (clean_goal assoc g)
    | None -> None
  in assoc, (g', List.map (clean_goal assoc) gs)

let goal_is_over ((g, _) : proof) : bool =
  match g with
  None -> true
  | Some _ -> false

let fill (hole : hlam) (e : hlam) : unit =
  match hole with
  Ref h -> h := e
  | _ -> raise (TacticFailed "not fillable")

let rec hlam_of_lam : lam -> hlam = function
  Fun ((x, t), e) ->
    let e = hlam_of_lam e in
    HFun ((x, t), e)
  | App (e1, e2) ->
    let e1 = hlam_of_lam e1 in
    let e2 = hlam_of_lam e2 in
    HApp (e1, e2)
  | Var x -> HVar x
  | Exf (e, t) ->
    let e = hlam_of_lam e in
    HExf (e, t)
  | Pair (l1, l2) ->
    HPair (hlam_of_lam l1, hlam_of_lam l2)
  | Left l ->
    HLeft (hlam_of_lam l)
  | Right l ->
    HRight (hlam_of_lam l)

let rec lam_of_hlam : hlam -> lam = function
  HFun ((x, t), e) ->
    let e = lam_of_hlam e in
    Fun ((x, t), e)
  | HApp (e1, e2) ->
    let e1 = lam_of_hlam e1 in
    let e2 = lam_of_hlam e2 in
    App (e1, e2)
  | HVar x -> Var x
  | HExf (e, t) ->
    let e = lam_of_hlam e in
    Exf (e, t)
  | HPair (h1, h2) ->
    Pair (lam_of_hlam h1, lam_of_hlam h2)
  | HLeft h ->
    Left (lam_of_hlam h)
  | HRight h ->
    Right (lam_of_hlam h)
  | Ref e_ref -> lam_of_hlam !e_ref
  | Hole -> raise (TacticFailed "can not translate unclosed terms")


let typecheck (e : lam) (expected_t : Types.ty) (cs : context) : bool =
  let gam_of_ctx : context -> Types.gam =
    let f = fun (_, var_id, _, ty) -> (var_id, ty) in
    List.map f
  in
  let g = gam_of_ctx cs in
  try Typing.typecheck g e expected_t
  with _ -> raise (TacticFailed "unable to type")

let rec get_term_by_id (hyp : id) : context -> hlam option =
  function
  [] -> None
  | (hyp',_,  h, _) :: _ when hyp' = hyp -> Some h
  | _ :: cs -> get_term_by_id hyp cs

let rec get_term_by_type (ty : Types.ty) : context -> hlam option =
  function
  [] -> None
  | (_, _, h, ty') :: _ when ty' = ty -> Some h
  | _ :: cs -> get_term_by_type ty cs

let next_goal (gs : goal list) : (goal option * goal list) =
  match gs with
  [] -> None, []
  | g :: gs -> Some g, gs

let tact_exact_term ((g, gs) : proof) (e : lam) : proof =
  match g with
  None -> raise (TacticFailed "no current goal")
  | Some (h, expected_t, cs) ->
    if typecheck e expected_t cs
    then
      begin
      fill h (hlam_of_lam e);
      next_goal gs
      end
    else raise (TacticFailed "type mismatch")

let tact_exact_proof ((g, gs) : proof) (hyp : id) : proof =
  match g with
  None -> raise (TacticFailed "no current goal")
  | Some (h, expected_t, cs) ->
    match get_term_by_id hyp cs with
    Some h' ->
      if typecheck (lam_of_hlam h') expected_t cs
      then
        begin
        fill h h';
        next_goal gs
        end
      else raise (TacticFailed "type mismatch")
    | None -> raise (TacticFailed "")

let tact_assumption ((g, gs) : proof) : proof =
  match g with
  None -> raise (TacticFailed "no current goal")
  | Some (h, goal_ty, cs) ->
    match get_term_by_type goal_ty cs with
    None -> (* failwith "assumption failed" *) (g, gs)
    | Some h' ->
      fill h h';
      next_goal gs

let tact_intro ((g, gs) : proof) : proof =
  match g with
  None -> raise (TacticFailed "no current goal")
  | Some (h, goal_ty, cs) ->
    match goal_ty with
    Arr (t1, t2) ->
      let (hyp_id, var_id) = get_fresh_hyp () in
      let cs = (hyp_id, var_id, HVar var_id, t1) :: cs in
      let new_h = Ref (ref Hole) in
      fill h (HFun ((var_id, t1), new_h));
      Some (new_h, t2, cs), gs
    | _ -> (* failwith "expected function" *) (* (g, gs) *)
            raise (TacticFailed "expected function")

let tact_cut ((g, gs) : proof) (new_t : Types.ty) : proof =
  match g with
  None -> raise (TacticFailed "no current goal")
  | Some (h, goal_ty, cs) ->
    (* subgoal 2 : new_t -> goal_ty *)
    let arrow_h = Ref (ref Hole) in
    let arrow_goal = (arrow_h, Arr (new_t, goal_ty), cs) in
    let gs = arrow_goal :: gs in
    (* subgoal 1 (main goal) : new_t *)
    let new_h = Ref (ref Hole) in
    fill h (HApp (arrow_h, new_h));
    Some (new_h, new_t, cs), gs

let tact_apply ((g, gs) : proof) (hyp_id : id) : proof =
  let rec is_implied (goal_ty : ty) (t : ty) : bool =
    match t with
    t when t = goal_ty -> true
    | Arr (_, t2) -> is_implied goal_ty t2
    | _ -> false
  in
  let rec generate_goals (goal_ty : ty) (impl_ty : ty) (impl_h : hlam)
  (goal_h : hlam) (h : hlam) (cs : context) (gs : goal list) : proof =
    match impl_ty with
    Arr (t1, t2) when t2 = goal_ty ->
      let sub_h = Ref (ref Hole) in
      let _ = fill h sub_h in
      let _ = fill goal_h impl_h in
      Some (sub_h, t1, cs), gs
    | Arr (t1, t2) ->
      let sub_h = Ref (ref Hole) in
      let new_h = Ref (ref Hole) in
      let _ = fill h sub_h in
      let impl_h = HApp (impl_h, new_h) in
      let gs = (sub_h, t1, cs) :: gs in
      generate_goals goal_ty t2 impl_h goal_h new_h cs gs
    | _ -> failwith "impossible"
  in
  let rec get_hyp : context -> (hlam * ty) = function
    [] -> raise (TacticFailed "no such hypothesis in context")
    | (hyp_id', _, h', t') :: _ when hyp_id = hyp_id' -> (h', t')
    | _ :: cs -> get_hyp cs
  in
  match g with
  None -> raise (TacticFailed "no current goal")
  | Some (goal_h, goal_ty, cs) ->
    let impl_h, impl_ty = get_hyp cs in
    let new_h = Ref (ref Hole) in
    let impl_h_2 = HApp (impl_h, new_h) in
    if is_implied goal_ty impl_ty
    then generate_goals goal_ty impl_ty impl_h_2 goal_h new_h cs gs
    else raise (TacticFailed "not an implication")

let tact_intros : proof -> proof =
  let rec push (p : proof) =
    try
      let p = tact_intro p in
      push p
    with TacticFailed _ -> p
  in push

let tact_split ((g, gs) : proof) : proof =
  match g with
  None -> raise (TacticFailed "no current goal")
  | Some (h, goal_ty, cs) ->
    match goal_ty with
    | And(t1, t2) ->
      let h1 = Ref (ref Hole) in
      let h2 = Ref (ref Hole) in
      fill h (HPair (h1, h2));
      Some (h1, t1, cs), (h2, t2, cs)::gs
    | _ -> raise (TacticFailed "Not a conjonction")


let tact_right ((g, gs) : proof) : proof =
  match g with
  None -> raise (TacticFailed "no current goal")
  | Some (h, goal_ty, cs) ->
    match goal_ty with
    | Or(_, t) ->
      let new_h = Ref (ref Hole) in
      fill h (HRight new_h);
      Some (new_h, t, cs), gs
    | _ -> raise (TacticFailed "Not a disjunction")

let tact_left ((g, gs) : proof) : proof =
  match g with
  None -> raise (TacticFailed "no current goal")
  | Some (h, goal_ty, cs) ->
    match goal_ty with
    | Or(t, _) ->
      let new_h = Ref (ref Hole) in
      fill h (HLeft new_h);
      Some (new_h, t, cs), gs
    | _ -> raise (TacticFailed "Not a disjunction")

let apply_tactic (p : proof) (t : tactic) : proof =
  match t with
  Exact_term e -> tact_exact_term p e
  | Exact_proof hyp -> tact_exact_proof p hyp
  | Intros -> tact_intros p
  | Intro -> tact_intro p
  | Assumption -> tact_assumption p
  | Cut t -> tact_cut p t
  | Apply h -> tact_apply p h
  | Split -> tact_split p
  | Right -> tact_right p
  | Left -> tact_left p


let tact_try (p : proof) (t : tactic) : proof =
  try apply_tactic p t
  with TacticFailed _ -> p