(*
TPPmark 2016

(i, j) の組でリスト L に対する入力として許されるもの全体と同型な型
Pos L を定義し， Pos L 型の引数を受け取るようにしました．実際の定義は，
そのような (i, j) の組の個数分の unit の直和になっています．

これだと「整数を受け取る」という問題文の要求に厳密には合っていない気が
しますが，「ちゃんと存在する要素の位置しか引数に取れないようにする」が
問題の本質的な部分と解釈して，このような方針にしてみました．
 *)

Require Import List Arith.
From mathcomp Require Import ssreflect.
Set Implicit Arguments.

(* 一般のリストに関係する識別子には小文字，リストのリストに関係する識
   別子には大文字で始まる名前を使うことにする *)


(* sum P [a1; a2; ...; an] = P a1 + P a2 + ... + P an *)
Definition sum A P : list A -> Type :=
  fun l => List.fold_right (fun a X => (P a + X)%type) Empty_set l.


(* 準備（一般のリストに対する remove） *)
Section List.

(* リストの要素の型 *)
Variable V : Type.

(* リスト中の要素の位置を表すデータの型．実体は要素数個の unit の直和． *)
Definition pos := @sum V (fun _ => unit).

(* リスト l 中の位置 p にある要素を削除する関数 *)
Fixpoint remove l : forall (p : pos l), list V :=
  match l with
  | nil => fun m => match m with end
  | v :: l' =>
    fun p => 
      match p with
      | inl _ => l'
      | inr p' => v :: remove l' p'
      end
  end.

(* remove すると長さが１減る *)
Lemma remove_length_decr l p : length l = S (length (remove l p)).
Proof.
  move : p; elim : l => [[] | v l'].
  move => IH [[] | p'] //.
  by simpl; rewrite (IH p').
Qed.

End List.


(* 本題 *)
Section ListList.

Variable V : Type.

(* リストのリストの中にある要素の位置の型．実体は pos の直和．
   Pos [l1, ..., ln] = pos l1 + ... + pos ln *)
Definition Pos := sum (@pos V).

(* L における位置 P の要素を削除する関数 (問題文の remove-lst) *)
Fixpoint Remove L : forall (P : Pos L), list (list V) :=
  match L with
  | nil => fun P => match P with end
  | l :: L =>
    fun P =>
      match P with
      | inl p => remove l p :: L
      | inr P => l :: Remove L P
      end
  end.

(* 戻り値に関する要求が満たされていることを確かめる *)

(* 外側のリストの長さは不変 *)
Lemma Remove_length_eq L P : length L = length (Remove L P).
Proof.
  move : P; elim : L => [[] | l L'].
  move => IH [p | P'] //.
  by simpl; rewrite (IH P').
Qed.

(* Pos L の値を受け取り，それが i 番目のリストの j 番目を表すなら，i
   を返す（pos L を返してもいいと思うが，自然数の方が簡単だった） *)
Fixpoint index_of (L : list (list V)) : Pos L -> nat :=
  match L with
  | nil => fun P => match P with end
  | l :: L' =>
    fun P =>
      match P with
      | inl _ => 0              (* 最初のリストの中を指している場合 *)
      | inr P' => S (index_of L' P') (* それ以外 *)
      end
  end.

(* 内側のリストの長さは，削除されたものだけ１減る *)
Lemma Remove_length_decr L P n :
  match nth_error L n, nth_error (Remove L P) n with
  | None, None => True          (* n が大きすぎる場合 *)
  | None, Some _                (* このパターンはありえない *)
  | Some _, None => False       (* 同上 *)
  | Some l, Some l' =>
    if beq_nat n (index_of L P)
    then length l = S (length l') (* n 番目のリストの要素が削除された場合 *)
    else length l = length l'     (* 削除されていない場合 *)
  end.
Proof.
  move : L P; elim : n => [ | n'].
  { case => /=; first by case.
    move => l L' [] => [p | P'] //=.
    by rewrite (remove_length_decr l p). }
  { move => IH [ | l L'] [p | P'] //=.
    by case (nth_error L' n').
    apply IH.
  }
Qed.

End ListList.


(* ----------------------------------------------------------------
   おまけ
   de Bruijn index の shift みたいなもの．「問題の背景」にあるようなこ
   とをしようとすると必要な気がした．
 *)
Section Shift.

Variable V : Type.

(* 一般のリスト上の pos に対する shift *)
(* 全単射になっているはず（だが証明していない） *)
Fixpoint shift (l : list V) : forall p, pos l -> option (pos (remove l p)) :=
  match l with
  | nil => fun _ _ => None
  | v :: l' =>
    fun p q =>
      match p, q with
      | inl _, inl _ => None
      | inl _, inr q' => Some q'
      | inr _, inl _ => Some (inl tt)
      | inr p', inr q' =>
        match shift _ p' q' with
        | None => None
        | Some r => Some (inr r)
        end
      end
  end.

(* shift が正しく定義されていそうなことを確かめてみる *)

(* None が返るのは削除された位置に対してのみ *)
Lemma shift_correct1 l p q : shift l p q = None -> p = q.
Proof.
  move : p q; elim : l; first inversion p.
  move => v l' IH [[] | p'] [[] | q'] => //=.
  case eqn : (shift l' p' q') => [] //=.
  by move => _; congr (inr _); apply IH.
Qed.

(* リストの要素を取り出す関数 *)
Fixpoint get l : pos l -> V :=
  match l with
  | nil => fun p => match p with end
  | v :: _ =>
    fun p =>
      match p with
      | inl _ => v
      | inr p' => get _ p'
      end
  end.

(* 位置が返ってきた場合，それは引数と同じ要素を指している *)
Lemma shift_correct l p q r : shift l p q = Some r -> get _ q = get _ r.
Proof.
  move : p q r; elim : l.
  inversion p.
  move => v l' IH [[] | p'] [[] | q'] /=; try congruence.
  by move => r [] <-.
  case eqn : (shift l' p' q') => [ | ] //=.
  move => [[] | r] // [] <-.
  by apply IH.
Qed.

(* リストのリストに対する shift *)
Fixpoint Shift C : forall P : Pos C, Pos C -> option (Pos (Remove C P)) :=
  match C with
  | nil => fun P => match P with end
  | l :: C' =>
    fun P Q =>
      match P, Q with
      | inl p, inl q =>
        match shift l p q with
        | None => None
        | Some r => Some (inl r)
        end
      | inl _, inr Q' => Some (inr Q')
      | inr _, inl q => Some (inl q)
      | inr P', inr Q' =>
        match Shift _ P' Q' with
        | None => None
        | Some R => Some (inr R)
        end
      end
  end.

Lemma Shift_correct1 L P Q : Shift L P Q = None -> P = Q.
Proof.
  move : P Q ; elim : L; first by inversion P.
  case => [ | v l].
  { move => L' IH [[] | P] [[] | Q] /=.
    case eqn : (Shift L' P Q) => [] //.
    by move => _; congr (inr _); apply IH.
  }
  { move => L IH.
    move => [p | P] [q | Q]; try (simpl; congruence).
    { case : p => [[] | ]; case : q => [[] | ] => //=.
      move => p q; case eqn : (shift l q p) => [] //.
      by rewrite (shift_correct1 l q p).
    }
    { simpl; case eqn : (Shift L P Q) => [] //.
      by move => _; congr (inr _); apply IH.
    }
  }
Qed.

Fixpoint Get L : Pos L -> V :=
  match L with
  | nil => fun P => match P with end
  | _ :: _ =>
    fun P =>
      match P with
      | inl p => get _ p
      | inr P' => Get _ P'
      end
  end.

Lemma Shift_correct2 L P Q Q' : Shift L P Q = Some Q' -> Get L Q = Get _ Q'.
Proof.
  move : P Q Q'; elim : L; first by inversion P.
  case => [ | v l].
  { move => L IH [[] | P] [[] | Q] [[] | Q'] /=.
    move => H; apply /IH => /=.
    move : H; case : (Shift L P Q); congruence.
  }
  { move => L IH.
    move => [p | P] [q | Q].
    { (* inl p, inl q *)
      move => [q' | Q'].
      { move => H; apply shift_correct.
        case eqn : (shift _ p q); move : eqn H => /= ->; congruence.
      }
      { by case eqn : (shift _ p q); move : eqn => /= ->. }
    }
    { (* inl p, inr Q *)
      by move => ? [] <-.
    }
    { (* inr P, inl q *)
      by move => ? [] <-.
    }
    { (* inr P, inr Q *)
      case eqn : (Shift _ P Q) => /=.
      { rewrite eqn.
        move => ? [] <- /=.
        by apply IH.
      }
      { by rewrite eqn. }
    }
  }
Qed.

End Shift.