From Coq Require Import Lia Arith.PeanoNat Classical_Prop Classes.RelationClasses.
From MMaps Require Import MMaps.
From DeBrLevel Require Import Level LevelInterface MapExtInterface MapExt
                              MapLevelInterface MapLevelKD.

Module IsLvlMapLVLD (Data : IsLvlETWL) (M : Interface.S Level)
                    (MO : MapLVLInterface Data M) <: IsLvlMapLVLDInterface Data M MO.

Include IsLvlMapKD Level Data M MO.
Import MO OP.P.

Lemma shift_new_notin (lb k : Lvl.t) (t : t) : lb ≥ (new_key t) → ¬ M.In lb (shift lb k t).
Proof.
  induction t using map_induction; intros Hgeq.
  - rewrite shift_Empty; auto.
    rewrite Empty_eq; auto.
    apply not_in_empty.
  - unfold Add in H0; rewrite H0 in *; clear H0.
    rewrite shift_add.
    rewrite new_key_add_max in Hgeq.
    rewrite add_in_iff.
    intros [Heq | HIn]; subst.
    -- unfold Level.shift in Heq.
       destruct (Lvl.leb_spec0 lb x) as [Hle | Hnle]; lia.
    -- apply IHt1; auto; lia.
Qed.

Lemma shift_max_refl (lb k : Lvl.t) (t : t) :
  lb ≥ (new_key t) → max_key (shift lb k t) = max_key t.
Proof.
  induction t using map_induction; intro Hleb.
  - rewrite shift_Empty; auto.
  - unfold Add in H0; rewrite H0 in *; clear H0.
    rewrite new_key_add_max in ×.
    rewrite shift_add.
    assert (Hle1: lb ≥ new_key t1) by lia.
    apply IHt1 in Hle1.
    do 2 rewrite max_key_add_max.
    rewrite Hle1; f_equal.
    rewrite Level.shift_wf_refl; auto.
    unfold Level.Wf.
    assert (HIn : M.In x (M.add x e t1)).
    { rewrite add_in_iff; now left. }
    apply new_key_in in HIn; lia.
Qed.

Lemma shift_new_refl (lb k : Lvl.t) (t : t) :
  lb ≥ (new_key t) → new_key (shift lb k t) = new_key t.
Proof.
  intro Hge; repeat rewrite new_key_unfold.
  destruct (M.is_empty t) eqn:HEmp.
  - rewrite shift_Empty.
    -- now rewrite HEmp.
    -- now apply is_empty_2.
  - destruct (M.is_empty (shift lb k t)) eqn:HEmp'.
    -- rewrite shift_Empty in HEmp'.
       + now rewrite HEmp in ×.
       + apply is_empty_2 in HEmp'.
         now rewrite <- shift_Empty_iff in HEmp'.
    -- f_equal; now apply shift_max_refl.
Qed.

Lemma shift_max_key_gt_iff (m n x : Lvl.t) (o : t) :
  max_key o > x ↔ max_key (shift m n o) > (Level.shift m n x).
Proof.
  revert x; induction o using map_induction; intro y.
  - split; intro Hmax.
    -- rewrite max_key_Empty in Hmax; auto; lia.
    -- rewrite shift_Empty in Hmax; auto.
       rewrite max_key_Empty in Hmax; auto; lia.
  - split; intro Hmax;
    unfold Add in H0; rewrite H0 in *; clear H0.
    -- rewrite shift_add.
       rewrite max_key_add_max.
       destruct (Lvl.leb_spec0 (max_key o1) x) as [Hle | Hgt].
       + rewrite max_key_add_l in Hmax; auto.
         apply (Level.shift_gt_iff m n) in Hmax; lia.
       + rewrite max_key_add_r in Hmax by lia.
         apply IHo1 in Hmax; lia.
    -- rewrite shift_add in Hmax.
       rewrite max_key_add_max.
       destruct (Lvl.leb_spec0 (max_key (shift m n o1)) (Level.shift m n x)) as [Hle' | Hgt'].
       + rewrite max_key_add_l in Hmax; auto.
         apply (Level.shift_gt_iff m n) in Hmax; lia.
       + rewrite max_key_add_r in Hmax by lia.
         apply IHo1 in Hmax; lia.
Qed.

Lemma shift_max_key_le (m n : Lvl.t) (o : t) :
  max_key o ≤ max_key (shift m n o).
Proof.
  induction o using map_induction.
  - now rewrite shift_Empty.
  - unfold Add in H0; rewrite H0; clear H0.
    rewrite shift_add.
    do 2 rewrite max_key_add_max.
    assert (x ≤ (Level.shift m n x)) by apply Level.shift_le.
    lia.
Qed.

Lemma shift_new_key_le (m n : Lvl.t) (o : t) :
  new_key o ≤ new_key (shift m n o).
Proof.
  repeat rewrite new_key_unfold.
  destruct (M.is_empty o) eqn:Hemp; try lia.
  destruct (M.is_empty (shift m n o)) eqn:Hemp'.
  - apply is_empty_2 in Hemp'.
    rewrite <- (shift_Empty_iff m n) in Hemp'.
    apply is_empty_1 in Hemp'.
    rewrite Hemp in Hemp'; now inversion Hemp'.
  - assert (max_key o ≤ max_key (shift m n o)); try lia.
    apply shift_max_key_le.
Qed.

Lemma Wf_in_iff (m n : Lvl.t) (x : M.key) (t : t) :
  Wf m t → M.In x (shift m n t) ↔ M.In x t.
Proof.
  revert x; induction t using map_induction; intros y Hvo; split; intro HIn.
  - rewrite shift_Empty in HIn; auto.
  - rewrite shift_Empty; auto.
  - unfold Add in *; rewrite H0 in *; clear H0.
    apply Wf_add_notin in Hvo as [Hvk [Hve Hv]]; auto.
    rewrite shift_add_notin in HIn; auto.
    rewrite add_in_iff in *; destruct HIn as [Heq | HIn]; subst.
    -- left; now rewrite Level.shift_wf_refl; auto.
    -- right. rewrite <- IHt1; eauto.
  - unfold Add in *; rewrite H0 in *; clear H0.
    apply Wf_add_notin in Hvo as [Hvk [Hve Hv]]; auto.
    rewrite shift_add_notin; auto.
    rewrite add_in_iff in *; destruct HIn as [Heq | HIn]; subst.
    -- left; now rewrite Level.shift_wf_refl; auto.
    -- right. rewrite IHt1; eauto.
Qed.

Lemma shift_find_Wf (lb k : Lvl.t) (x : M.key) (m m' : t) :
  Level.Wf lb x → M.In x m →
  M.find x m = M.find x m' → M.find x (shift lb k m) = M.find x (shift lb k m').
Proof.
  intros Hvk HIn Hfi.
  destruct HIn as [v HfV]; apply find_1 in HfV.
  rewrite <- (Level.shift_wf_refl lb k x); auto.
  apply (shift_find_iff lb k) in HfV as HfV1.
  rewrite HfV1; symmetry.
  rewrite <- shift_find_iff.
  now rewrite <- Hfi.
Qed.

Lemma shift_in_new_key (n : Lvl.t) (x : M.key) (t : t) :
 M.In x (shift (new_key t) n t) ↔ M.In x t.
Proof.
  split; intro HIn.
  - apply shift_in_e in HIn as H.
    destruct H as [y Heq]; subst.
    rewrite <- shift_in_iff in HIn.
    apply new_key_in in HIn as Hlt.
    rewrite Level.shift_wf_refl; auto.
  - apply new_key_in in HIn as Hlt.
    rewrite <- (Level.shift_wf_refl (new_key t) n x); auto.
    now rewrite <- shift_in_iff.
Qed.

End IsLvlMapLVLD.

Module IsBdlLvlMapLVLD
  (Data : IsBdlLvlETWL) (M : Interface.S Level)
  (MO : MapLVLInterface Data M) <: IsBdlLvlMapLVLDInterface Data M MO.

Include IsLvlMapLVLD Data M MO.
Import MO OP.P.

Lemma shift_wf_refl (lb k : Lvl.t) (t : t) : Wf lb t → eq (shift lb k t) t.
Proof.
  induction t using map_induction; intro Hvt.
  - now apply shift_Empty.
  - unfold Add in H0; rewrite H0 in *; clear H0.
    apply Wf_add_notin in Hvt as [Hvx [Hve Hvt]]; auto.
    rewrite shift_add.
    rewrite IHt1; auto.
    rewrite Level.shift_wf_refl; auto.

    replace (Data.shift lb k e) with e; try reflexivity.
    apply Data.eq_leibniz.
    now rewrite Data.shift_wf_refl; auto.
Qed.

End IsBdlLvlMapLVLD.

Module MakeLvlMapLVLD (Data : IsLvlETWL) <: IsLvlET.

Module Raw := MMaps.OrdList.Make Level.
Module Ext := MapETLVL Data Raw.
Include IsLvlMapLVLD Data Raw Ext.
Include OP.P.

End MakeLvlMapLVLD.

Module MakeBdlLvlMapLVLD (Data : IsBdlLvlETWL) <: IsBdlLvlET.

Module Raw := MMaps.OrdList.Make Level.
Module Ext := MapETLVL Data Raw.
Include IsBdlLvlMapLVLD Data Raw Ext.
Include OP.P.

End MakeBdlLvlMapLVLD.