From Undecidability.L Require Import L_facts Tactics.LTactics.

From Complexity.L.AbstractMachines Require Import AbstractHeapMachineDef UnfoldTailRec FunctionalDefinitions.
From Complexity.L.AbstractMachines.Computable Require Import Unfolding HeapMachine Shared.

From Undecidability.L.Datatypes Require Import Lists.
From Complexity.L.Datatypes Require Import LBinNums.
From Undecidability.L.Functions Require Import UnboundIteration Proc.
From Complexity.L.Functions Require Import BinNums BinNumsCompare.

Import AbstractHeapMachineDef.clos_notation .

Definition extractRes (sigma : state) :=
  match sigma with
    ([],[g],_) => Some (g)
  | _ => None
  end.

Definition evalForTime__step : N * state -> N * state + option (clos * list heapEntry) :=
  fun '(i,sigma) =>
    let '(_,_,H):= sigma in
    match extractRes sigma with
      Some g => inr (Some (g,H))
    | _ => if (0 <? i)%N then
            match heapStep sigma with
              Some sigma' => inl (N.pred i,sigma')
            | _ => inr None
            end
          else inr None
    end.

Definition init__evalForTime (fuel:N) (s:term) :=
  ((4 * fuel +1)%N,init s).

Definition evalForTime (fuel : N) ( s:term) : option (clos * list heapEntry) :=
  if closedb s then
    let '(fuel,x) := init__evalForTime fuel s in
    
    match loopSum (N.to_nat fuel + 1) evalForTime__step (fuel,x) with
      Some g => g
    | _ => None
    end
  else None.

Lemma spec_extractRes sigma:
  match extractRes sigma with
    Some g => exists H, sigma = ([],[g],H)
  | None => forall g H, sigma <> ([],[g],H)
  end.
Proof.
  unfold extractRes.
  all:repeat (let eq := fresh "eq" in destruct _ eqn:eq).
  all:try congruence. eexists _;repeat f_equal. congruence.
Qed.

Arguments extractRes : simpl never.

Import AbstractHeapMachineDef.

Lemma spec_evalForTime__step i sigma:
  match evalForTime__step (i,sigma) with
    inl (i',sigma') => i = N.succ i' /\ step sigma sigma' /\ extractRes sigma = None
  | inr None => (i = 0%N \/ terminal step sigma) /\ extractRes sigma = None
  | inr (Some (g,H)) => sigma = ([],[g],H)
  end.
Proof.
  specialize (heapStep_sound sigma) as R.
  specialize (spec_extractRes sigma) as H''.
  all:cbn.
  all:destruct sigma as [[T V] H];cbn [snd].
  all:destruct (N.ltb_spec0 0 i);[|replace i with (0%N) by Lia.nia].
  all:repeat (let eq := fresh "eq" in destruct _ eqn:eq).
  all:try discriminate.
  all:subst.
  all:repeat match goal with
               H : inl _ = inl _ |- _ => inv H
             | H : inr _ = inr _ |- _ => inv H
             end.
  1:now split;(tauto + Lia.nia).
  all:try (destruct H'' as (?&H'');inv H'').
  all:easy.
Qed.

Instance termT_extractRes : computableTime' extractRes (fun _ _ => (23,tt)).
Proof.
  unfold state.
  extract. solverec.
Qed.

Definition time_evalForTime__step x :=
  let '(i,x) := x in
  let '(T,V,H):=x in
  match extractRes x with
    Some g => 0
  | _ => N.size_nat i * 12
        + if (0 <? i)%N
          then heapStep_time T H
          else 0
  end + 81.
Import N.
Instance termT_evalForTime__step : computableTime' evalForTime__step (fun x _ => (time_evalForTime__step x,tt)).
Proof.
  unfold state.
  extract. intros (i&((T&V)&H)). unfold time_evalForTime__step.   all:solverec.
Qed.

Arguments evalForTime__step : simpl never.

Import L_facts ARS. Import AbstractMachines.AbstractHeapMachineDef.

Instance termT_init__evalForTime : computableTime' init__evalForTime (fun fuel (_:unit) => (1,fun s (_:unit) => (size s * 108 + N.size_nat fuel * 84 + 244,tt))).
Proof.
  eapply computableTimeExt with (x:=fun fuel s => ((fuel + fuel + fuel + fuel +1)%N,init s)).
  -unfold init__evalForTime. repeat intro. hnf. f_equal. Lia.nia.
  -extract.
   recRel_prettify2. reflexivity.
   change (N.size_nat 1) with 1. ring_simplify.

   repeat rewrite N_size_nat_add_leq.
   rewrite !Nat.max_l. all:try Lia.lia.
Qed.

Tactic Notation "destruct" "*" "eqn" :=
  let E := fresh "eq" in destruct * eqn:E.

Definition R__step := fun '(i,s) '(i',s') => i = N.succ i' /\ step s s'.

Lemma pow_Rstep k i i' x x':
  ARS.pow R__step k (i,x) (i',x')
  <-> (i = i' + N.of_nat k)%N /\ ARS.pow step k x x'.
Proof.
  induction k in i',x'|-*.
  {unfold pow. cbn. split. now intros [=]. intros. f_equal; intuition try solve [lia|easy]. }
  replace (S k) with (k + 1) in * by lia.
  rewrite !pow_add. unfold rcomp.
  repeat setoid_rewrite <- rcomp_1.
  split.
  -intros ([]&(->&?)%IHk&?&?). intuition. lia. eauto.
  -intros (->&(?&?&?)).
   eexists (_,_). rewrite IHk. cbn. repeat split. 2,3:eassumption. lia.
Qed.

Import Undecidability.L.AbstractMachines.LargestVar Complexity.L.AbstractMachines.AbstractHeapMachine.

Lemma time_uiter_evalForTime__step s fuel:
  uiterTime evalForTime__step (fun x (_:unit) => (time_evalForTime__step x,tt)) (N.to_nat fuel + 1) (fuel,init s) <=
  (N.to_nat fuel + 1) * (N.size_nat fuel * 12 + heapStep_timeBound (largestVar s) (N.to_nat fuel) + 90).
Proof.
  rewrite @uiterTime_computesRel
    with (X:=(N*state)%type)(R:= R__step)
         (t__step := N.size_nat fuel * 12 + heapStep_timeBound (largestVar s) (N.to_nat fuel) + 81)
         (t__end := 0).
  2:{
    intros fuel' (i'&x') ? R.
    eapply pow_Rstep in R as (->&R). cbn [fst].
    specialize spec_evalForTime__step with (i:=i') (sigma:=x') as H'.
    specialize (spec_extractRes x') as Hx'.
    -destruct x' as [[T V] Hp].
     cbn [fst snd] in *.
     unfold time_evalForTime__step.
     destruct extractRes.
     +destruct evalForTime__step as [[]|]. now destruct H'. split. 2:tauto. lia.
     +split.
      *rewrite N_size_nat_monotone with (n':=(i' + N.of_nat fuel')%N). 2:Lia.nia.
       destruct N.ltb_spec0 with (x:=0%N) (y:=i'%N). 2:Lia.nia.
       rewrite heapStep_timeBound_le. 2:eassumption.
       rewrite heapStep_timeBound_mono with (k':=(N.to_nat (i' + N.of_nat fuel'))). 2:Lia.nia. Lia.nia.
      *destruct evalForTime__step as [[]|]. 2:easy.
       split. 2:easy. easy.
  }
  Lia.nia.
Qed.

Lemma sound_evalForTime__step fuel x: {res & loopSum (N.to_nat fuel +1 ) evalForTime__step (fuel,x) = Some res}.
Proof.
  remember (N.to_nat fuel) as n eqn:eqn.
  induction n in eqn,fuel,x|-*.
  all:cbn.
  all:specialize (spec_evalForTime__step fuel x) as H'.
  all:destruct evalForTime__step as [[]|].
  1:exfalso;Lia.nia.
  1,3:eauto.
  {destruct H' as (->&?&?).
   edestruct IHn with (fuel:=n0). Lia.nia.
   eauto.
  }
Qed.

Import HOAS_Notations.

Definition term_evalForTime : term := Eval cbn [convert TH]in
  [L_HOAS λ fuel s, (!!(extT (closedb)) s)
                  (λ _, (!!(uiter evalForTime__step)
                          (!!(extT (init__evalForTime)) fuel s)))
                  (λ _, !!(enc (None (A:=clos * list heapEntry)))) !!I ].

Definition t__evalForTime maxVar (size:nat) fuel :=
  let fuel' := (4*fuel + 1) in
  size *139 + N.size_nat (N.of_nat fuel) * 84+
  (fuel' + 1) * (N.size_nat (N.of_nat fuel') * 12
                 + heapStep_timeBound maxVar fuel' +90) +264.

Instance evalForTime_comp : computableTime' evalForTime (fun fuel _ => (1,fun s _ => (t__evalForTime (largestVar s)(size s) (N.to_nat fuel),tt))).
Proof.
  exists term_evalForTime. unfold term_evalForTime,evalForTime.
  Intern.cstep. Intern.cstep. Intern.cstep.
  destruct closedb. 1,2:reflexivity. Unshelve.
  4:{ eapply uiter_sound. unfold evalForTime,init__evalForTime. edestruct sound_evalForTime__step as (?&eq').
      rewrite eq'. reflexivity. }
  2:exact True.
  {intros fuel _. recRel_prettify2. easy.
   all:unfold t__evalForTime. 2:lia. unfold init__evalForTime. rewrite time_uiter_evalForTime__step.
   pose (fuel':=(4*fuel + 1)%N).
   fold fuel'.
   replace (4 * N.to_nat fuel + 1) with (N.to_nat fuel') by (unfold fuel';lia).
   rewrite !Nnat.N2Nat.id. lia.
  }
Qed.

Lemma rel_evalForTime__step (x : N * state) : match evalForTime__step x with
                                            | inl y => R__step x y
                                            | inr _ => terminal R__step x
                                            end.
Proof.
  destruct x as [i x]. assert (H':=spec_evalForTime__step i x).
  destruct evalForTime__step as [[]|[[]|]].
  -destruct H' as (->&?&?).
   hnf. easy.
  -intros [i' x'] R'. destruct R' as (->&R').
   assert (Hx:=spec_extractRes x). destruct extractRes.
   +destruct Hx as (?&->).
    inv R'.
   +easy.
  -intros [] []. destruct H' as [[->|Ter] ?]. easy.
   edestruct Ter. easy.
Qed.

Lemma sound_evalForTime__step2' i sigma g H:
  loopSum (N.to_nat i + 1) evalForTime__step (i,sigma) = Some (Some (g,H))
  <-> (exists k, k <= N.to_nat i /\ ARS.pow step k sigma ([],[g],H)).
Proof.
  split.
  -intros H'.
   eapply loopSum_rel_correct2 with (R:=R__step) in H' as (k&(x'&i')&?&eq1&R&Ter&eq2). 2:exact rel_evalForTime__step.
   eapply pow_Rstep in R as (->&R).
   eassert (H'':=spec_evalForTime__step _ _ ).
   unfold state in *.
   rewrite eq2 in H''.
   eassert (H':=spec_extractRes i').
   destruct extractRes;cbn;try congruence.
   destruct H' as (?&->). inv H''. eexists;split. 2:eauto. lia.
  -intros (k'&?&?).
   induction k' in i,sigma,k',H1,H0 |-*.
   +inv H1. replace (N.to_nat i + 1) with (1 + N.to_nat i). cbn.
    specialize (spec_evalForTime__step i ([],[g],H)). destruct evalForTime__step as [[]|[[]|]]. all:easy.
   +replace (S k') with (1+k') in H1 by easy.
    assert (H1':=H1).
    eapply pow_add in H1 as (sigma'&R&H1).
    eapply rcomp_1 in R.
    assert (exists i', N.to_nat i = S (N.to_nat i')) as (i'&eq) by (exists (i-1)%N;lia).
    rewrite eq in *. cbn.
    specialize (spec_evalForTime__step i sigma). destruct evalForTime__step as [[? sigma'']|[[]|]].
    *intros (->&?&?). replace n with i' in * by lia.
     apply IHk'. lia.
     specialize parametrized_semi_confluence with (X:=state) (t1:=([],[g],H) ) (3:=H2) (2:=H1') as (?&?&?&?&?&?&?&?);[now eauto using functional_uc,step_functional|].
     destruct x. 2:{ destruct H6 as (?&?&?);easy. }
     inv H6.
     replace k' with x0 by lia. easy.
    *intros ->. easy.
    *intros [[|Ter ]]. easy. edestruct Ter;easy.
Qed.

Lemma evalForTime_spec (s:term) (fuel : N):
  match evalForTime fuel s with
  |Some (g,H) =>
   closed s
   /\ (exists k, k <= N.to_nat (4* fuel+1) /\ ARS.pow step k (init s) ([],[g],H))
   /\ exists t, reprC H g t /\ s ⇓(<= N.to_nat fuel) t
  | None => ~ closed s \/ ~ exists t, s ⇓(<= N.to_nat fuel) t
  end.
Proof.
  unfold evalForTime. destruct (closedb_spec s). 2:easy.
  cbn [init__evalForTime].
  destruct loopSum as [[[]|]|] eqn:eq.
  -split. easy.
   eapply sound_evalForTime__step2' in eq. split. easy.
   destruct eq as (?&?&?).
   edestruct soundness as (?&?&?&?&eq&?&?&?). 2:split. 1,2:eassumption. now intros ? ?.
   inv eq.
   rewrite ResourceMeasures.timeBS_evalIn in H1.
   do 2 esplit. eassumption. eapply evalIn_evalLe. 2:eassumption. lia.
  -right. intros (?&(?&?&R')%evalLe_evalIn).
   rewrite <- ResourceMeasures.timeBS_evalIn in R'.
   eapply correctTime in R' as (?&?&?&?). 2:easy.
   eassert (H':=sound_evalForTime__step2' _ _ _ _). rewrite eq in H'. destruct H' as [_ H']. discriminate H'.
   repeat eexists. 2:eassumption. lia.
  -edestruct sound_evalForTime__step as (?&eq'). rewrite eq' in eq. easy.
Qed.

Lemma mono_t__evalForTime maxVar maxVar' x x' size size' :
  maxVar <= maxVar' -> x <= x' -> size <= size' -> t__evalForTime maxVar size x <= t__evalForTime maxVar' size' x'.
Proof.
  intros H1 H2 H3.
  unfold t__evalForTime.
  repeat (lazymatch goal with
            |- _ + _ <= _ + _ => eapply plus_le_compat
          | |- _ * _ <= _ * _ => eapply mult_le_compat
          | |- _ => first [eassumption | reflexivity | eapply N_size_nat_monotone | eapply unfoldBool_time_mono | Lia.nia |eapply heapStep_timeBound_mono']
          end).
Qed.

Lemma suplin_t__evalForTime maxVar size x : x <= t__evalForTime maxVar size x.
Proof.
  unfold t__evalForTime. Lia.nia.
Qed.