Add LoadPath "/Users/erutuf/Documents/mit/fiat/src" as Fiat.
Require Import Fiat.Common Fiat.Computation
  Fiat.ADTRefinement Fiat.ADTNotation Fiat.ADTRefinement.BuildADTRefinements.
Require Import List Omega.

Set Implicit Arguments.
Set Regular Subst Tactic.
Import ListNotations.

Inductive removed A : nat -> list A -> list A -> Prop :=
| removed_O : forall a l, removed O (a :: l) l
| removed_Sn : forall n a l l', removed n l l' -> removed (S n) (a :: l) (a :: l').

Lemma removed_len : forall A n (l l' : list A), removed n l l' ->
  length l = S (length l').
Proof with simpl in *.
  induction n; intros...
  - inversion H. subst. auto.
  - inversion H. subst... f_equal. auto.
Qed.

Definition remove_spec A (n : nat)(l : list A) : Comp (list A) :=
  { l' | removed n l l' }%comp.

Ltac return_impl := repeat
  match goal with
  | H : forall v, (ret ?x) ↝ v -> _ |- _ => specialize (H _ (@ReturnComputes _ _))
  end.
Ltac refines := unfold refine in *; intros; return_impl;
  repeat computes_to_econstructor; repeat computes_to_inv; subst.
Ltac monad_simpl := autosetoid_rewrite with refine_monad;
  try simplify_with_applied_monad_laws; simpl.

Lemma remove_ex A n (l : list A) (H : n < length l) :
  {l' | refine (remove_spec n l) (ret l')}.
Proof with simpl in *.
  unfold remove_spec. revert l H.
  induction n; intros; destruct l; simpl in *; [omega|..].
  - exists l. refines. constructor.
  - omega.
  - assert (n < length l) as H0 by omega.
    apply IHn in H0. destruct H0.
    exists (a :: x). refines. constructor. auto.
Defined.

Definition remove A n (l : list A) (H : n < length l) :=
  ` (@remove_ex A n l H).

Inductive removed_lst A : nat -> nat -> list (list A) -> list (list A) -> Prop :=
| removed_lst_O : forall m l l' lst, removed m l l' -> removed_lst O m (l :: lst) (l' :: lst)
| removed_lst_Sn : forall n m l lst lst',
    removed_lst n m lst lst' -> removed_lst (S n) m (l :: lst) (l :: lst').

Inductive lst_len A : nat -> list (list A) -> list (list A) -> Prop :=
| lst_len_O : forall l l' lst lst',
    length l = S (length l') ->
    Forall2 (fun l1 l2 => length l1 = length l2) lst lst' ->
    lst_len O (l :: lst) (l' :: lst')
| lst_len_Sn : forall n l l' lst lst',
    lst_len n lst lst' ->
    length l = length l' ->
    lst_len (S n) (l :: lst) (l' :: lst').

Lemma removed_lst_len : forall A n m (lst lst' : list (list A)),
  removed_lst n m lst lst' -> lst_len n lst lst'.
Proof with simpl in *.
  induction n; intros...
  - inversion H; subst; constructor.
    + eauto using removed_len.
    + clear H H0. induction lst0; auto.
  - inversion H; subst. constructor; eauto.
Qed.

Definition remove_lst_spec A n m (lst : list (list A)) : Comp (list (list A)) :=
  { lst' | removed_lst n m lst lst' }%comp.

Inductive cond A : nat -> nat -> list (list A) -> Type :=
| cond_O : forall m l lst, m < length l -> cond O m (l :: lst)
| cond_Sn : forall n m l lst, cond n m lst -> cond (S n) m (l :: lst).

Definition remove_lst_O A m l (lst : list (list A)) :=
  l' <- remove_spec m l;
  ret (l' :: lst).

Lemma remove_lst_O_ex : forall A m l (lst : list (list A)),
  refine (remove_lst_spec O m (l :: lst)) (remove_lst_O m l lst).
Proof with simpl in *.
  intros. unfold remove_lst_O. refines.
  inversion H; subst; repeat constructor; auto.
Qed.

Definition remove_lst_ex A n m (lst : list (list A)) (H : cond n m lst) :
  { lst' | refine (remove_lst_spec n m lst) (ret lst') }.
Proof with simpl in *.
  revert m lst H. induction n; intros; inversion H; subst...
  - destruct (@remove_ex A m l H0).
    eexists. rewrite remove_lst_O_ex. unfold remove_lst_O.
    rewrite r. monad_simpl.
    finish honing.
  - apply IHn in X. destruct X.
    exists (l :: x). refines. constructor. auto.
Defined.

Definition remove_lst A n m (lst : list (list A)) (H : cond n m lst) :=
  ` (@remove_lst_ex A n m lst H).

Definition hoge := [[1; 2; 3]; [2; 3]].
Lemma hoge_cond : cond 0 1 hoge.
Proof.
  unfold hoge. repeat constructor.
Defined.

Eval compute in (@remove_lst nat 0 1 hoge hoge_cond).

Definition hoge_cond2 : cond 1 0 hoge.
Proof.
  unfold hoge. repeat constructor.
Defined.

Eval compute in (@remove_lst nat 1 0 hoge hoge_cond2).