From Undecidability Require Import TM.Code.ProgrammingTools LM_heap_def.
From Undecidability.TM.L Require Import Alphabets.
From Undecidability Require Import TM.Code.ListTM TM.Code.CasePair TM.Code.CaseSum TM.Code.CaseNat Hoare.Hoare.

Set Default Proof Using "Type".
Local Arguments plus : simpl never.
Local Arguments mult : simpl never.

Section Lookup.



  Variable sigLookup : finType.

  Variable retr_clos_lookup : Retract sigHClos sigLookup.
  Variable retr_heap_lookup : Retract sigHeap sigLookup.


  Definition retr_nat_clos_ad : Retract sigNat sigHClos :=
    Retract_sigPair_X _ _.
  Definition retr_nat_lookup_clos_ad : Retract sigNat sigLookup :=
    ComposeRetract retr_clos_lookup retr_nat_clos_ad.

  Definition retr_nat_clos_var : Retract sigNat sigHClos :=
    Retract_sigPair_Y _ _.
  Definition retr_nat_lookup_clos_var : Retract sigNat sigLookup :=
    ComposeRetract retr_clos_lookup retr_nat_clos_var.

  Definition retr_nat_heap_entry : Retract sigNat sigHeap :=
    Retract_sigList_X (Retract_sigOption_X (Retract_sigPair_Y _ (Retract_id _))).
  Local Definition retr_nat_lookup_entry : Retract sigNat sigLookup :=
    ComposeRetract retr_heap_lookup retr_nat_heap_entry.


  Definition retr_clos_heap : Retract sigHClos sigHeap := _.
  Definition retr_clos_lookup_heap : Retract sigHClos sigLookup := ComposeRetract retr_heap_lookup retr_clos_heap.

  Definition retr_hent_heap : Retract sigHEntr sigHeap := _.
  Local Definition retr_hent_lookup : Retract sigHEntr sigLookup := ComposeRetract retr_heap_lookup retr_hent_heap.

  Definition : Retract sigHEntr' sigHeap := _.
  Local Definition : Retract sigHEntr' sigLookup := ComposeRetract retr_heap_lookup .


  Definition Lookup_Step : pTM sigLookup^+ (option bool) 5 :=
    If (Nth' retr_heap_lookup retr_nat_lookup_clos_ad @ [|; ; ; |])
       (If (CaseOption retr_hent_lookup @ [||])
           (CasePair sigHClos_fin sigHAdd_fin @ [|; |];;
            If (CaseNat retr_nat_lookup_clos_var @ [||])
               (Return (CopyValue _ @ [|; |];;
                        Translate retr_nat_lookup_entry retr_nat_lookup_clos_ad @ [||];;
                        Reset _ @ [||];;
                        Reset _ @ [| |])
                       None)
               (Return (Reset _ @ [||];;
                        Reset _ @ [||];;
                        Translate retr_clos_lookup_heap retr_clos_lookup @ [||])
                       (Some true)))
           (Return Nop (Some false)))
       (Return Nop (Some false))
  .

  Definition Lookup_Step_size (H : Heap) (a n : ) : Vector.t () 5 :=
    match nth_error H a with
    | Some (Some (g, b))
      match n with
      | S n'
        [| H a @>;
           H a @> >> CopyValue_size b;
           S;
           H a @> >> CasePair_size1 g >> Reset_size g;
           H a @> >> CaseOption_size_Some >> CasePair_size0 g >> Reset_size b|]
      | 0
        [| H a @>;
           H a @>;
           Reset_size 0;
           H a @> >> CasePair_size1 g;
           H a @> >> CaseOption_size_Some >> CasePair_size0 g >> Reset_size b|]
      end
    | _ default
    end.


  Local Definition Lookup_Step_steps_CaseNat (n: ) (e': HClos * HAdd) :=
    let (g,b) := (fst e', snd e') in
    match n with
    | S _ 1 + CopyValue_steps b + 1 + Translate_steps b + 1 + Reset_steps b + Reset_steps g
    | O 1 + Reset_steps b + 1 + Reset_steps 0 + Translate_steps g
    end.

  Local Definition Lookup_Step_steps_CaseOption (n:) (e: HEntr) :=
    match e with
    | Some ((g, b) as e') 1 + CasePair_steps g + 1 + CaseNat_steps + Lookup_Step_steps_CaseNat n e'
    | None 0
    end.

  Local Definition Lookup_Step_steps_Nth' H a n :=
    match nth_error H a with
    | Some e 1 + CaseOption_steps + Lookup_Step_steps_CaseOption n e
    | None 0
    end.

  Definition Lookup_Step_steps (H: Heap) (a: HAdd) (n: ) :=
    1 + H a + Lookup_Step_steps_Nth' H a n.


  Lemma Lookup_Step_SpecT_space H a n ss:
  TripleT
    ≃≃([],withSpace [| ≃(_) H ; ≃(retr_nat_lookup_clos_ad) a;≃(retr_nat_lookup_clos_var) n;Void;Void |] ss)
    (Lookup_Step_steps H a n) Lookup_Step
    ( yout ≃≃([yout = match nth_error (A:=HEntr) H a with Some (Some _ ) match n with 0 Some true | _ None end | _ Some false end]
       ,withSpace
       match nth_error (A:=HEntr) H a with
        Some (Some (g,b) )
        match n with
        | 0 [| ≃(_) H; Void;Void;≃(retr_clos_lookup) g;Void|]
        | S n' [| ≃(_) H;≃(retr_nat_lookup_clos_ad) b;≃(retr_nat_lookup_clos_var) n';Void;Void|]
        end
       | _ SpecVTrue
       end (appSize (Lookup_Step_size H a n) ss))).
  Proof.
    unfold Lookup_Step. remember (Lookup_Step_size H a n) as F eqn:HF.
    eapply If_SpecTReg with (:= match nth_error H a with Some (Some (g,b)) _ | Some _ _ | None _ end) (:=0).
    now hsteps_cbn.
    2:{ destruct nth_error. hintros [=]. cbn. hsteps_cbn. tspec_ext. }
    {
      unfold Lookup_Step_size in HF.
      cbn. destruct nth_error as [ h | ];hintros [=];[].
      eapply If_SpecTReg with (:= match h with Some (g,b) _ | None _ end) (:=0).
      { hsteps_cbn. cbn. tspec_ext. }
      2:{ cbn. destruct h; hintros [=]. hsteps_cbn. cbn. tspec_ext. }
      2:{ cbn. intros ? . destruct h as [[]| ]. 2:reflexivity. shelve. }
      cbn. destruct h as [[g b] | ];hintros [=];[].
      hstep. { hsteps_cbn. cbn. tspec_ext. }
      2:{cbn. unfold CasePair_steps. reflexivity. }
      intros ?. hstep. { hsteps_cbn. cbn. tspec_ext. }
      - cbn. hintros Hn. destruct n. easy.
        hsteps_cbn. cbn.
        { eapply ConsequenceT_pre. refine (Translate_SpecT_size _ _ _ _ [| _|]). 3:tspec_ext. reflexivity. }
        1-3:reflexivity.
        subst F. cbn. tspec_ext.
      - cbn. hintros .
        hsteps_cbn. cbn. { eapply ConsequenceT_pre. refine (Translate_SpecT_size _ _ _ _ [| _|]). 3:tspec_ext. reflexivity. }
        1-2:reflexivity.
        subst F. cbn; tspec_ext.
      - cbn. intros Hb. refine (_ : _ match n with 0 _ | _ _ end).
        destruct , n. 1,4:exfalso;clear - Hb;. all:reflexivity.
    }
    Unshelve. 5:reflexivity. 2,3:exact 0.
    cbn. intros ? . unfold Lookup_Step_steps,Lookup_Step_steps_Nth'. destruct nth_error as [[[[] ]| ]| ]. all:cbn. 2,3:nia. unfold CasePair_steps. destruct n; cbn. all:rewrite !Nat.add_assoc. all:reflexivity.
  Qed.


  Definition Lookup := While Lookup_Step.

  Fixpoint Lookup_size (H : Heap) (a n : ) {struct n} : Vector.t ( ) 5 :=
    match nth_error H a with
    | Some (Some (g, b))
      match n with
      | S n' Lookup_Step_size H a n >>> Lookup_size H b n'
      | 0 Lookup_Step_size H a n
      end
    | _ default
    end.

  Lemma Lookup_size_eq (H : Heap) (a n : ) :
    Lookup_size H a n =
    match nth_error H a with
    | Some (Some (g, b))
      match n with
      | S n' Lookup_Step_size H a n >>> Lookup_size H b n'
      | 0 Lookup_Step_size H a n
      end
    | _ default
    end.
  Proof. destruct n; auto. Qed.

  Definition Lookup_Rel : pRel sigLookup^+ bool 5 :=
     tin '(yout, tout)
       (H: Heap) (a n: ) (s0 s1 s2 s3 s4 : ),
        let size := Lookup_size H a n in
        tin[@] ≃(;) H
        tin[@] ≃(retr_nat_lookup_clos_ad ; ) a
        tin[@] ≃(retr_nat_lookup_clos_var; ) n
        isVoid_size tin[@] isVoid_size tin[@]
        match yout with
        | true
           g,
          lookup H a n = Some g
          tout[@] ≃(;size @> ) H
          isVoid_size tout[@] (size @> )
          isVoid_size tout[@] (size @> )
          tout[@] ≃(retr_clos_lookup; size @> ) g
          isVoid_size tout[@] (size @> )
        | false
          lookup H a n = None
        end.

  Arguments Lookup_Step_size : simpl never.

  Fixpoint Lookup_steps (H : Heap) (a : HAdd) (n : ) : :=
    match nth_error H a with
    | Some (Some (g, b))
      match n with
      | 0 Lookup_Step_steps H a n
      | S n' 1 + Lookup_Step_steps H a n + Lookup_steps H b n'
      end
    | _ Lookup_Step_steps H a n
    end.


  Lemma Lookup_SpecT_space H a n ss:
  TripleT
    ≃≃([],withSpace [| ≃(_) H ; ≃(retr_nat_lookup_clos_ad) a;≃(retr_nat_lookup_clos_var) n;Void;Void |] ss)
    (Lookup_steps H a n) Lookup
    ( yout ≃≃([yout = match lookup H a n with Some _ true | _ false end]
    , withSpace match lookup H a n with Some g [| ≃(_) H;Void;Void;≃(retr_clos_lookup) g; Void|] | _ SpecVTrue end (appSize (Lookup_size H a n) ss))).
  Proof.
    refine (While_SpecTReg (PRE := '(a,n,ss) (_,_))(INV := '(a,n,ss) y ([y = match nth_error H a with
    | Some (Some _) match n with | 0 Some true | S _ None end | _ Some false end],_)) (POST := '(a,n,ss) y (_,_))
       (f__step := '(a,n,ss) Lookup_Step_steps H a n ) (f__loop := '(a,n,ss) _ ) _ _ ((a,n,ss)));clear a n ss;intros [[a n] ss].
    { eapply ConsequenceT. eapply Lookup_Step_SpecT_space. 2:intros. 1,2:cbn - [appSize SpecVTrue]. 1,2:now tspec_ext. reflexivity. }
    all:cbn - [SpecVTrue appSize Lookup_size].
    remember (Lookup_size H a n) as F eqn:HF. remember (Lookup_steps H a n) as F' eqn:HF'. split.
    -destruct n;unfold Lookup_size in HF;unfold Lookup_steps in HF';fold Lookup_size in HF;fold Lookup_steps in HF'.
      +intros ? H'. cbn[lookup Lookup_Step_size]. destruct nth_error as [[[]| ]| ] eqn:Hnth. all:split;[ | subst F';reflexivity].
      all:revert H';intros [= ]. now rewrite HF. 1-2:tspec_ext.
      +intros ? H'. cbn [lookup]. unfold Lookup_Step_steps. destruct nth_error as [[[]| ]| ] eqn:Hnth. easy.
        all:split;[ | subst F';reflexivity]. all:inv H'. all:tspec_ext.
    - destruct (nth_error H a) as [[[g' b]| ]| ] eqn:Hnth. 2-3:easy. destruct n as [ | n]. easy. intros _.
      cbn [lookup]. rewrite Hnth.
      unfold Lookup_size in HF;fold Lookup_size in HF. rewrite Hnth in HF.
      eexists (b,n,_). repeat apply conj.
      + subst F. cbn. tspec_ext.
      + reflexivity.
      +intros. subst F. reflexivity.
  Qed.


  Lemma Lookup_Realise : Lookup Lookup_Rel.
  Proof.
    repeat (eapply RealiseIntroAll;intro). eapply Realise_monotone.
    -eapply TripleT_Realise. eapply Lookup_SpecT_space with (ss:=[| _;_;_;_;_|]).
    -cbn. intros ? [] H **. modpon H.
    {unfold "≃≃",withSpace;cbn. intros i; destruct_fin i;cbn. all:eassumption. }
    repeat destruct _;unfold "≃≃",withSpace in H;cbn in H.
    all:destruct H as [Heq H].
    2,3:discriminate Heq. 2:easy.
    all:specializeFin H. eexists;repeat split; easy.
  Qed.


  Definition Lookup_T : tRel sigLookup^+ 5 :=
     tin k
       (H: Heap) (a n: ),
        tin[@] H
        tin[@] ≃(retr_nat_lookup_clos_ad) a
        tin[@] ≃(retr_nat_lookup_clos_var) n
        isVoid tin[@] isVoid tin[@]
        Lookup_steps H a n k.

  Lemma Lookup_Terminates : Lookup Lookup_T.
  Proof.
    repeat (eapply TerminatesInIntroEx;intro). eapply TerminatesIn_monotone.
    -eapply TripleT_TerminatesIn. eapply TripleT_RemoveSpace,Lookup_SpecT_space.
    -intros ? k H **. modpon H.
    split. 2:eassumption.
    unfold "≃≃",withSpace;cbn. intros i; destruct_fin i;cbn. all:assumption.
  Qed.


End Lookup.

Arguments Lookup_steps : simpl never.
Arguments Lookup_size : simpl never.