From Stdlib Require Import List Bool Arith Lia PeanoNat.
Import ListNotations.
From LambdaSub Require Import ExecReducer Binding.

Inductive mtrm : Type :=
| mBVar : nat -> mtrm
| mFVar : atom -> mtrm
| mMVar : atom -> list atom -> mtrm
| mApp : mtrm -> mtrm -> mtrm
| mLam : mtrm -> mtrm
| mESub : mtrm -> mtrm -> mtrm.

Fixpoint mfvs (t : mtrm) : list atom :=
  match t with
  | mBVar _ => []
  | mFVar x => [x]
  | mMVar x d => x :: d
  | mApp a b => mfvs a ++ mfvs b
  | mLam a => mfvs a
  | mESub a b => mfvs a ++ mfvs b
  end.

Fixpoint lift_m (k : nat) (t : mtrm) : mtrm :=
  match t with
  | mBVar n => if Nat.ltb n k then mBVar n else mBVar (S n)
  | mFVar x => mFVar x
  | mMVar x d => mMVar x d
  | mApp a b => mApp (lift_m k a) (lift_m k b)
  | mLam a => mLam (lift_m (S k) a)
  | mESub a b => mESub (lift_m (S k) a) (lift_m k b)
  end.

Fixpoint open_rec_m (k : nat) (u : mtrm) (t : mtrm) : mtrm :=
  match t with
  | mBVar n => if Nat.eqb n k then lift_m k u else mBVar n
  | mFVar x => mFVar x
  | mMVar x d => mMVar x d
  | mApp a b => mApp (open_rec_m k u a) (open_rec_m k u b)
  | mLam a => mLam (open_rec_m (S k) u a)
  | mESub a b => mESub (open_rec_m (S k) u a) (open_rec_m k u b)
  end.

Fixpoint close_rec_m (x : atom) (k : nat) (t : mtrm) : mtrm :=
  match t with
  | mBVar n => mBVar n
  | mFVar y => if Nat.eqb y x then mBVar k else mFVar y
  | mMVar y d => mMVar y d
  | mApp a b => mApp (close_rec_m x k a) (close_rec_m x k b)
  | mLam a => mLam (close_rec_m x (S k) a)
  | mESub a b => mESub (close_rec_m x (S k) a) (close_rec_m x k b)
  end.

Fixpoint of_trm (t : trm) : mtrm :=
  match t with
  | BVar n => mBVar n
  | FVar x => mFVar x
  | App a b => mApp (of_trm a) (of_trm b)
  | Lam a => mLam (of_trm a)
  | ESub a b => mESub (of_trm a) (of_trm b)
  end.

Lemma of_trm_fvs : forall t, mfvs (of_trm t) = fvs t.
Proof. induction t; simpl; [reflexivity | reflexivity | rewrite IHt1, IHt2; reflexivity | rewrite IHt; reflexivity | rewrite IHt1, IHt2; reflexivity]. Qed.

Lemma of_trm_lift : forall t k, of_trm (lift k t) = lift_m k (of_trm t).
Proof.
  induction t; intros k; simpl.
  - destruct (Nat.ltb n k); reflexivity.
  - reflexivity.
  - rewrite IHt1, IHt2. reflexivity.
  - rewrite IHt. reflexivity.
  - rewrite IHt1, IHt2. reflexivity.
Qed.

Lemma of_trm_open : forall t k u,
  of_trm (open_rec k u t) = open_rec_m k (of_trm u) (of_trm t).
Proof.
  induction t; intros k u; simpl.
  - destruct (Nat.eqb n k) eqn:E; [| reflexivity].
    apply Nat.eqb_eq in E. subst n. simpl. rewrite of_trm_lift. reflexivity.
  - reflexivity.
  - rewrite IHt1, IHt2. reflexivity.
  - rewrite IHt. reflexivity.
  - rewrite IHt1, IHt2. reflexivity.
Qed.

Lemma of_trm_close : forall t x k,
  of_trm (close_rec x k t) = close_rec_m x k (of_trm t).
Proof.
  induction t; intros x k; simpl.
  - reflexivity.
  - destruct (Nat.eqb a x); reflexivity.
  - rewrite IHt1, IHt2. reflexivity.
  - rewrite IHt. reflexivity.
  - rewrite IHt1, IHt2. reflexivity.
Qed.

Lemma of_trm_injective : forall t t', of_trm t = of_trm t' -> t = t'.
Proof.
  induction t; intros t' H; destruct t'; simpl in H; try discriminate.
  - injection H as H. f_equal. exact H.
  - injection H as H. f_equal. exact H.
  - injection H as H1 H2. f_equal; [apply IHt1 | apply IHt2]; assumption.
  - injection H as H. f_equal. apply IHt. exact H.
  - injection H as H1 H2. f_equal; [apply IHt1 | apply IHt2]; assumption.
Qed.

Print Assumptions of_trm_fvs.
Print Assumptions of_trm_lift.
Print Assumptions of_trm_open.
Print Assumptions of_trm_close.
Print Assumptions of_trm_injective.