Require Import Lia.
Require Import Coq.Program.Equality.
Require Import PyLevelLang.Language.
Require Import PyLevelLang.Interpret.
Require Import coqutil.Map.Interface coqutil.Map.Properties coqutil.Tactics.Tactics.
Require Import Coq.Lists.List.
From Coq Require Import FunctionalExtensionality.

Local Open Scope Z_scope.

Lemma length_eval_range : forall l s, length (eval_range s l) = l.
Proof.
  induction l; simpl; auto.
Qed.

Section WithMap.
  Context {locals: map.map string {t & interp_type t}} {locals_ok: map.ok locals}.

  Fixpoint listify {t : type} (l : list (expr t)) : expr (TList t) :=
    match l with
    | nil => EAtom (ANil t)
    | x :: xs => EBinop (OCons t) x (listify xs)
    end.

  Lemma listify_correct (t : type) (l : locals) (xs : list (expr t)) :
    interp_expr l (listify xs) = map (interp_expr l) xs.
  Proof.
    induction xs; try easy.
    simpl.
    rewrite IHxs.
    reflexivity.
  Qed.

  Fixpoint reify (t : type) : interp_type t -> expr t :=
    match t in type return interp_type t -> expr t with
    | TInt => fun n => EAtom (AInt n)
    | TBool => fun b => EAtom (ABool b)
    | TString => fun s => EAtom (AString s)
    | TPair s t1 t2 => fun p => EBinop (OPair s t1 t2) (reify t1 (fst p)) (reify t2 (snd p))
    | TEmpty => fun _ => EAtom AEmpty
    | TList t => fun l => listify (map (reify t) l)
    end.

  Lemma reify_correct (t : type) (l : locals) (c : interp_type t) : interp_expr l (reify t c) = c.
  Proof.
    induction t; intros; try easy.
    - destruct c. simpl. rewrite IHt1. rewrite IHt2. reflexivity.
    - destruct c. reflexivity.
    - induction c.
      + reflexivity.
      + simpl. rewrite IHt. simpl in IHc. rewrite IHc. reflexivity.
  Qed.

  Definition invert_atom {t : type} (e : expr t) : option (interp_type t) :=
    match e with
    | EAtom a => Some (interp_atom a)
    | _ => None
    end.

  Lemma invert_atom_correct {t : type} (l : locals) (e : expr t) :
    forall c, invert_atom e = Some c -> interp_expr l e = c.
  Proof.
    intros.
    destruct e; inversion H.
    simpl.
    reflexivity.
  Qed.

  Definition as_const {t : type} (e : expr t) : option (interp_type t) :=
    match e with
    | EAtom a => Some (interp_atom a)
    | EUnop o e1 => match invert_atom e1 with
                    | Some c1 => Some (interp_unop o c1)
                    | _ => None
                    end
    | EBinop o e1 e2 => match invert_atom e1, invert_atom e2 with
                        | Some c1, Some c2 => Some (interp_binop o c1 c2)
                        | _, _ => None
                        end
    | EIf a e1 e2 => match invert_atom a with
                     | Some c1 => if c1 then invert_atom e1 else invert_atom e2
                     | _ => None
                     end
    | _ => None
    end.

  Lemma as_const_correct {t : type} (l : locals) (e : expr t) :
    forall c, as_const e = Some c -> interp_expr l e = c.
  Proof.
    destruct e; simpl; intros; try solve [inversion H]; auto.
    - inversion H. reflexivity.
    - destruct e; try inversion H. reflexivity.
    - destruct e1; destruct e2; try inversion H. reflexivity.
    - dependent induction e1; simpl; simpl in *; try solve [inversion H];
      dependent induction a; simpl; simpl in *; try solve [inversion H];
      destruct b; simpl; apply invert_atom_correct; apply H.
  Qed.

  Definition constfold_head {t} (e : expr t) : expr t :=
    match t with
    | TInt | TBool | TString | TEmpty => match as_const e with
                                         | Some v => reify _ v
                                         | _ => e
                                         end
    | _ => e
    end.

  Lemma constfold_head_correct l {t} e :
    interp_expr l (@constfold_head t e) = interp_expr l e.
  Proof.
    cbv [constfold_head].
    repeat destruct_one_match; auto; try erewrite reify_correct, as_const_correct; auto.
  Qed.

  Fixpoint eLength {t : type} (e : expr (TList t)) : expr TInt :=
    match e with
    | EBinop (OCons t') eh et => EBinop OPlus (EAtom (AInt 1)) (eLength et)
    | EBinop ORange s e => constfold_head (EIf (constfold_head (EBinop OLess e s)) (EAtom (AInt 0)) (constfold_head (EBinop OMinus e s)))
    | _ => match invert_atom e with
           | Some nil => EAtom (AInt 0)
           | _ => EUnop (OLength _) e
           end
    end.

  Lemma eLength_correct {t : type} (l : locals) (e : expr (TList t)) :
    interp_expr l (eLength e) = interp_expr l (EUnop (OLength t) e).
  Proof.
    dependent induction e; cbn [eLength]; try reflexivity.
    - case invert_atom eqn:?; trivial. destruct i; trivial.
      apply invert_atom_correct with (l:=l) in Heqo.
      cbn in *. rewrite Heqo. reflexivity.
    - dependent induction o; cbn [invert_atom]; trivial.
      + cbn [interp_expr interp_binop interp_unop interp_atom Datatypes.length].
        rewrite IHe2; trivial.
        cbn [interp_expr interp_binop interp_unop interp_atom Datatypes.length].
        lia.
      + repeat (rewrite constfold_head_correct; cbn [interp_expr interp_binop interp_unop interp_atom Datatypes.length]).
        destruct_one_match; rewrite length_eval_range; lia.
  Qed.

  Definition invert_EPair {X A B} (e : expr (TPair X A B)) : option (expr A * expr B) :=
    match e in expr t
    return option match t with TPair _ A' B' => expr A' * expr B' | _ => Empty_set end
    with @EBinop _a _b _ab op e1 e2 =>
       match op in binop a b t
       return expr a -> expr b -> option match t with TPair _ A' B' => expr A' * expr B' | _ => Empty_set end
       with OPair s a b => fun e1' e2' => Some (e1', e2') | _ => fun _ _ => None
       end e1 e2
    | _ => None end.

  (* Partial evaluation happens before constant folding *)
  Definition eUnop {t1 t2 : type} (o : unop t1 t2) (e : expr t1) : expr t2 :=
    constfold_head (
    match o in unop t1 t2 return expr t1 -> expr t2 with
    | OLength t => eLength
    | OFst s t1 t2 as o => fun e => match invert_EPair e with Some (x, _) => x | _ => EUnop o e end
    | OSnd s t1 t2 as o => fun e => match invert_EPair e with Some (_, x) => x | _ => EUnop o e end
    | o => EUnop o
    end e).

  Lemma invert_EPair_correct {X A B} e e1 e2 :
    @invert_EPair X A B e = Some (e1, e2)
    <-> e = EBinop (OPair X A B) e1 e2.
  Proof.
    clear dependent locals.
    dependent induction e; cbn; intuition; try congruence;
    dependent induction o; inversion H;
    try apply Eqdep.EqdepTheory.inj_pair2 in H1, H2; try exact type_eq_dec;
    congruence.
  Qed.

  Lemma EUnop_correct {t1 t2 : type} (l : locals) (o : unop t1 t2) (e : expr t1)
    : interp_expr l (eUnop o e) = interp_expr l (EUnop o e).
  Proof.
    cbv [eUnop]; rewrite constfold_head_correct; case o in *; trivial.
    { apply eLength_correct. }
    { destruct invert_EPair as [[]|] eqn:H; try apply invert_EPair_correct in H; rewrite ?H; trivial. }
    { destruct invert_EPair as [[]|] eqn:H; try apply invert_EPair_correct in H; rewrite ?H; trivial. }
  Qed.

  Definition partial_head {t} (e : expr t) : expr t :=
    match e with
    | EUnop o e1 => eUnop o e1
    | e => e
    end.

  Section fold_expr.
    Context (f : forall {t}, expr t -> expr t). Local Arguments f {_}.
    Fixpoint fold_expr {t : type} (e : expr t) : expr t :=
      f
      match e in expr t return expr t with
      | EUnop o e1 => EUnop o (fold_expr e1)
      | EBinop o e1 e2 => EBinop o (fold_expr e1) (fold_expr e2)
      | EFlatmap e1 x e2 => EFlatmap (fold_expr e1) x (fold_expr e2)
      | EFold e1 e2 x y e3 => EFold (fold_expr e1) (fold_expr e2) x y (fold_expr e3)
      | EIf e1 e2 e3 => EIf (fold_expr e1) (fold_expr e2) (fold_expr e3)
      | ELet x e1 e2 => ELet x (fold_expr e1) (fold_expr e2)
      | (EVar _ _ | ELoc _ _ | EAtom _) as e => e
      end.
  End fold_expr.

  Definition partial {t : type} := @fold_expr (@partial_head) t.

  Fixpoint constant_folding' {t : type} (e : expr t) : expr t * option (interp_type t) :=
    match e with
    | EVar _ x => (EVar _ x, None)
    | ELoc _ x => (ELoc _ x, None)
    | EAtom a => (EAtom a, Some (interp_atom a))
    | EUnop o e1 =>
        let (e1', v) := constant_folding' e1 in
        match v with
        | Some v' => let r := interp_unop o v' in (reify _ r, Some r)
        | _ => (EUnop o e1', None)
        end
    | EBinop o e1 e2 =>
        let (e1', v1) := constant_folding' e1 in
        let (e2', v2) := constant_folding' e2 in
        match v1, v2 with
        | Some v1', Some v2' => let r := interp_binop o v1' v2' in (reify _ r, Some r)
        | _, _ => (EBinop o e1' e2', None)
        end
    | EFlatmap e1 x e2 => (EFlatmap (fst (constant_folding' e1)) x (fst (constant_folding' e2)), None)
    | EFold e1 e2 x y e3 => (EFold (fst(constant_folding' e1)) (fst(constant_folding' e2)) x y (fst(constant_folding' e3)), None)
    | EIf e1 e2 e3 => (EIf (fst (constant_folding' e1)) (fst (constant_folding' e2)) (fst (constant_folding' e3)), None)
    | ELet x e1 e2 => (ELet x (fst (constant_folding' e1)) (fst (constant_folding' e2)), None)
    end.

  Lemma constant_folding'_snd_correct {t : type} (l : locals) (e : expr t) :
    forall c, snd (constant_folding' e) = Some c -> interp_expr l e = c.
  Proof.
    generalize dependent l.
    induction e; intros; simpl; try easy.

    - inversion H. reflexivity.

    - simpl in H.
      destruct (constant_folding' e). destruct o0; inversion H.
      rewrite (IHe l i); easy.

    - simpl in H.
      destruct (constant_folding' e1). destruct (constant_folding' e2). destruct o0; destruct o1; inversion H.
      rewrite (IHe1 l i), (IHe2 l i0); easy.
  Qed.

  Definition constant_folding {t : type} (e : expr t) : expr t := fst (constant_folding' e).

  Lemma constant_folding_correct {t : type} (l : locals) (e : expr t) :
    interp_expr l (constant_folding e) = interp_expr l e.
  Proof.
    generalize dependent l.
    unfold constant_folding.
    induction e; intros; simpl; try easy.

    - destruct (constant_folding' e) eqn:E. destruct o0; simpl.
      + rewrite reify_correct.
        f_equal.
        simpl in IHe.
        rewrite (constant_folding'_snd_correct l e i); try easy.
        rewrite E. reflexivity.
      + simpl in IHe. rewrite IHe. reflexivity.

    - destruct (constant_folding' e1) eqn:E1. destruct (constant_folding' e2) eqn:E2. destruct o0; destruct o1; simpl.
      + rewrite reify_correct.
        f_equal.
        * simpl in IHe1.
          rewrite (constant_folding'_snd_correct l e1 i); try easy. rewrite E1. reflexivity.
        * simpl in IHe2.
          rewrite (constant_folding'_snd_correct l e2 i0); try easy. rewrite E2. reflexivity.

      + rewrite IHe1, IHe2. reflexivity.
      + rewrite IHe1, IHe2. reflexivity.
      + rewrite IHe1, IHe2. reflexivity.

    - rewrite IHe1.
      assert (H: (fun y : interp_type t1 => interp_expr (set_local l x y) (fst (constant_folding' e2)))
      = (fun y : interp_type t1 => interp_expr (set_local l x y) e2)).
      { apply functional_extensionality. intros. apply IHe2. }
      rewrite <- H.
      reflexivity.

    - rewrite IHe1.
      f_equal; eauto.
      apply functional_extensionality. intros.
      apply functional_extensionality. intros.
      easy.

    - rewrite IHe1, IHe2, IHe3. reflexivity.

    - rewrite IHe1, IHe2. reflexivity.
  Qed.

  Fixpoint branch_elim {t : type} (e : expr t) : expr t :=
    match e in expr t' return expr t' with
    | EVar _ x => EVar _ x
    | ELoc _ x => ELoc _ x
    | EAtom a => EAtom a
    | EUnop o e1 => EUnop o (branch_elim e1)
    | EBinop o e1 e2 => EBinop o (branch_elim e1) (branch_elim e2)
    | EIf e1 e2 e3 =>
        let e1' := branch_elim e1 in
        let e2' := branch_elim e2 in
        let e3' := branch_elim e3 in
        match as_const e1' with
        | Some true => e2'
        | Some false => e3'
        | _ => EIf e1' e2' e3'
        end
    | EFlatmap e1 x e2 => EFlatmap (branch_elim e1) x (branch_elim e2)
    | EFold e1 e2 x y e3 => EFold (branch_elim e1) (branch_elim e2) x y (branch_elim e3)
    | ELet x e1 e2 => ELet x (branch_elim e1) (branch_elim e2)
    end.

  Lemma branch_elim_correct {t : type} (l : locals) (e : expr t)
    : interp_expr l (branch_elim e) = interp_expr l e.
  Proof.
    generalize dependent l.
    induction e; simpl; intros;
    try rewrite <- IHe; try rewrite <- IHe1; try rewrite <- IHe2; try rewrite <- IHe3;
    try reflexivity.

    - assert (H:(fun y : interp_type t1 => interp_expr (set_local l x y) (branch_elim e2))
              = (fun y => interp_expr (set_local l x y) e2)).
      { apply functional_extensionality. intros. rewrite IHe2. reflexivity. }
      rewrite H. reflexivity.

    - f_equal.
      do 2 (apply functional_extensionality; intros).
      easy.

    - destruct (as_const (branch_elim e1)) eqn:H; try easy.
      destruct i; rewrite (as_const_correct l (branch_elim e1) _ H); reflexivity.
  Qed.

  Fixpoint is_name_used {t : type} (x : string) (e : expr t) : bool :=
    match e with
    | EVar _ x' => eqb x' x
    | ELoc _ x' => eqb x' x
    | EAtom _ => false
    | EUnop _ e1 => is_name_used x e1
    | EBinop _ e1 e2 => is_name_used x e1 || is_name_used x e2
    | EFlatmap e1 x' e2 => eqb x' x || is_name_used x e1 || is_name_used x e2
    | EFold e1 e2 x' y' e3 => eqb x' x || eqb y' x || is_name_used x e1 || is_name_used x e2 || is_name_used x e3
    | EIf e1 e2 e3 => is_name_used x e1 || is_name_used x e2 || is_name_used x e3
    | ELet x' e1 e2 => eqb x' x || is_name_used x e1 || is_name_used x e2
    end.

  Lemma set_local_same {tx ty : type} (x : string) (vx : interp_type tx) (vy : interp_type ty) (l : locals)
    : set_local (set_local l x vx) x vy = set_local l x vy.
  Proof.
    unfold set_local.
    apply map.map_ext.
    intros.
    rewrite map.get_put_dec.
    destruct (eqb x k) eqn:E.
    - rewrite eqb_eq in E. rewrite E. rewrite map.get_put_same. reflexivity.
    - rewrite eqb_neq in E. rewrite map.get_put_diff; auto.
      rewrite map.get_put_diff; auto.
  Qed.

  Lemma set_local_comm_diff {tx ty : type} (x y : string) (vx : interp_type tx) (vy : interp_type ty) (l : locals)
    : x <> y -> set_local (set_local l y vy) x vx = set_local (set_local l x vx) y vy.
  Proof.
    unfold set_local.
    intros.
    apply map.map_ext.
    intros.
    rewrite 4 map.get_put_dec.
    repeat destruct_one_match; try rewrite map.get_put_diff; tauto.
  Qed.

  Lemma is_name_used_correct {t : type} {t' : type} (l : locals) (e : expr t) (x : string):
    forall y : interp_type t', is_name_used x e = false -> interp_expr l e = interp_expr (set_local l x y) e.
  Proof.
    generalize dependent l.
    induction e; intros; simpl in H; simpl.

    - unfold get_local, set_local.
      rewrite map.get_put_diff; try easy.
      apply eqb_neq, H.

    - unfold get_local, set_local.
      rewrite map.get_put_diff; try easy.
      apply eqb_neq, H.

    - reflexivity.

    - rewrite <- IHe; easy.

    - rewrite <- IHe1, <- IHe2; destruct (is_name_used x e2); destruct (is_name_used x e1); try easy.

    - destruct (is_name_used x e1); destruct (is_name_used x e2); destruct (eqb x0 x) eqn:Hx; simpl in H; try easy.
      assert (Hf:(fun y0 : interp_type t1 => interp_expr (set_local l x0 y0) e2)
              = (fun y0 : interp_type t1 => interp_expr (set_local (set_local l x y) x0 y0) e2)).
              { apply functional_extensionality. intros. rewrite set_local_comm_diff.
                + apply IHe2. reflexivity.
                + apply eqb_neq, Hx.
              }
              rewrite Hf. rewrite <- IHe1; easy.

    - destruct (eqb x0 x) eqn:H0x;
      destruct (eqb y x) eqn:Hyx;
      destruct (is_name_used x e1) eqn:He1;
      destruct (is_name_used x e2) eqn:He2;
      destruct (is_name_used x e3) eqn:He3;
      try discriminate H.
      f_equal; eauto.
      do 2 (apply functional_extensionality; intros).
      rewrite <- set_local_comm_diff with (x := x); cycle 1.
      { apply eqb_neq. rewrite eqb_sym. assumption. }
      rewrite <- set_local_comm_diff with (x := x); cycle 1.
      { apply eqb_neq. rewrite eqb_sym. assumption. }
      apply IHe3. reflexivity.

    - rewrite <- IHe1, <- IHe2, <- IHe3;
      destruct (is_name_used x e1);
      destruct (is_name_used x e2);
      destruct (is_name_used x e3);
      easy.

    - rewrite set_local_comm_diff;
      destruct (is_name_used x e1);
      destruct (is_name_used x e2);
      destruct (eqb x0 x) eqn:Hx;
      try easy.
      + rewrite <- IHe1, <- IHe2; easy.
      + apply eqb_neq, Hx.
  Qed.

  Fixpoint unused_name_elim {t : type} (e : expr t) : expr t :=
    match e with
    | EVar _ x => EVar _ x
    | ELoc _ x => ELoc _ x
    | EAtom a => EAtom a
    | EUnop o e1 => EUnop o (unused_name_elim e1)
    | EBinop o e1 e2 => EBinop o (unused_name_elim e1) (unused_name_elim e2)
    | EFlatmap e1 x e2 => EFlatmap (unused_name_elim e1) x (unused_name_elim e2)
    | EFold e1 e2 x y e3 => EFold (unused_name_elim e1) (unused_name_elim e2) x y (unused_name_elim e3)
    | EIf e1 e2 e3 => EIf (unused_name_elim e1) (unused_name_elim e2) (unused_name_elim e3)
    | ELet x e1 e2 => if is_name_used x e2 then ELet x (unused_name_elim e1) (unused_name_elim e2) else unused_name_elim e2
    end.

  Lemma unused_name_elim_correct {t : type} (l : locals) (e : expr t)
    : interp_expr l (unused_name_elim e) = interp_expr l e.
  Proof.
    generalize dependent l.
    induction e; try easy; intros; simpl.

    - rewrite IHe. reflexivity.

    - rewrite IHe1, IHe2. reflexivity.

    - assert (H:(fun y : interp_type t1 => interp_expr (set_local l x y) (unused_name_elim e2))
              = (fun y => interp_expr (set_local l x y) e2)).
      { apply functional_extensionality. intros. rewrite IHe2. reflexivity. }
      rewrite IHe1, H. reflexivity.

    - f_equal; eauto.
      do 2 (apply functional_extensionality; intros).
      eauto.

    - rewrite IHe1, IHe2, IHe3. reflexivity.

    - destruct (is_name_used x e2) eqn:E.
      + simpl. rewrite IHe1, IHe2. reflexivity.
      + rewrite <- is_name_used_correct; easy.
  Qed.

  Definition flatmap_flatmap_head {t : type} (e : expr t) : expr t :=
    match e with
    | EFlatmap l x f => match l with
                        | EFlatmap l' y g => if orb (is_name_used y f) (eqb x y)
                                              then EFlatmap l x f
                                              else EFlatmap l' y (EFlatmap g x f)
                        | l' => EFlatmap l' x f
                        end
    | e' => e'
    end.

  Lemma flat_map_flat_map :
    forall {A B C} (l : list A) (f : B -> list C)  (g : A -> list B),
    flat_map f (flat_map g l) = flat_map (fun x => flat_map f (g x)) l.
  Proof.
    induction l; auto.
    intros. simpl. rewrite flat_map_app. rewrite IHl. reflexivity.
  Qed.

  Lemma flatmap_flatmap_head_correct {t : type} (l : locals) (e : expr t)
    : interp_expr l (flatmap_flatmap_head e) = interp_expr l e.
  Proof.
    destruct e; auto.
    dependent induction e1; auto; simpl.
    destruct_one_match; auto;
    simpl. rewrite flat_map_flat_map. f_equal. apply functional_extensionality.
    intros. f_equal. apply functional_extensionality.
    intros. symmetry. rewrite set_local_comm_diff.
    + apply is_name_used_correct. apply E.
    + apply E.
  Qed.

  Definition flatmap_flatmap {t} : expr t -> expr t := fold_expr (@flatmap_flatmap_head).

  Goal forall s x, eUnop (OLength _) (eUnop (OFst s _ _)
    (EBinop (OPair s _ _)
      (EBinop ORange (EAtom (AInt 1)) (EAtom (AInt 5)))
      (EVar TInt x)))
  = EAtom (AInt 4).
  Proof. cbv. solve [trivial]. Abort.

End WithMap.


(** * Annotative Probabilistic Models *)
Require Import Reals.
Require Import List.
Require Import ListSet.

Open Scope R_scope.

Require Import Probabilities.
Require Import State.
Require Import RatExpr.
Require Import TransitionMatrix.Definitions.
Require Import TransitionMatrix.Real.
Require Import TransitionMatrix.Parametric.
Require Import Markov.DTMC.
Require Import Markov.PMC.
Require Import ReliabilityModels.Definitions.

(** ** Annotative PMC **)

Definition annotative_state (m: TransitionMatrix RatExpr) (s: State): Prop :=
    forall r: StateMaps.t RatExpr, StateMaps.MapsTo s r m ->
        (expr_is_stochastic_row r
        \/
        (exists s1 s2: State, exists x: Var,
            s1 <> s2
            /\ StateMaps.MapsTo s1 (OneVar x) r
            /\ StateMaps.MapsTo s2 (Sub (Const 1) (OneVar x)) r
            /\ forall s': State,
                StateMaps.In s' (StateMaps.remove s2 (StateMaps.remove s1 r)) ->
                  StateMaps.MapsTo s' (Const 0) (StateMaps.remove s2 (StateMaps.remove s1 r)) )).

Definition annotative_pmc (p: PMC) : Prop :=
    forall s: State, In s p.(S) -> annotative_state p.(P) s.

Lemma eval_annotative_state_is_stochastic:
    forall (m: TransitionMatrix RatExpr) (s: State) (r: StateMaps.t RatExpr) (u: Evaluation),
        StateMaps.MapsTo s r m ->
        annotative_state m s ->
        stochastic_evaluation u (var_set m)
        -> is_stochastic_row (eval_row r u).
Proof.
    intros m s r u H_s_mapsto_m H_s_annot H_u_stochastic.
    unfold annotative_state in H_s_annot.
    assert (H_s_mapsto_m':= H_s_mapsto_m).
    apply H_s_annot in H_s_mapsto_m.
    destruct H_s_mapsto_m.
    - (* all transitions are constant values *)
      apply eval_stochastic_row. apply H.
    - (* transitions are x and 1-x *)
      destruct H as [s1 [s2 [x [Hs1_s2 [Hx [Hx_compl H]]]]]].
      clear H_s_annot.
      unfold is_stochastic_row. split.
      + (* expressions evaluate to valid probabilities *)
        intros s' v H_s'_value.
        unfold stochastic_evaluation in H_u_stochastic.
        assert (H_x_in_r: In x (var_set' r)).
        { unfold var_set'.
          apply StateMaps.elements_1 in Hx. elim Hx.
          - (* x is in the first element. *)
            intros.
            unfold var_set_list'.
            repeat red in H0. destruct H0 as [H_fst H_snd].
            rewrite <- H_snd. simpl.
            apply set_union_intro1. simpl. left. reflexivity.
          - (* x is somewhere in the remaining elements. *)
            intros.
            unfold var_set_list'.
            apply set_union_intro2. assumption.
        }
        apply vars_injection with (m:=m) (s:=s) in H_x_in_r.
          Focus 2. assumption.
        apply H_u_stochastic in H_x_in_r.
        apply commutative_eval_mapsto in H_s'_value.
        elim H_s'_value. intros e [H_s'_mapsto_e H_v_eval_e].
        { (* Let us perform case analysis of s': *)
          destruct (State.eq_dec s1 s').
          - (* s' == s1 *)
            apply StateMaps.MapsTo_1 with (x:=s1) (y:=s') in Hx.
              Focus 2. auto.
            apply StateMapsFacts.MapsTo_fun with (e:=OneVar x) in H_s'_mapsto_e.
              Focus 2. assumption.
            rewrite H_v_eval_e.
            rewrite <- H_s'_mapsto_e.
            simpl. assumption.
          - (* s' <> s1 *)
            destruct (State.eq_dec s2 s').
            + (* s' == s2 *)
              apply StateMaps.MapsTo_1 with (x:=s2) (y:=s') in Hx_compl.
                Focus 2. auto.
              apply StateMapsFacts.MapsTo_fun with (e:=(Sub (Const 1) (OneVar x))) in H_s'_mapsto_e.
                Focus 2. assumption.
              rewrite H_v_eval_e.
              rewrite <- H_s'_mapsto_e.
              simpl.
              apply valid_prob_1_minus_p. assumption.
            + (* s' <> s2 *)
              apply StateMaps.remove_2 with (x:=s1) in H_s'_mapsto_e.
                Focus 2. assumption.
              apply StateMaps.remove_2 with (x:=s2) in H_s'_mapsto_e.
                Focus 2. assumption.
              apply StateMapsFacts.MapsTo_fun with (e:=Const 0) in H_s'_mapsto_e.
                Focus 2. apply H. apply mapsto_in in H_s'_mapsto_e. assumption.
              rewrite H_v_eval_e.
              rewrite <- H_s'_mapsto_e. simpl.
              unfold is_valid_prob. split.
              * apply Rle_refl.
              * apply Rle_0_1.
        }
      + (* sum of row is 1 *)
        rewrite <- eval_sum_in_map.
        rewrite sum_in_map with (s:=s1) (e:=OneVar x).
        Focus 2. apply Hx.
        rewrite sum_in_map with (s:=s2) (e:=Sub (Const 1) (OneVar x)).
        Focus 2. apply StateMaps.remove_2 with (x:=s1) in Hx_compl.
        Focus 2. apply Hs1_s2.
        apply Hx_compl.
        remember (StateMaps.remove s2 (StateMaps.remove s1 r)) as rest.
        apply sum_zero_in_map in H.
        rewrite H. simpl. remember (u x) as a.
        rewrite Rplus_0_r. rewrite Rplus_minus. reflexivity.
Qed.

Definition AnnotativeEvaluationFactory (pf: PresenceFunction): EvaluationFactory :=
    fun (c: SPL.Configuration) (x: Var) => if pf x c then 1 else 0.

(** Definition 6 (Annotative probabilistic model) *)
Record AnnotativeProbModel : Type := { pmc: PMC;
                                       pf: PresenceFunction;
                                       w: EvaluationFactory;
                                       FM: SPL.FM }.

Notation "[| fm |]" := (SPL.fm_semantics fm).
Definition wf_annot_prob_model (apm: AnnotativeProbModel) : Prop :=
    wf_PMC apm.(pmc)
    /\ annotative_pmc apm.(pmc)
    /\ ( forall (c: SPL.Configuration) (x: Var),
           (In c [| apm.(FM) |]) -> (In x apm.(pmc).(X))
           ->  apm.(w) = AnnotativeEvaluationFactory apm.(pf) ).

(** Lemma 2 (Evaluation well-definedness for annotative models) *)
Lemma evaluation_well_definedness:
    forall apm: AnnotativeProbModel, wf_annot_prob_model apm ->
        forall c: SPL.Configuration,
            In c [|apm.(FM)|] -> well_defined_evaluation apm.(pmc) (apm.(w) c).
Proof.
    intros.
    destruct H as [[Hs0 [HTinS [HSinP HXinP]]] [HannotPMC HannotEvalFactory]].
    unfold well_defined_evaluation.
    remember (eval_pmc (pmc apm) (w apm c)) as d.
    unfold wf_DTMC.
    intuition. (* Break each proposition in the conjunction into a separate goal. *)
    - (* s0 in S *)
      rewrite -> Heqd. simpl. apply Hs0.
    - (* T in S *)
      rewrite -> Heqd. simpl. apply HTinS.
    - (* StateMaps.In s (P d) *)
      rewrite Heqd in H. unfold eval_pmc in H. unfold S in H.
      rewrite Heqd. unfold eval_pmc. unfold P.
      apply HSinP in H. apply eval_matrix_in. apply H.
    - (* In s (S d) *)
      rewrite Heqd in H. unfold eval_pmc in H. unfold P in H.
      rewrite Heqd. unfold eval_pmc. unfold S.
      apply HSinP. apply eval_matrix_in in H. apply H.
    - (* all transitions are valid probabilities *)
      unfold is_stochastic_matrix. intros s r H_mapsto.
      rewrite Heqd in H_mapsto. unfold eval_pmc in H_mapsto. unfold P in H_mapsto.
      apply eval_matrix_evaluated_expr in H_mapsto. elim H_mapsto.
      intros e [H_mapsto_e H_r]. rewrite H_r.
      apply eval_annotative_state_is_stochastic with (m:=P (pmc apm)) (s:=s).
      + (* MapsTo s e (P apm) ? *)
        assumption.
      + (* annotative_state ? *)
        unfold annotative_pmc in HannotPMC.
        apply mapsto_in in H_mapsto_e.
        rewrite <- HSinP in H_mapsto_e. apply HannotPMC in H_mapsto_e.
        assumption.
      + (* w is a stochastic evaluation? *)
        unfold stochastic_evaluation. intros.
        apply HannotEvalFactory with (x:=x) in H0.
        { (* is valid probability? *)
          rewrite H0.
          unfold AnnotativeEvaluationFactory.
          induction (pf apm x c).
          - unfold is_valid_prob. split.
              + apply Rle_0_1.
              + apply Rle_refl.
          - unfold is_valid_prob. split.
              + apply Rle_refl.
              + apply Rle_0_1.
        }
        { (* is x in X? *)
          rewrite HXinP. assumption.
        }
Qed.


(** **********************************************************

Benedikt Ahrens, Ralph Matthes

SubstitutionSystems

2015

************************************************************)

(** **********************************************************

Contents :

-    Definition of heterogeneous substitution systems
-    Various lemmas about the substitution ("bracket") operation
-    Definition of precategory of substitution systems

simplified notion of HSS by Ralph Matthes (2022, 2023)
the file is very close to the homonymous file in the parent directory
basically, the changes start with [bracket_property] and are then propagated throughout

************************************************************)

Require Import UniMath.Foundations.PartD.

Require Import UniMath.CategoryTheory.Core.Categories.
Require Import UniMath.CategoryTheory.Core.Functors.
Require Import UniMath.CategoryTheory.Core.NaturalTransformations.
Require Import UniMath.CategoryTheory.FunctorCategory.
Local Open Scope cat.
Require Import UniMath.CategoryTheory.whiskering.
Require Import UniMath.CategoryTheory.FunctorAlgebras.
Require Import UniMath.CategoryTheory.limits.bincoproducts.
Require Import UniMath.CategoryTheory.PointedFunctors.
Require Import UniMath.CategoryTheory.PrecategoryBinProduct.
Require Import UniMath.CategoryTheory.HorizontalComposition.
Require Import UniMath.CategoryTheory.PointedFunctorsComposition.
Require Import UniMath.SubstitutionSystems.Signatures.
Require Import UniMath.SubstitutionSystems.Notation.
Local Open Scope subsys.

Section fix_a_category.

(** ** Some variables and assumptions *)

Context (C : category) (CP : BinCoproducts C).

Local Notation "'EndC'":= ([C, C]) .
Let CPEndC : BinCoproducts EndC := BinCoproducts_functor_precat _ _ CP.

(** The category of pointed endofunctors on [C] *)
Local Notation "'Ptd'" := (category_Ptd C).
(** The category of endofunctors on [C] *)
Local Notation "'EndC'":= ([C, C]) .

(** preparations for the definition of hss based on a functor [H] *)
Section prep_hss.

Context (H : functor [C, C] [C, C]).

Definition Id_H
: functor EndC EndC
  := BinCoproduct_of_functors _ _ CPEndC
                       (constant_functor _ _ (functor_identity _ : EndC))
                       H.

(* An Id_H algebra is a pointed functor *)

Definition eta_from_alg (T : algebra_ob Id_H) : EndC ⟦ functor_identity _,  `T ⟧.
Proof.
  exact (BinCoproductIn1 (CPEndC _ _) · alg_map _ T).
Defined.

Local Notation η := eta_from_alg.

Definition ptd_from_alg (T : algebra_ob Id_H) : Ptd.
Proof.
  exists (pr1 T).
  exact (η T).
Defined.

Definition tau_from_alg (T : algebra_ob Id_H) : EndC ⟦H `T, `T⟧.
Proof.
  exact (BinCoproductIn2 (CPEndC _ _) · alg_map _ T).
Defined.
Local Notation τ := tau_from_alg.

(*
Coercion functor_from_algebra_ob (X : algebra_ob _ Id_H) : functor C C  := pr1 X.
*)

Local Notation "f ⊕ g" := (BinCoproductOfArrows _ (CPEndC _ _ ) (CPEndC _ _ ) f g).

(** analysis of the "bracket operation" individually for each pointed functor *)
Section fix_a_pointed_functor.

  Context {Z: Ptd}.
  Context (θ : @PrestrengthForSignatureAtPoint C C C H Z).

Definition bracket_property (T : algebra_ob Id_H) (f : U Z --> `T)
           (h : `T • (U Z)  --> `T) : UU
  :=
    alg_map _ T •• (U Z) · h =
          (identity (U Z) ⊕ θ `T) ·
          (identity (U Z) ⊕ #H h) ·
          (BinCoproductArrow  (CPEndC _ _ ) f (tau_from_alg T)).

Definition bracket_at (T : algebra_ob Id_H) (f : U Z --> `T): UU :=
  ∃! h : `T • (U Z)  --> `T, bracket_property T f h.

Definition bracket_property_parts (T : algebra_ob Id_H) (f : U Z --> `T)
           (h : `T • (U Z)  --> `T) : UU
  :=
    (f = η T •• (U Z) · h) ×
     (θ `T · #H h · τ T  = τ T •• (U Z) ·  h).

Definition bracket_parts_at (T : algebra_ob Id_H) (f : U Z --> `T) : UU :=
   ∃! h : `T • (U Z)  --> `T, bracket_property_parts T f h.

(* show that for any h of suitable type, the following are equivalent *)

Lemma parts_from_whole (T : algebra_ob Id_H) (f : U Z --> `T)
      (h :  `T • (U Z)  --> `T) :
  bracket_property T f h → bracket_property_parts T f h.
Proof.
  intro Hyp.
  split.
  + unfold eta_from_alg.
    apply nat_trans_eq_alt.
    intro c.
    simpl.
    unfold coproduct_nat_trans_in1_data.
    assert (Hyp_inst := nat_trans_eq_pointwise Hyp c); clear Hyp.
    apply (maponpaths (λ m, BinCoproductIn1 (CP _ _)· m)) in Hyp_inst.
    match goal with |[ H1 : _  = ?f |- _ = _   ] =>
         intermediate_path (f) end.
    * clear Hyp_inst.
      rewrite <- assoc.
      apply BinCoproductIn1Commutes_right_in_ctx_dir.
      rewrite id_left.
      apply BinCoproductIn1Commutes_right_in_ctx_dir.
      rewrite id_left.
      apply BinCoproductIn1Commutes_right_dir.
      apply idpath.
    * rewrite <- Hyp_inst; clear Hyp_inst.
      rewrite <- assoc.
      apply idpath.
  + unfold tau_from_alg.
    apply nat_trans_eq_alt.
    intro c.
    simpl.
    unfold coproduct_nat_trans_in2_data.
    assert (Hyp_inst := nat_trans_eq_pointwise Hyp c); clear Hyp.
    apply (maponpaths (λ m,  BinCoproductIn2 (CP _ _)· m)) in Hyp_inst.
    match goal with |[ H1 : _  = ?f |- _ = _   ] =>
         intermediate_path (f) end.
    * clear Hyp_inst.
      do 2 rewrite <- assoc.
      apply BinCoproductIn2Commutes_right_in_ctx_dir.
      simpl.
      rewrite <- assoc.
      apply maponpaths.
      apply BinCoproductIn2Commutes_right_in_ctx_dir.
      simpl.
      rewrite <- assoc.
      apply maponpaths.
      unfold tau_from_alg.
      apply BinCoproductIn2Commutes_right_dir.
      apply idpath.
    * rewrite <- Hyp_inst; clear Hyp_inst.
      rewrite <- assoc.
      apply idpath.
Qed.

Lemma whole_from_parts (T : algebra_ob Id_H) (f : U Z --> `T)
      (h :  `T • (U Z)  --> `T) :
  bracket_property_parts T f h → bracket_property T f h.
Proof.
  intros [Hyp1 Hyp2].
  apply nat_trans_eq_alt.
  intro c.
  apply BinCoproductArrow_eq_cor.
  + clear Hyp2.
    assert (Hyp1_inst := nat_trans_eq_pointwise Hyp1 c); clear Hyp1.
    rewrite <- assoc.
    apply BinCoproductIn1Commutes_right_in_ctx_dir.
    rewrite id_left.
    apply BinCoproductIn1Commutes_right_in_ctx_dir.
    rewrite id_left.
    apply BinCoproductIn1Commutes_right_dir.
    simpl. simpl in Hyp1_inst.
    rewrite Hyp1_inst.
    simpl.
    apply assoc.
  + clear Hyp1.
    assert (Hyp2_inst := nat_trans_eq_pointwise Hyp2 c); clear Hyp2.
    rewrite <- assoc.
    apply BinCoproductIn2Commutes_right_in_ctx_dir.
    simpl.
    rewrite assoc.
    eapply pathscomp0.
    * eapply pathsinv0.
      exact Hyp2_inst.
    * clear Hyp2_inst.
      simpl.
      do 2 rewrite <- assoc.
      apply maponpaths.
      apply BinCoproductIn2Commutes_right_in_ctx_dir.
      simpl.
      rewrite <- assoc.
      apply maponpaths.
      apply BinCoproductIn2Commutes_right_dir.
      apply idpath.
Qed.

(* show bracket_parts_point is logically equivalent to bracket_point, then
   use it to show that bracket_parts is equivalent to bracket using [weqonsecfibers:
  ∏ (X : UU) (P Q : X → UU),
  (∏ x : X, P x ≃ Q x) → (∏ x : X, P x) ≃ (∏ x : X, Q x)] *)

End fix_a_pointed_functor.

Section instantiate_with_identity.

    Context (T : algebra_ob Id_H).
    Context (θ : @PrestrengthForSignatureAtPoint C C C H (ptd_from_alg T)).

Definition bracket_property_parts_identity_nicer (h : `T • `T  --> `T) : UU
  := (identity `T = η T •• `T · h) × (θ `T · #H h · τ T  = τ T •• `T ·  h).

Lemma bracket_property_parts_identity_nicer_impl1 (h : `T • `T  --> `T):
  bracket_property_parts θ T (identity _) h -> bracket_property_parts_identity_nicer h.
Proof.
  intro Hyp. induction Hyp as [Hyp1 Hyp2].
  split.
  - etrans.
    2: { exact Hyp1. }
    apply nat_trans_eq_alt.
    intro c.
    apply idpath.
  - etrans.
    { exact Hyp2. }
    apply idpath.
Qed.

(** basically the same proof *)
Lemma bracket_property_parts_identity_nicer_impl2 (h : `T • `T  --> `T):
  bracket_property_parts_identity_nicer h -> bracket_property_parts θ T (identity _) h.
  intro Hyp. induction Hyp as [Hyp1 Hyp2].
  split.
  - etrans.
    2: { exact Hyp1. }
    apply nat_trans_eq_alt.
    intro c.
    apply idpath.
  - etrans.
    { exact Hyp2. }
    apply idpath.
Qed.

End instantiate_with_identity.

(** the notion one would be looking for: an algebra and a substitution operation that does lookup on
variables and behaves homomorphically elsewhere, as instructed by the pre-strength at that point *)
Definition heterogeneous_substitution : UU := ∑ (T: algebra_ob Id_H),
  ∑ (θ : @PrestrengthForSignatureAtPoint C C C H (ptd_from_alg T)), bracket_at θ T (identity _).

Coercion alg_from_hetsubst (T : heterogeneous_substitution) : algebra_ob Id_H := pr1 T.
Definition θ_from_hetsubst (T : heterogeneous_substitution) : @PrestrengthForSignatureAtPoint C C C H (ptd_from_alg T)
  := pr1 (pr2 T).

(** we write prejoin as a warning that the monad laws are not necessarily fulfilled *)
Definition prejoin_from_hetsubst (T : heterogeneous_substitution) : `T • `T --> `T
  := pr1 (pr1 (pr2 (pr2 T))).

Lemma prejoin_from_hetsubst_η (T : heterogeneous_substitution) :
  identity `T = η T •• `T · (prejoin_from_hetsubst T).
Proof.
  refine (pr1 (bracket_property_parts_identity_nicer_impl1 T (θ_from_hetsubst T) _ _)).
  apply parts_from_whole.
  exact (pr2 (pr1 (pr2 (pr2 T)))).
Qed.

Lemma prejoin_from_hetsubst_τ (T : heterogeneous_substitution) :
  θ_from_hetsubst T `T · #H (prejoin_from_hetsubst T) · τ T  = τ T •• `T ·  (prejoin_from_hetsubst T).
Proof.
  refine (pr2 (bracket_property_parts_identity_nicer_impl1 T (θ_from_hetsubst T) _ _)).
  apply parts_from_whole.
  exact (pr2 (pr1 (pr2 (pr2 T)))).
Qed.

Section fix_a_prestrength.

  Context (θ :  @PrestrengthForSignature C C C H).

Definition bracket (T : algebra_ob Id_H) : UU
  := ∏ (Z : Ptd) (f : U Z --> `T), bracket_at (nat_trans_fix_snd_arg _ _ _ _ _ θ Z) T f.

Lemma isaprop_bracket (T : algebra_ob Id_H) : isaprop (bracket T).
Proof.
  apply impred_isaprop; intro Z.
  apply impred_isaprop; intro f.
  apply isapropiscontr.
Qed.

Definition bracket_parts (T : algebra_ob Id_H) : UU
  := ∏ (Z : Ptd) (f : U Z --> `T), bracket_parts_at (nat_trans_fix_snd_arg _ _ _ _ _ θ Z) T f.

End fix_a_prestrength.

End prep_hss.

Arguments ptd_from_alg {_} _ .
Arguments prejoin_from_hetsubst {_} _ .
Arguments prejoin_from_hetsubst_η {_} _ .
Arguments prejoin_from_hetsubst_τ {_} _ .
Arguments bracket_parts {_} _ _ .

  Section def_hss.

  Context (H : Presignature C C C).

  Local Notation η := (eta_from_alg H).
  Local Notation τ := (tau_from_alg H).

  Let θ : PrestrengthForSignature H := theta H.
  Let Id_H := Id_H H.

(** the notion of a heterogeneous substitution system that asks for more operations to uniquely exist *)
Definition hss : UU := ∑ (T: algebra_ob Id_H), bracket H θ T.

Coercion hetsubst_from_hss (T : hss) : heterogeneous_substitution H.
Proof.
  exists (pr1 T).
  use tpair.
  - apply (nat_trans_fix_snd_arg _ _ _ _ _ θ).
  - apply (pr2 T).
Defined.

Definition fbracket (T : hss) (Z : Ptd) (f : U Z --> `T)
  : `T • (U Z) --> `T
  := pr1 (pr1 (pr2 T Z f)).

Notation "⦃ f ⦄_{ Z }" := (fbracket _ Z f)(at level 0).

(** The bracket operation [fbracket] is unique *)

Definition fbracket_unique_pointwise (T : hss) {Z : Ptd} (f : U Z --> `T)
  : ∏ (α : functor_composite (U Z) `T ⟹ pr1 `T),
     (∏ c : C, pr1 f c = pr1 (η T) (pr1 (U Z) c) · α c) →
     (∏ c : C, pr1 (θ (`T ⊗ Z))  c · pr1 (#H α) c · pr1 (τ T) c =
        pr1 (τ T) (pr1 (U Z) c) · α c)
     →
     α = ⦃f⦄_{Z}.
Proof.
  intros α H1 H2.
  apply path_to_ctr.
  apply whole_from_parts.
  split.
  - apply nat_trans_eq_alt; assumption.
  - apply nat_trans_eq_alt; assumption.
Qed.

Definition fbracket_unique (T : hss) {Z : Ptd} (f : U Z --> `T)
: ∏ α : (*functor_composite (C:=C)*) `T • (U Z)  --> `T,
    bracket_property_parts H (nat_trans_fix_snd_arg _ _ _ _ _ θ Z) T f α
   →
   α = ⦃f⦄_{Z}.
Proof.
  intros α [H1 H2].
  apply path_to_ctr.
  apply whole_from_parts.
  split;  assumption.
Qed.

Definition fbracket_unique_target_pointwise (T : hss) {Z : Ptd} (f : U Z --> `T)
  : ∏ α : `T • U Z --> `T,
        bracket_property_parts H (nat_trans_fix_snd_arg _ _ _ _ _ θ Z) T f α
   →
   ∏ c, pr1 α c =  pr1 ⦃f⦄_{Z} c.
Proof.
  intros α H12.
  set (t:= fbracket_unique _ _ α H12).
  apply (nat_trans_eq_weq (homset_property C) _ _ t).
Qed.

(** Properties of [fbracket] by definition: commutative diagrams *)

Lemma fbracket_η (T : hss) : ∏ {Z : Ptd} (f : U Z --> `T),
   f = η T •• U Z · ⦃f⦄_{Z}.
Proof.
  intros Z f.
  (* assert (H' := parts_from_whole T Z f (fbracket _ f)) . *)
  exact (pr1 (parts_from_whole _ _ _ _ _ (pr2 (pr1 (pr2 T Z f))))).
Qed.

Lemma fbracket_τ (T : hss) : ∏ {Z : Ptd} (f : U Z --> `T),
    θ (`T ⊗ Z) · #H ⦃f⦄_{Z} · τ T
    =
    τ T •• U Z · ⦃f⦄_{Z}.
Proof.
  intros Z f.
  exact (pr2 (parts_from_whole _ _ _ _ _ (pr2 (pr1 (pr2 T Z f))))).
Qed.

(** [fbracket] is also natural *)

Lemma fbracket_natural (T : hss) {Z Z' : Ptd} (f : Z --> Z') (g : U Z' --> `T)
:  (`T ∘ #U f : EndC⟦ `T • U Z , `T • U Z' ⟧) · ⦃g⦄_{Z'} = ⦃#U f · g⦄_{Z}.
Proof.
  apply fbracket_unique_pointwise.
  - simpl. intro c.
    rewrite assoc.
    pose proof (nat_trans_ax (η T)) as H'.
    simpl in H'.
    rewrite <- H'; clear H'.
    rewrite <- assoc.
    apply maponpaths.
    pose proof (nat_trans_eq_weq (homset_property C) _ _  (fbracket_η T g)) as X.
    simpl in X. exact (X _ ).
  - intro c; simpl.
    assert (H':=nat_trans_ax (τ T)).
    simpl in H'.
    eapply pathscomp0. 2: apply assoc'.
    eapply pathscomp0.
    2: { apply cancel_postcomposition. apply H'. }
    clear H'.
    set (H':=fbracket_τ T g).
    simpl in H'.
    assert (X:= nat_trans_eq_pointwise H' c).
    simpl in X.
    rewrite  <- assoc.
    rewrite  <- assoc.
    transitivity (  # (pr1 (H ((`T)))) (pr1 f c) ·
                     (pr1 (θ ((`T) ⊗ Z')) c)· pr1 (# H ⦃g⦄_{Z'}) c· pr1 (τ T) c).
    2: { rewrite <- assoc.
         rewrite <- assoc.
         apply maponpaths.
         repeat rewrite assoc.
         apply X.
    }
    clear X.
    set (A:=θ_nat_2_pointwise).
    simpl in *.
    set (A':= A C C C H θ (`T) Z Z').
    simpl in A'.
    set (A2:= A' f).
    clearbody A2; clear A'; clear A.
    rewrite A2; clear A2.
    repeat rewrite <- assoc.
    apply maponpaths.
    simpl.
    repeat rewrite assoc.
    apply cancel_postcomposition.
    rewrite (functor_comp H).
    apply cancel_postcomposition.
    clear H'.
    set (A:=horcomp_id_postwhisker C C C).
    rewrite A; try apply homset_property.
    apply idpath.
Qed.

(** As a consequence of naturality, we can compute [fbracket f] from [fbracket identity] *)

Lemma compute_fbracket (T : hss) : ∏ {Z : Ptd} (f : Z --> ptd_from_alg T),
    ⦃#U f⦄_{Z} = (`T ∘ # U f : EndC⟦`T • U Z , `T • U (ptd_from_alg T)⟧) · ⦃identity (U (ptd_from_alg T))⦄_{ptd_from_alg T}.
Proof.
  intros Z f.
  assert (A : f = f · identity _ ).
  { rewrite id_right; apply idpath. }
  rewrite A.
  rewrite functor_comp.
  rewrite <- fbracket_natural.
  rewrite id_right.
  apply idpath.
Qed.

Section from_identity_to_hss.
(** the operations of an hss can be obtained through this formula from
    just a heterogeneous substitution *)

  Context (T : algebra_ob Id_H).
  Context (prejoin : bracket_at H (nat_trans_fix_snd_arg _ _ _ _ _ θ (ptd_from_alg T)) T (identity _)).

  Let T0 : heterogeneous_substitution H :=
    T ,, (nat_trans_fix_snd_arg _ _ _ _ _ θ (ptd_from_alg T) ,, prejoin).

Lemma heterogeneous_substitution_into_bracket {Z : Ptd} (f : Z --> ptd_from_alg T0) :
  bracket_property H (nat_trans_fix_snd_arg _ _ _ _ _ θ Z) T0 (#U f)
    ((` T0 ∘ #U f : EndC ⟦ `T0 • U Z , `T0 • U (ptd_from_alg T0) ⟧) · prejoin_from_hetsubst T0).
Proof.
  apply whole_from_parts.
  split.
  - apply nat_trans_eq_alt.
    intro c.
    induction f as [f pt].
    simpl.
    assert (alg_map_nat := nat_trans_ax (alg_map Id_H T0) _ _ (pr1 f c)).
    etrans.
    2: { rewrite <- assoc. apply maponpaths. rewrite assoc.
         apply cancel_postcomposition.
         exact alg_map_nat.
    }
    clear alg_map_nat.
    etrans.
    2: { do 2 rewrite assoc. do 2 apply cancel_postcomposition.
         apply pathsinv0. unfold Id_H. simpl. apply BinCoproductIn1Commutes. }
    simpl.
    etrans.
    { apply pathsinv0.
      apply id_right.
    }
    do 2 rewrite <- assoc.
    apply maponpaths.
    rewrite assoc.
    assert (prejoin_ok := prejoin_from_hetsubst_η T0).
    apply (maponpaths pr1) in prejoin_ok.
    apply toforallpaths in prejoin_ok.
    apply prejoin_ok.
  - rewrite functor_comp.
    apply nat_trans_eq_alt.
    intro c.
    induction f as [f pt].
    simpl.
    assert (alg_map_nat := nat_trans_ax (alg_map Id_H T0) _ _ (pr1 f c)).
    etrans.
    2: { rewrite <- assoc. apply maponpaths. rewrite assoc.
         apply cancel_postcomposition.
         exact alg_map_nat.
    }
    etrans.
    2: { do 2 rewrite assoc. do 2 apply cancel_postcomposition.
         apply pathsinv0. unfold Id_H. simpl. apply BinCoproductIn2Commutes. }
    assert (prejoin_ok := prejoin_from_hetsubst_τ T0).
    apply (maponpaths pr1) in prejoin_ok.
    apply toforallpaths in prejoin_ok.
    assert (prejoin_ok_inst := prejoin_ok c).
    simpl in prejoin_ok_inst.
    etrans.
    { repeat rewrite assoc. do 3 apply cancel_postcomposition.
      apply pathsinv0.
      assert (theta_nat_2 := θ_nat_2_pointwise _ _ _ H θ `T0 _ _ (f,,pt) c).
      rewrite horcomp_id_postwhisker in theta_nat_2; try apply homset_property.
      apply theta_nat_2.
    }
    etrans.
    { repeat rewrite <- assoc. apply maponpaths.
      rewrite assoc.
      exact prejoin_ok_inst.
    }
    clear prejoin_ok prejoin_ok_inst.
    repeat rewrite assoc.
    apply idpath.
Qed.

End from_identity_to_hss.

(** ** Morphisms of heterogeneous substitution systems *)

(** A morphism [f] of pointed functors is an algebra morphism when... *)
(*
Definition isAlgMor {T T' : Alg} (f : T --> T') : UU :=
  #H (# U f) · τ T' = compose (C:=EndC) (τ T) (#U f).

Lemma isaprop_isAlgMor (T T' : Alg) (f : T --> T') : isaprop (isAlgMor f).
Proof.
  apply isaset_nat_trans.
  apply hs.
Qed.
*)

(** a little preparation for much later *)
Lemma τ_part_of_alg_mor  (T T' : @algebra_ob [C, C] Id_H)
  (β : @algebra_mor [C, C] Id_H T T'): #H β · τ T' = compose (C:=EndC) (τ T) β.
Proof.
  assert (β_is_alg_mor := pr2 β).
  simpl in β_is_alg_mor.
  assert (β_is_alg_mor_inst := maponpaths (fun m:EndC⟦_,_⟧ => (BinCoproductIn2 (CPEndC _ _))· m)
                                          β_is_alg_mor); clear β_is_alg_mor.
  simpl in β_is_alg_mor_inst.
  apply nat_trans_eq_alt.
  intro c.
  assert (β_is_alg_mor_inst':= nat_trans_eq_pointwise β_is_alg_mor_inst c); clear β_is_alg_mor_inst.
  simpl in β_is_alg_mor_inst'.
  rewrite assoc in β_is_alg_mor_inst'.
  eapply pathscomp0.
  2: { eapply pathsinv0.
       exact β_is_alg_mor_inst'. }
  clear β_is_alg_mor_inst'.
  apply BinCoproductIn2Commutes_right_in_ctx_dir.
  simpl.
  rewrite <- assoc.
  apply idpath.
Qed.

(** A morphism [β] of pointed functors is a bracket morphism when... *)

Lemma is_ptd_mor_alg_mor (T T' : @algebra_ob [C, C] Id_H)
  (β : @algebra_mor [C, C] Id_H T T') :
  @is_ptd_mor C (ptd_from_alg T) (ptd_from_alg T') (pr1 β).
Proof.
  simpl.
  unfold is_ptd_mor. simpl.
  intro c.
  rewrite <- assoc.
  assert (X:=pr2 β).
  assert (X':= nat_trans_eq_pointwise X c).
  simpl in *.
  etrans. { apply maponpaths. apply X'. }
  unfold coproduct_nat_trans_in1_data.
  repeat rewrite assoc.
  unfold coproduct_nat_trans_data.
  etrans.
  { apply cancel_postcomposition.
    apply BinCoproductIn1Commutes. }
  simpl.
  repeat rewrite <- assoc.
  apply id_left.
Qed.

Definition ptd_from_alg_mor {T T' : algebra_ob Id_H} (β : algebra_mor _ T T')
: ptd_from_alg T --> ptd_from_alg T'.
Proof.
  exists (pr1 β).
  apply is_ptd_mor_alg_mor.
Defined.

(** show functor laws for [ptd_from_alg] and [ptd_from_alg_mor] *)

Definition ptd_from_alg_functor_data : functor_data (category_FunctorAlg Id_H) Ptd.
Proof.
  exists ptd_from_alg.
  intros T T' β.
  apply ptd_from_alg_mor.
  exact β.
Defined.

Lemma is_functor_ptd_from_alg_functor_data : is_functor ptd_from_alg_functor_data.
Proof.
  split; simpl; intros.
  + unfold functor_idax.
    intro T.
    (* match goal with | [ |- ?l = _ ] => let ty:= (type of l) in idtac ty end. *)
    apply (invmap (eq_ptd_mor_cat _ _ _)).
    apply (invmap (eq_ptd_mor  _ _)).
    (* match goal with | [ |- ?l = _ ] => let ty:= (type of l) in idtac ty end. *)
    apply idpath.
  + unfold functor_compax.
    intros T T' T'' β β'.
    apply (invmap (eq_ptd_mor_cat _ _ _)).
    apply (invmap (eq_ptd_mor _ _)).
    apply idpath.
Qed.

Definition ptd_from_alg_functor: functor (category_FunctorAlg Id_H) Ptd :=
  tpair _ _ is_functor_ptd_from_alg_functor_data.

Definition isbracketMor {T T' : hss} (β : algebra_mor _ T T') : UU :=
    ∏ (Z : Ptd) (f : U Z --> `T),
      ⦃f⦄_{Z} · β = β •• U Z · ⦃f · #U (# ptd_from_alg_functor β)⦄_{Z}.

Lemma isaprop_isbracketMor (T T' : hss) (β : algebra_mor _ T T') : isaprop (isbracketMor β).
Proof.
  do 2 (apply impred; intro).
  apply isaset_nat_trans.
  apply homset_property.
Qed.

(** A morphism of hss is a pointed morphism that is compatible with both
    [τ] and [fbracket] *)

Definition ishssMor {T T' : hss} (β : algebra_mor _ T T') : UU
  :=   isbracketMor β.

Definition hssMor (T T' : hss) : UU
  := ∑ β : algebra_mor _ T T', ishssMor β.

Coercion ptd_mor_from_hssMor (T T' : hss) (β : hssMor T T') : algebra_mor _ T T' := pr1 β.

(*
Definition isAlgMor_hssMor {T T' : hss} (β : hssMor T T')
  : isAlgMor β := pr1 (pr2 β).
*)
Definition isbracketMor_hssMor {T T' : hss} (β : hssMor T T')
  : isbracketMor β := pr2 β.

(** **** Equality of morphisms of hss *)

Section hssMor_equality.

(** Show that equality of hssMor is equality of underlying nat. transformations *)

Variables T T' : hss.
Variables β β' : hssMor T T'.
Definition hssMor_eq1 : β = β' ≃ (pr1 β = pr1 β').
Proof.
  apply subtypeInjectivity.
  intro γ.
  apply isaprop_isbracketMor.
Defined.

Definition hssMor_eq : β = β' ≃ (β : EndC ⟦ _ , _ ⟧) = β'.
Proof.
  eapply weqcomp.
  - apply hssMor_eq1.
  - apply subtypeInjectivity.
    intro.
    apply isaset_nat_trans.
    apply homset_property.
Defined.

End hssMor_equality.

Lemma isaset_hssMor (T T' : hss) : isaset (hssMor T T').
Proof.
  intros β β'.
  apply (isofhlevelweqb _ (hssMor_eq _ _ β β')).
  apply isaset_nat_trans.
  apply homset_property.
Qed.

(** ** The precategory of hss *)

(** *** Identity morphism of hss *)

Lemma ishssMor_id (T : hss) : ishssMor (identity (C:=category_FunctorAlg _) (pr1 T)).
Proof.
  unfold ishssMor.
  unfold isbracketMor.
  intros Z f.
  rewrite id_right.
  rewrite functor_id.
  rewrite id_right.
  apply pathsinv0.
  set (H2:=pre_composition_functor _ _ C (U Z)).
  set (H2' := functor_id H2). simpl in H2'.
  simpl.
  rewrite H2'.
  rewrite (@id_left EndC).
  apply idpath.
Qed.

Definition hssMor_id (T : hss) : hssMor _ _ := tpair _ _ (ishssMor_id T).

(** *** Composition of morphisms of hss *)

Lemma ishssMor_comp {T T' T'' : hss} (β : hssMor T T') (γ : hssMor T' T'')
  : ishssMor (compose (C:=category_FunctorAlg _) (pr1 β)  (pr1 γ)).
Proof.
  unfold ishssMor.
  unfold isbracketMor.
  intros Z f.
  eapply pathscomp0; [apply assoc|].
  (* match goal with | [|- ?l = _ ] => assert (Hyp : l = fbracket T f· pr1 β· pr1 γ) end. *)
  etrans.
  { apply cancel_postcomposition.
    apply isbracketMor_hssMor. }
  rewrite <- assoc.
  etrans.
  { apply maponpaths.
    apply isbracketMor_hssMor. }
  rewrite assoc.
  do 2 rewrite functor_comp.
  rewrite assoc.
  apply cancel_postcomposition.
  apply pathsinv0, (functor_comp (pre_composition_functor _ _ C (U Z)) ).
Qed.

Definition hssMor_comp {T T' T'' : hss} (β : hssMor T T') (γ : hssMor T' T'')
  : hssMor T T'' := tpair _ _ (ishssMor_comp β γ).

Definition hss_obmor : precategory_ob_mor.
Proof.
  exists hss.
  exact hssMor.
Defined.

Definition hss_precategory_data : precategory_data.
Proof.
  exists hss_obmor.
  split.
  - exact hssMor_id.
  - exact @hssMor_comp.
Defined.

Lemma is_precategory_hss : is_precategory hss_precategory_data.
Proof.
  apply is_precategory_one_assoc_to_two.
  repeat split; intros.
  - apply (invmap (hssMor_eq _ _ _ _ )).
    apply (@id_left EndC).
  - apply (invmap (hssMor_eq _ _ _ _ )).
    apply (@id_right EndC).
  - apply (invmap (hssMor_eq _ _ _ _ )).
    apply (@assoc EndC).
Qed.

Definition hss_precategory : precategory := tpair _ _ is_precategory_hss.

Lemma has_homsets_precategory_hss : has_homsets hss_precategory.
Proof.
  red.
  intros T T'.
  apply isaset_hssMor.
Qed.

Definition hss_category : category := hss_precategory ,, has_homsets_precategory_hss.

End def_hss.

End fix_a_category.

Arguments hss {_} _ _ .
Arguments hssMor {_ _ _ } _ _ .
Arguments fbracket {_ _ _} _ _ _ .
Arguments fbracket_η {_ _ _} _ {_} _ .
Arguments fbracket_τ {_ _ _} _ {_} _ .
Arguments fbracket_unique_target_pointwise {_ _ _ } _ {_ _ _} _ _.
Arguments fbracket_unique {_ _ _ } _ {_} _ {_} _ .
(* Arguments Alg {_ _} _. *)
Arguments hss_precategory {_} _ _ .
Arguments hss_category {_} _ _ .
Arguments eta_from_alg {_ _ _} _.
Arguments tau_from_alg {_ _ _} _.
Arguments ptd_from_alg {_ _ _} _.
Arguments ptd_from_alg_functor {_} _ _ .
Arguments bracket_property {_ _ _ _ } _ _ _ _ .
Arguments bracket_property_parts {_ _ _ _} _ _ _ _ .
Arguments bracket {_ _ _} _ _.
Arguments prejoin_from_hetsubst {_ _ _} _ .
Arguments prejoin_from_hetsubst_η {_ _ _} _ .
Arguments prejoin_from_hetsubst_τ {_ _ _} _ .

Notation τ := tau_from_alg.
Notation η := eta_from_alg.


Require Import Lia Div2 Even PeanoNat.

Class St := mkSt { x : nat; r : nat; h : nat; m : nat }.

Require Import Statement Wpr Wfd CbC.

Definition sqr n := n * n.

Open Scope cbc_scope.
Open Scope stmt_scope.

Program Definition SquareRoot :=
  ⟨ fun '(mkSt x r h m) '(mkSt x' r' h' m')  => (sqr r' <= x < sqr (1 + r') /\ x = x') ⟩
  :{
      '(mkSt x r h m) ::= (mkSt x 0 (1 + x) m);;
      ⟨fun '(mkSt x r h m) '(mkSt x' r' h' m') => (sqr r <= sqr r' <= x /\ x < sqr h' <= sqr h) /\ x = x' /\ 1 + r' = h' ⟩ :{
        While (fun '(mkSt x r h m) => 1 + r <> h) Do
           ⟨ fun '(mkSt x r h m) '(mkSt x' r' h' m') => (sqr r <= sqr r' <= x /\ x < sqr h' <= sqr h) /\ (r < r' \/ h > h') /\ x = x' ⟩
            :{
               ⟨ fun '(mkSt x r h m) '(mkSt x' r' h' m') => r < m' < h /\ x = x' /\ r = r' /\ h = h' ⟩
               :{
                  '(mkSt x r h m) ::= (mkSt x r h (r + div2 (h - r)))
                };;
                If (fun '(mkSt x r h m) => sqr m <= x) Then
                  '(mkSt x r h m) ::= (mkSt x m h m)
                Else
                 '(mkSt x r h m) ::= (mkSt x r m m)
                End
             }
        Done
     }
  }.

Fail Check SquareRoot.

Next Obligation.
  intros (x,r,h,m) ((x',r',h',m'),(HHrh,_)); simpl in *.
  split; auto.
  set (k := PeanoNat.Nat.div2 (h-r)).
  assert (2*r < 2*r + 2*k < 2*h).
  { rewrite <- (Nat.double_twice k).
    destruct (even_or_odd (h-r)); subst k.
    { rewrite <- (even_double _ H). lia. }
    { remember (h - r); destruct H.
      rewrite <- (even_div2 _ H), <- (even_double _ H).
      unfold sqr in *; nia.
    }
  }
  lia.
Qed.

Next Obligation.
  intros (x,r,h,m) ((x',r',h',m'),HH); simpl in *.
  split.
  { intros (xx,rx,hx,mx); simpl.
    intros [ HH1' [HH2 [HH3' HH4 ]]]; subst xx rx hx.
    assert (sqr mx <= x \/ sqr mx > x) as [ HH1 | HH1 ] by lia.
    { left; unfold sqr in *.
      split; auto; split; try lia.
      split.
      { split; try nia. }
      { split; auto; lia. }
    }
    { right; split; try lia.
      split; try lia.
      unfold sqr in *. nia.
    }
  }
  { exists (mkSt x r h (r + 1)).
    repeat split; auto; try lia.
    unfold sqr in *; nia.
  }
Qed.

Next Obligation.
  split.
  { apply (Wfd.by_inclusion _ _ (fun '(mkSt x' r' h' m') '(mkSt x r h m)  =>
                                   r < h /\ (r < r' < h' /\ h' <= h \/ h' < h /\ r <= r' < h'))).
    { apply (Wfd.by_nat_variant _ _ (fun '(mkSt x r h m) => h - r)).
      intros (x,r,h,m) (x',r',h',m'). lia.
    }
    { intros (x,r,h,m) (x',r',h',m').
      unfold sqr; nia.
    }
  }
  { intros (x,r,h,m) (x',r',h',m') (HHrh2,(HHx,(HHrh3,HHeq))); subst x'.
    split.
    { intros (xx,rx,hx,mx) [ HH | (HH1,HH2) ].
      { unfold sqr in *; nia. }
      { inversion HH2; subst xx rx hx mx. lia. }
    }
    { assert (S r' = h' \/ S r' <> h') as [ HH | HH ] by lia.
      { exists (mkSt x r' h' m'); auto. }
      { assert (sqr (S r') <= x \/ x < sqr (S r')) as [ HH1 | HH1 ] by lia.
        { exists (mkSt x (S r') h' m'); left; split; auto. unfold sqr in *; nia. }
        {  exists (mkSt x r' (S r') m').
           left; split; auto.
           split; auto.
           { split.
             { split.
               { lia. }
               { clear HHrh3.
                 unfold sqr in *; nia.
               }
             }
             { split; auto.
               destruct HHx as (HHx1,HHx2).
               clear -HHx1 HHx2 HH.
               unfold sqr in *; nia.
             }
           }
           { split; auto; right; unfold sqr in *; nia. }
        }
      }
    }
  }
Qed.

Next Obligation.
  intros (x,r,h,m) ((x',r',h',m'),HH); simpl in *.
  destruct HH as ((HH1,HH2),(HH3,HH4)).
  assert (S r = h \/ S r <> h) by lia.
  destruct H.
  { right; split; auto; split; auto; split; auto.
    { subst; unfold sqr in *; nia. }
    { lia. }
  }
  { left; split; auto.
    split.
    { intros (xx,rx,hx,mx); intros ((HHx',(HHprog,HHxx)),HHrh); subst xx.
      split; auto; split; auto.
      lia.
    }
    { exists (mkSt x' r' h' m').
      split; auto.
      { unfold sqr in *.
        repeat split; auto.
        all : nia.
      }
    }
  }
Qed.

Next Obligation.
  intros (x,r,h,m) ((x',r',h',m'),(HH,HHx)); subst x'; simpl in *.
  split.
  { intros (xx,rx,hx,mx); intros ((HHx,HHx'),(HHprog,HHxx)); subst.
    unfold sqr in *; simpl; split; auto.
    nia.
  }
  { exists (mkSt x r' (S r') m').
    split; auto.
    { unfold sqr in *.
      split.
      { nia. }
      { split.
        { nia. }
        { assert (r' <= x) by nia.
          clear HH.
          nia.
        }
      }
    }
  }
Qed.

Definition proof : (Πc SquareRoot) ⊑ (Πa SquareRoot) := CbC.soundness SquareRoot.
Check proof. (*  →  proof : (Π2 SquareRoot) ⊑ (Π1 SquareRoot) *)
Print Assumptions proof. (* → Closed under the global context *)

Close Scope stmt_scope.
Close Scope cbc_scope.


Require Import HoTT.
Require Import UnivalenceAxiom.

Require Import Denotation.
Require Import HoTTEx.

Module Type Types.
  Parameter type : Type.
  Parameter denotationType : Denotation type Type.
  Parameter isHSetType : forall T:type, IsHSet ⟦ T ⟧.

  Parameter constant : type -> Type.
  Parameter unary : type -> type -> Type.
  Parameter binary : type -> type -> type -> Type.

  Parameter denotationConstant : forall S, Denotation (constant S) ⟦ S ⟧.
  Parameter denotationUnary : forall S T, Denotation (unary S T) (⟦ S ⟧ -> ⟦ T ⟧).
  Parameter denotationBinary : forall S T U, Denotation (binary S T U) (⟦ S ⟧ -> ⟦ T ⟧ -> ⟦ U ⟧).
End Types.

Inductive Tree (A:Type) :=
| node (t0 t1 : Tree A)
| leaf (a:A)
| empty
.

Arguments node {_} _ _.
Arguments leaf {_} _.
Arguments empty {_}.

Module SQLDenotation (T : Types).
  Import T.
  (* definition of inductive types in
     modules isn't supported by Coq in this case *)
  (* NOTE ideally schemas would have little structure, and users
          of this library would be free to implement products of
          schemas etc as they wish. All this structure is useful
          though because it introduces more computation in our
          proofs, and thus leads to much simpler proofs. *)
  Definition Schema := Tree type.
  Definition singleton := @leaf type.
  Notation "s0 ++ s1" := (node s0 s1).

  Fixpoint Tuple (s:Schema) : Type.
    refine (match s with
    | node t0 t1 => (Tuple t0) * (Tuple t1)
    | leaf T => ⟦T⟧
    | empty => Unit
    end).
  Defined.

  Instance isHSetTuple s : IsHSet (Tuple s).
    induction s; refine (_).
  Defined.

  (* NOTE we could have `Tuple s -> nat` for the bag semantics,
          but this is more general as we support infinite types.
          this is useful for projection *)
  Definition Relation s := Tuple s -> Type.
  Definition Query Γ s := Tuple Γ -> Relation s.
End SQLDenotation.

Module Type Schemas (T : Types).
  Import T.
  Module TD := SQLDenotation T.
  Import TD.
  Export TD.
End Schemas.

Module Type Relations (T : Types) (S : Schemas T).
  Import T S.

  Parameter relation : Schema -> Type.
  Parameter denotationRelation : forall s, Denotation (relation s) (Relation s).
End Relations.

Module Type Aggregators (T : Types) (S : Schemas T).
  Import T S.

  Parameter aggregator : type -> type -> Type.
  Parameter denotationAggregator : forall S T, Denotation (aggregator S T) (Relation (leaf S) -> ⟦ T ⟧).
End Aggregators.

(* We have SQL depend on modules instead of type class instances
   because type classes lead to bad unfolding behavior of
   mutually inductive fixpoints. *)
Module SQL (T : Types) (S : Schemas T) (R : Relations T S) (A : Aggregators T S).
  Import T S R A.

  Inductive SQL : Schema -> Schema -> Type :=
  | table {Γ s} : relation s -> SQL Γ s
  | union {Γ s} : SQL Γ s -> SQL Γ s -> SQL Γ s
  | minus {Γ s} : SQL Γ s -> SQL Γ s -> SQL Γ s
  | select  {Γ s} : Pred (Γ ++ s) -> SQL Γ s -> SQL Γ s
  | product {Γ s0 s1} : SQL Γ s0 -> SQL Γ s1 -> SQL Γ (s0 ++ s1)
  | project {Γ s s'} : Proj (Γ ++ s) s' -> SQL Γ s -> SQL Γ s'
  | distinct {Γ s} : SQL Γ s -> SQL Γ s
  | castSQL {Γ Γ' s} : Proj Γ Γ' -> SQL Γ' s -> SQL Γ s

  with Pred : Schema -> Type :=
  | inhabited {Γ s} : SQL Γ s -> Pred Γ
  | equal {Γ T} : Expr Γ T -> Expr Γ T -> Pred Γ
  | negate {Γ} : Pred Γ -> Pred Γ
  | and {Γ} : Pred Γ -> Pred Γ -> Pred Γ
  | or {Γ} : Pred Γ -> Pred Γ -> Pred Γ
  | false {Γ} : Pred Γ
  | true {Γ} : Pred Γ
  | castPred {Γ Γ'} : Proj Γ Γ' -> Pred Γ' -> Pred Γ

  with Proj : Schema -> Schema -> Type :=
  | combine  {Γ Γ0 Γ1} : Proj Γ Γ0 -> Proj Γ Γ1 -> Proj Γ (Γ0 ++ Γ1)
  | left  {Γ0 Γ1} : Proj (Γ0 ++ Γ1) Γ0
  | right  {Γ0 Γ1} : Proj (Γ0 ++ Γ1) Γ1
  | compose  {Γ Γ' Γ''} : Proj Γ Γ' -> Proj Γ' Γ'' -> Proj Γ Γ''
  | star     {Γ} : Proj Γ Γ
  | e2p {T Γ} : Expr Γ T -> Proj Γ (singleton T)
  | erase    {Γ} : Proj Γ empty

  with Expr : Schema -> type -> Type :=
  | variable {T Γ} (c:Proj Γ (leaf T)) : Expr Γ T
  | aggregate {Γ S T} : aggregator S T -> SQL Γ (singleton S) -> Expr Γ T
  | constantExpr {Γ S} : constant S -> Expr Γ S
  | unaryExpr {Γ S T} : unary S T -> Expr Γ S -> Expr Γ T
  | binaryExpr {Γ S T U} : binary S T U -> Expr Γ S -> Expr Γ T -> Expr Γ U
  | castExpr {Γ Γ' T} : Proj Γ Γ' -> Expr Γ' T -> Expr Γ T
  .

  Fixpoint denoteSQL {Γ s} (a : SQL Γ s) : Query Γ s
  with     denotePred {Γ} (slct : Pred Γ) : Tuple Γ -> hProp
  with     denoteProj {Γ Γ'} (cast: Proj Γ Γ') {struct cast} : Tuple Γ -> Tuple Γ'
  with     denoteExpr {Γ T} (e : Expr Γ T) : Tuple Γ -> ⟦T⟧.
    - refine (
      match a in SQL Γ s return Query Γ s with
      | table _ _ r => fun _ => ⟦r⟧
      | union _ _ a a' => fun g t => denoteSQL _ _ a g t + denoteSQL _ _ a' g t
      | minus _ _ a a' => fun g t => denoteSQL _ _ a g t * hnot (merely (denoteSQL _ _ a' g t))
      | select _ _ slct a => fun g t => denotePred _ slct (g, t) *
                                    denoteSQL _ _ a g t
      | product _ _ _ a0 a1 => fun g t => denoteSQL _ _ a0 g (fst t) *
                                      denoteSQL _ _ a1 g (snd t)
      | project _ _ _ proj a => fun g t' => ∑ t, denoteSQL _ _ a g t *
                                             BuildhProp (denoteProj _ _ proj (g, t) = t')
      | distinct _ _ a => fun g t => merely (denoteSQL _ _ a g t)
      | castSQL _ _ _ f a => fun g t => denoteSQL _ _ a (denoteProj _ _ f g) t
      end).
    - refine (
      match slct in Pred Γ return Tuple Γ -> hProp with
      | inhabited _ _ a => fun g => hexists (fun t => denoteSQL _ _ a g t)
      | equal _ _ e0 e1 => fun g => BuildhProp (denoteExpr _ _ e0 g = denoteExpr _ _ e1 g)
      | negate _ slct => fun g => hnot (denotePred _ slct g)
      | and _ slct0 slct1 => fun g => hand (denotePred _ slct0 g) (denotePred _ slct1 g)
      | or _ slct0 slct1 => fun g => hor (denotePred _ slct0 g) (denotePred _ slct1 g)
      | false _ => fun _ => False
      | true _ => fun _ => True
      | castPred _ _ f slct => fun g => denotePred _ slct (denoteProj _ _ f g)
      end).
    - refine (
      match cast in Proj s s' return Tuple s -> Tuple s' with
      | combine _ _ _ c c' => fun t => (denoteProj _ _ c t, denoteProj _ _ c' t)
      | left _ _ => fst
      | right _ _ => snd
      | compose _ _ _ c c' => fun t => denoteProj _ _ c' (denoteProj _ _ c t)
      | star _ => fun t => t
      | e2p _ _ e => fun t => denoteExpr _ _ e t
      | erase _ => fun _ => tt
      end).
    - refine (
      match e in Expr Γ T return Tuple Γ -> ⟦T⟧ with
      | variable T _ c => fun g => denoteProj _ _ c g
      | aggregate _ _ _ aggr a => fun g => ⟦aggr⟧ (denoteSQL _ _ a g)
      | constantExpr _ _ c => fun _ => ⟦ c ⟧
      | unaryExpr _ _ _ f e => fun g => ⟦ f ⟧ (denoteExpr _ _ e g)
      | binaryExpr _ _ _ _ f e0 e1 => fun g =>  ⟦ f ⟧ (denoteExpr _ _ e0 g) (denoteExpr _ _ e1 g)
      | castExpr _ _ _ f e => fun g => denoteExpr _ _ e (denoteProj _ _ f g)
      end).
  Defined.

  Global Instance denotationProj {Γ Γ'} : Denotation (Proj Γ Γ') (Tuple Γ -> Tuple Γ') := {|
    denote := denoteProj
  |}.

  Global Instance denotationSQL Γ s : Denotation (SQL Γ s) _ := {|
    denote := denoteSQL
  |}.

  Global Instance denotationPred Γ : Denotation (Pred Γ) _ := {|
    denote := denotePred
  |}.

  Global Instance denotationExpr Γ T : Denotation (Expr Γ T) _ := {|
    denote := denoteExpr
  |}.

  Definition Column T Γ := Proj Γ (leaf T).

  Definition getLeaf {T Γ} : Column T Γ -> Proj Γ (singleton T) :=
    fun c => e2p (variable c).

  Definition Projection Γ s s' := Proj (Γ ++ s) s'.

  Arguments Projection _ _ _ /.

  Definition projectStar {Γ s} : Projection Γ s s := right.

  Definition projectSingleton {Γ T s} (e:Expr (Γ ++ s) T) : Projection Γ s (singleton T) := e2p e.

  Definition projectNil {Γ s} : Projection Γ s empty := erase.

  Definition projectCons {Γ T s s'} (e:Expr (Γ ++ s) T) (proj:Projection Γ s s') : Projection Γ s (node (singleton T) s') :=
    combine (projectSingleton e) proj.

  Notation "p1 ⋅ p2" := (compose p1 p2) (at level 45).
  Notation "Γ ⊢ x : s" := (x:(SQL Γ s)) (at level 45, x at level 45).
  Notation "a 'WHERE' c" := (select c a) (at level 45, c at level 45).
  Notation "'SELECT' '*' a" := (a) (at level 45).
  Notation "'SELECT1' proj a" := (project proj a) (at level 45).
  Notation "'SELECT2' proj0 , proj1 a" := (project (projectCons proj0 (projectCons proj1 projectNil)) a) (at level 45).
  Notation "'FROM1' a" := (a) (at level 45).
  Notation "'FROM2' a , b" := (product a b) (at level 45).  (* TODO kill, and rename FROM2 to FROM *)
  Notation "'FROM' a , b , .. , c" := (product .. (product a b) .. c) (at level 46, c at level 44).
  Notation "a 'UNION' 'ALL' b" := (union a b) (at level 45).
  Notation "a 'MINUS' b" := (minus a b) (at level 45).
  Notation "'EXISTS' s" := (inhabited s) (at level 45).
  Notation "s0 'AND' s1" := (and s0 s1) (at level 45).
  Notation "s0 'OR' s1" := (or s0 s1) (at level 45).
  Notation "'NOT' s0" := (negate s0) (at level 45).
  Notation "'FALSE'" := (false) (at level 45).
  Notation "'TRUE'" := (true) (at level 45).
  Notation "'DISTINCT' s" := (distinct s) (at level 45).

  Instance isHSetSQL `{forall s (r:relation s) t, IsHSet (⟦ r ⟧ t)} Γ s q g t : IsHSet (⟦ Γ ⊢ q : s ⟧ g t).
    induction q; refine (_).
  Defined.

  (* A group by projection has access to the context, a selected tuple t,
     and a query that returns all tuples that match t *)
  Definition GroupByProj Γ s s' := SQL (Γ ++ s) s -> Proj (Γ ++ s) s'.

  Definition combineGroupByProj {Γ s s' s''}
    (p1:GroupByProj Γ s s') (p2:GroupByProj Γ s s'') : GroupByProj Γ s (s' ++ s'') :=
    fun a => combine (p1 a) (p2 a).

  Definition plainGroupByProj {Γ s s'} (p:Proj (Γ ++ s) s') : GroupByProj Γ s s' := fun _ => p.
  Arguments plainGroupByProj [_ _ _] _ / _.

  Definition aggregatorGroupByProj {Γ s S T} : aggregator S T -> Expr (Γ ++ s) S -> GroupByProj Γ s (leaf T).
    intros agg e a.
    refine (e2p (aggregate agg (SELECT1 (e2p _) FROM1 a))).
    refine (castExpr (combine (left⋅left) right) e).
  Defined.
  Arguments aggregatorGroupByProj [_ _ _ _] _ _ / _.

  Fixpoint tupleEqual {Γ s} : Proj Γ s -> Proj Γ s -> Pred Γ :=
    match s with
    | s0 ++ s1 => fun p0 p1 =>
        tupleEqual (p0 ⋅ left)  (p1 ⋅ left) AND
        tupleEqual (p0 ⋅ right) (p1 ⋅ right)
    | leaf T => fun p0 p1 =>
        equal (variable p0) (variable p1)
    | empty => fun _ _ => TRUE
    end.

  Definition groupBy {Γ s s' C} (proj : GroupByProj Γ s s')
                     (a : SQL Γ s) (c : Proj (Γ ++ s) C) : SQL Γ s'.
    refine (DISTINCT SELECT1 proj (
              SELECT * FROM1 castSQL left a
              WHERE tupleEqual (left ⋅ c) ((combine (left⋅left) right) ⋅ c))
            FROM1 a).
  Defined.

  Arguments groupBy [_ _ _ _] _ _ _ /.

  (* Semi-join with single equality predicate. *)
  Definition semiJoin1 {Γ s1 s2 ty} (a: SQL Γ s1)
             (b: SQL Γ s2) (c1: Column ty s1)
             (c2: Column ty s2) : SQL Γ s1.
    pose (c1' := @variable ty ((Γ ++ s1) ++ s2) (left⋅right⋅c1)).
    pose (c2' := @variable ty ((Γ ++ s1) ++ s2) (right⋅c2)).
    refine (SELECT * FROM1 a WHERE EXISTS (SELECT * FROM1 _ WHERE (equal c1' c2'))).
    refine (castSQL left b).
  Defined.

  Definition semiJoin {Γ s1 s2} (a: SQL Γ s1) (b: SQL Γ s2)
             (slct: Pred (Γ ++ s1 ++ s2)): SQL Γ s1.
    refine (SELECT * FROM1 a WHERE EXISTS (SELECT * FROM1 (castSQL left b) WHERE (castPred _ slct) )).
    refine (combine (left⋅left) (combine (left⋅right) right)).
  Defined.

  (* SEMI JOIN with single equality predicate *)
  Notation "a 'SEMI_JOIN1' b 'ON' 'SEQ' c1 , c2 " := (semiJoin1 a b c1 c2) (at level 45).

  (* SEMI JOIN with predicate *)
  Notation "a 'SEMI_JOIN' b 'ON' slct " := (semiJoin a b slct) (at level 45).

  Notation "'PLAIN' ( e )" := (plainGroupByProj (e2p e)).
  Notation "'SELECT' proj 'FROM1' a 'GROUP' 'BY' v " := (groupBy proj a v) (at level 45).

End SQL.


(* (c) Copyright ? *)

(*****************************************************************************
  Matrix derivation (differentiation).

 lift_to E A == lift A : 'M[R]_(m,n) to 'M[E]_(m,n), where E : lalgType R.
                Matrix derivation is commutative with multiplication by a
                (lifted) constant matrix.
      lift A == lift_to _ A

  All the following results are parametrized on f : {derMorph R -> V}, which is
  a derivation on the elements. Derivation on matrices are simply defined as
  (map_mx f). V : bimodType R R. The following notations are used:
          \d == f
         \\d == map_mx f

  The main results about matrix derivation are:
        dmM :      \\d (A *m B) = \\d A *m B + A *m \\d B
        dmI :             \\d I = 0
       dmcs :     \\d (c *ml: A) = c *: 1 *ml: \\d A
    dmscale : \\d (c *ml: A) = (c *ml: I) *ml \\d A
        dmc :      \\d (lift C) = 0
       dmcl : \\d (lift C *m A) = lift C *ml \\d A
              (and the symmetric version dmcr)
  dm_lrkron : \\d (A *o B) = \\d A *or B + A *ol \\d B
dm_lkron1mx : \\d (I *o A) = I *ol \\d A
dm_rkronmx1 : \\d (A *o I) = \\d A *or I

  \\d on non-empty square matrices inherit the derMorh, derAdd and linearDer
  structure from \d.

******************************************************************************)

Set Implicit Arguments.
Unset Strict Implicit.
Unset Printing Implicit Defensive.

Require Import ssreflect ssrfun ssrbool eqtype ssrnat div seq choice fintype.
Require Import finfun bigop prime binomial.

Require Import ssralg.
Import GRing.Theory.
Open Local Scope ring_scope.

Require Import matrix.

Require Import mxutil.
Import Notations.
Require Import bimodule.
Require Import derivation.
Require Import mxmodule.
Import Notations.

Section Ring.

(* Element type *)
Variable E : ringType.
Variable D : bimodType E E.
Variable der : {derAdd D}.

Notation "\d" := der.
Notation "\\d" := (@map_mx _ _ \d _ _).

Section AnyMatrix.

Variable m n r : nat.
Implicit Types A : 'M[E]_(m, n).
Implicit Types B : 'M[E]_(n, r).

(* Matrix derivation is derivative (has Lebniz product rule) for matrix multiplication *)
Lemma dmM A B : \\d (A *m B) = \\d A *mr B + A *ml \\d B.
Proof.
  by apply/matrixP => i k; rewrite !mxE !raddf_sum -big_split; apply eq_bigr => j; rewrite /= !mxE derM.
Qed.

End AnyMatrix.

Lemma dmI n : \\d I = 0 :> 'M_n.
Proof.
  apply: (addIr (\\d (I *m I))).
  by rewrite add0r {1}mul1mx dmM lmul1mx rmulmx1.
Qed.

(*here*)

Section SquareMatrix.

Variable n' : nat.
Local Notation n := n'.+1.

(* Can only register for DerMorph here because DerMorph is only defined for rings, not graded rings *)
Canonical dm_der := DerMorphFor 'M[D]_n^m (@dmM n n n).
(* Now that \\d has a canonical DerMorph structure, and since it already has a canonical Additive structure thanks to map_mx_additive, packDerAdd will automatically generate a DerAdd structure *)
Canonical dm_derAdd := packDerAdd 'M[D]_n^m \\d.

End SquareMatrix.

End Ring.

Section Lalgebra.

Variable R : ringType.
Variable E : lalgType R.
Variable D : bimodType E E.
Variable der : {linearDer E -> D}.

Notation "\d" := der.
Notation "\\d" := (@map_mx _ _ \d _ _).

Lemma dmcs m c (A : 'M[E]_m) : \\d (c *ml: A) = c%:A *ml: \\d A .
Proof. by apply/matrixP => i j; rewrite !mxE linearZ /=. Qed.

Lemma dmscale m c (A : 'M[E]_m) : \\d (c *ml: A) = (c *ml: I) *ml \\d A.
Proof.
  by rewrite dmcs -lmul_scalar_mx -scalar_gscalar -scalemx1 scale_lscale.
Qed.

Variable n' : nat.
Local Notation n := n'.+1.

Lemma dmscale' c (A : 'M[E]_n^s) : \\d (c *: A) = c *: (1 : 'M[E]_n^s) *: (\\d A : 'M[D]_n^m).
Proof. by rewrite dmscale. Qed.

Canonical dm_linear := AddScale ('M[E]_n^s -> 'M[D]_n^m) dmscale'.

End Lalgebra.

Section DerKronProd.

Variable E : comRingType.
Variable D : comBimodType E.
Variable der : {derAdd D}.
Notation "\d" := (DerAdd.apply der).
Notation "\\d" := (map_mx \d).

Section Main.

Lemma dmdelta m n i j : \\d (delta_mx i j) = 0 :> 'M[D]_(m,n).
Proof.
  apply/matrixP=> i' j'; rewrite !mxE /=.
  case (_ && _).
  - by rewrite der1.
  - by rewrite der0.
Qed.

Variable m1 n1 m2 n2 : nat.
Implicit Types A : 'M[E]_(m1,n1).
Implicit Types B : 'M[E]_(m2,n2).

(* Matrix derivation is also derivative (has Lebniz product rule) for Kronecker product, like it is for matrix multiplication *)
Lemma dm_lrkron A B : \\d (A *o B) = \\d A *or B + A *ol \\d B.
Proof.
  apply/matrixP => i j.
  case/mxvec_indexP: i => n1i n2i.
  case/mxvec_indexP: j => m1i m2i.
  by rewrite mxE kronE mxE lkronE rkronE derM !mxE.
Qed.

End Main.

Section Corollaries.

Variable m n r : nat.
Implicit Types A : 'M[E]_(m,n).

Lemma dm_lkron1mx A : \\d (I *o A) = (I : 'M_(r,_)) *ol \\d A.
Proof.
  by rewrite dm_lrkron dmI rkron0mx add0r.
Qed.

Lemma dm_rkronmx1 A : \\d (A *o I) = \\d A *or (I : 'M_(_,r)).
Proof.
  by rewrite dm_lrkron dmI lkronmx0 addr0.
Qed.

End Corollaries.

End DerKronProd.

(* Lift a matrix of 'M[R]_(m,n) into 'M[E]_(m,n), where E : lalgTyp R *)
Section Lift.

Variable R : ringType.
Variable E : lalgType R.
Variable m n r : nat.
Implicit Types C : 'M[R]_(m,n).
Implicit Types D : 'M[R]_(n,r).

Notation lift := (map_mx (in_alg E)).

Lemma lift_mul C D : lift (C *m D) = lift C *m lift D.
Proof.
  apply/matrixP=> i j; rewrite !mxE raddf_sum.
  apply eq_bigr => k.
  by rewrite !mxE -scalerAl mul1r scalerA.
Qed.

Lemma lift_vec C : lift (vec C) = vec (lift C).
  by rewrite map_vec.
Qed.

End Lift.

Local Notation lift := (map_mx (in_alg _)).
Local Notation lift_to E := (map_mx (in_alg E)).

(* Matrix derivation is commutative with multiplication by a (lifted) constant matrix *)
Section DmLift.

Variable R : ringType.
Variable E : algType R.
Variable V : bimodType E E.
Variable der : {linearDer E -> V}.
Notation "\d" := (LinearDer.apply der).
Notation "\\d" := (map_mx \d).
Variable m n r : nat.
Implicit Types C : 'M[R]_(m,n).
Implicit Types D : 'M[R]_(n,r).

Lemma dmc C : \\d (lift_to E C) = 0.
Proof.
  by apply/matrixP=> i j; rewrite !mxE linearZ /= /cscale der1 scaler0.
Qed.

Lemma dmcl C (A : 'M[E]_(_, r)) : \\d (lift C *m A) = lift C *ml \\d A.
Proof.
  by rewrite dmM dmc rmul0mx add0r.
Qed.

Lemma dmcr (A : 'M[E]_(r, _)) C : \\d (A *m lift C) = \\d A *mr lift C.
Proof.
  by rewrite dmM dmc lmulmx0 addr0.
Qed.

End DmLift.

Module Notations.

Notation lift_to E := (map_mx (in_alg E)).
Notation lift := (map_mx (in_alg _)).

End Notations.


From Hammer Require Import Tactics Reflect Hints.
From Hammer Require Import Hammer.

From mathcomp Require Import ssreflect ssrfun zmodp ssrbool ssrnat ssralg seq fintype finalg tuple eqtype.

From Coq Require Import FunctionalExtensionality.
From Coq Require Import Relation_Definitions RelationClasses.

Require Import CoqArith.Utils.

Require Import Program.

Section Sigma_1_1_Internal.

Inductive PolyTerm {A : nat -> nat} : Type :=
| PolyVar : nat -> PolyTerm
| PolyFun : forall i, (|[A i]| -> PolyTerm) -> PolyTerm
| PolyMinusOne : PolyTerm
| PolyPlusOne : PolyTerm
| PolyZero : PolyTerm
| PolyPlus : PolyTerm -> PolyTerm -> PolyTerm
| PolyTimes : PolyTerm -> PolyTerm -> PolyTerm
| PolyInd : PolyTerm -> PolyTerm -> PolyTerm.

Inductive ZerothOrderFormula {A} : Type :=
| ZONot : ZerothOrderFormula -> ZerothOrderFormula
| ZOAnd : ZerothOrderFormula -> ZerothOrderFormula -> ZerothOrderFormula
| ZOOr : ZerothOrderFormula -> ZerothOrderFormula -> ZerothOrderFormula
| ZOImp : ZerothOrderFormula -> ZerothOrderFormula -> ZerothOrderFormula
| ZOEq : @PolyTerm A -> @PolyTerm A -> ZerothOrderFormula.

Inductive QuantifiedFormula {A} : Type :=
| ZO : @ZerothOrderFormula A
    -> QuantifiedFormula
| QExists : forall (bs : seq (@PolyTerm A))
                   (y : @PolyTerm A),
            @QuantifiedFormula (ExtendAt0 (length bs) A)
         -> QuantifiedFormula
| QForall : @PolyTerm A
         -> QuantifiedFormula
         -> QuantifiedFormula.

End Sigma_1_1_Internal.

Section Sigma_1_1_Denotation.

Variable (FSize : nat).

Record Sigma11Model : Type :=
  mkSigma11Model {
      V_F : nat -> 'F_FSize;
      F_S : nat -> { a & (|[a]| -> 'F_FSize) -> option 'F_FSize };
  }.

Definition indFun {p} (x y : 'F_p) : 'F_p := if (x < y) then 1%R else 0%R.

Theorem FinFieEq {p} {a b : 'F_p} (e : a == b = true) : a = b.
Proof.
  destruct a, b.
  apply EEConvert in e; simpl in e.
  destruct e.
  by replace i with i0;[|apply eq_irrelevance].
Qed.

Fixpoint Poly_Denote (M : Sigma11Model)
  (r : @PolyTerm (fun x => projT1 (F_S M x))) : option ('F_FSize) :=
  match r with
  | PolyVar m => Some (V_F M m)
  | PolyFun i t =>
    let ds := option_fun (fun x => Poly_Denote M (t x)) in
    obind (projT2 (F_S M i)) ds
  | PolyMinusOne => Some (-1)%R
  | PolyPlusOne => Some 1%R
  | PolyZero => Some 0%R
  | PolyPlus r1 r2 =>
    let d1 := Poly_Denote M r1 in
    let d2 := Poly_Denote M r2 in
    obind (fun r1 => obind (fun r2 => Some (r1 + r2)%R) d2) d1
  | PolyTimes r1 r2 =>
    let d1 := Poly_Denote M r1 in
    let d2 := Poly_Denote M r2 in
    obind (fun r1 => obind (fun r2 => Some (r1 * r2)%R) d2) d1
  | PolyInd r1 r2 =>
    let r1 := Poly_Denote M r1 in
    let r2 := Poly_Denote M r2 in
    obind (fun r1 => obind (fun r2 => Some (indFun r1 r2)) r2) r1
  end.

Fixpoint ZerothOrder_Denote (M : Sigma11Model)
  (f : @ZerothOrderFormula (fun x => projT1 (F_S M x))) : option bool :=
  match f with
  | ZONot f =>
    let d := ZerothOrder_Denote M f in
    obind (fun d => Some (negb d)) d
  | ZOAnd f1 f2 =>
    let d1 := ZerothOrder_Denote M f1 in
    let d2 := ZerothOrder_Denote M f2 in
    obind (fun r1 => obind (fun r2 => Some (r1 && r2)) d2) d1
  | ZOOr f1 f2 =>
    let d1 := ZerothOrder_Denote M f1 in
    let d2 := ZerothOrder_Denote M f2 in
    obind (fun r1 => obind (fun r2 => Some (r1 || r2)) d2) d1
  | ZOImp f1 f2 =>
    let d1 := ZerothOrder_Denote M f1 in
    let d2 := ZerothOrder_Denote M f2 in
    obind (fun r1 => obind (fun r2 => Some (r1 ==> r2)) d2) d1
  | ZOEq r1 r2 =>
    let d1 := Poly_Denote M r1 in
    let d2 := Poly_Denote M r2 in
    obind (fun r1 => obind (fun r2 => Some (r1 == r2)) d2) d1
  end.

Definition AddModelV (M : Sigma11Model) (r : 'F_FSize) : Sigma11Model :=
  {| V_F := ExtendAt0 r (V_F M); F_S := F_S M |}.

Definition AddModelF  (M : Sigma11Model) (f : { newA & (|[newA]| -> 'F_FSize) -> option ('F_FSize)})  :
  Sigma11Model := {| V_F := V_F M; F_S := ExtendAt0 f (F_S M) |}.

Program Fixpoint FunBounds
  (M : Sigma11Model) {a}
  (ins : |[a]| -> 'F_FSize) (out : 'F_FSize)
  (insB : |[a]| -> @PolyTerm (fun x => projT1 (F_S M x))) (outB : @PolyTerm (fun x => projT1 (F_S M x))) : bool :=
  match a with
  | 0 =>
    match Poly_Denote M outB with
    | None => false
    | Some oB => out < oB
    end
  | n.+1 =>
    match Poly_Denote M (insB 0) with
    | None => false
    | Some iB => (ins 0 < iB) && @FunBounds (AddModelV M (ins 0)) n (fun x => ins (x.+1)) out (fun x => insB (x.+1)) outB
    end
  end.

Definition Fun_Bound_Check
  (M : Sigma11Model)
  {n : nat}
  (bs : |[n]| -> @PolyTerm (fun x => projT1 (F_S M x)))
  (y : @PolyTerm (fun x => projT1 (F_S M x)))
  (f : (|[n]| -> 'F_FSize) -> option ('F_FSize)) : Prop :=
forall (ins : |[n]| -> 'F_FSize) (out : 'F_FSize),
  f ins == Some out ->
  FunBounds M ins out bs y == true.

Program Fixpoint QuantifiedFormula_Denote (M : Sigma11Model) (f : @QuantifiedFormula (fun x => projT1 (F_S M x))) : Prop :=
  match f with
  | ZO z => ZerothOrder_Denote M z == Some true
  | QExists bs y f =>
    exists (F : (|[length bs]| -> 'F_FSize) -> option ('F_FSize)),
    Fun_Bound_Check M (lnth bs) y F /\ QuantifiedFormula_Denote (AddModelF M (existT _ (length bs) F)) f
  | QForall b f =>
    match Poly_Denote M b with
    | None => False
    | Some p' => forall (r : 'F_FSize), r < p' -> QuantifiedFormula_Denote (AddModelV M r) f
    end
  end.
Next Obligation.
  apply functional_extensionality;move=> [|x]; auto.
Qed.

End Sigma_1_1_Denotation.


(*
 This is the definition of formal syntax for Dan Grossman's Thesis,
  "SAFE PROGRAMMING AT THE C LEVEL OF ABSTRACTION".

  This defines the notions of extending contexts.

*)

Require Import List.
Export ListNotations.
Require Import ZArith.
Require Import Init.Datatypes.

Require Export FormalSyntax.
Require Export TacticNotations.
Require Export CpdtTactics.
Require Export Case.
Require Export StaticSemanticsKindingAndContextWellFormedness.

Inductive ExtendedByD : Delta -> Delta -> Prop :=
  | ExtendedByD_nil   : forall (d' : Delta),
                          ExtendedByD [] d'
  | ExtendedByD_left  :
      forall (alpha : TVar) (k : Kappa) (dtail : Delta) (d' : Delta),
        getD d' alpha = Some k ->
        ExtendedByD dtail d' ->
        ExtendedByD ((alpha,k) :: dtail) d'.

(* TODO Why can't I check that if x is in D and D' it agrees? *)
(* But then I need a delete. *)
(*  | ExtendedByD_right  :
      forall (alpha : TVar) (k : Kappa) (d: Delta) (d'tail : Delta),
        getD d alpha = Some k ->
        ExtendedByD (deleteD alpha k d) d'tail ->
        ExtendedByD d ((alpha,k) :: d'tail).
*)

Inductive ExtendedByG : Gamma -> Gamma -> Prop :=
  | ExtendedByG_nil   : forall (d' : Gamma),
                          ExtendedByG [] d'
  | ExtendedByG_left  :
      forall (x : EVar) (t : Tau) (dtail : Gamma) (d' : Gamma),
        getG d' x = Some t ->
        ExtendedByG dtail d' ->
        ExtendedByG ((x,t) :: dtail) d'.


Require Import Terms.
Require Import Model.
Require Import Patterns.
Require Import Satisfaction.
Require Import Closure.
Require Import Relation_Definitions.
Require Import List.

Module Type ObjectLanguage.
  Declare Module Export Base : Terms.
  Declare Module Export BaseM : Model Base.
  Declare Module Export BaseP : Patterns Base.
  Declare Module Export BaseS : Satisfaction Base BaseM BaseP.
End ObjectLanguage.

Module Type Reachability (Export Base : ObjectLanguage).

  Definition SimpleRule := (Pattern * Pattern)%type.

  Definition SimpleSystem := list SimpleRule.

  Definition CondRule := (SimpleRule * SimpleSystem)%type.

  Definition CondSystem := CondRule -> Prop.

  Fixpoint mapFst (l : list SimpleRule) :=
    match l with
      | nil => nil
      | cons hd tl =>
        cons (fst hd) (mapFst tl)
    end.

  Fixpoint mapSnd (l : list SimpleRule) :=
    match l with
      | nil => nil
      | cons hd tl =>
        cons (snd hd) (mapFst tl)
    end.

  Inductive InTSList (R : relation M) : list M -> list M -> Prop :=
  | InTSListNil :
    InTSList R nil nil
  | InTSListCons :
    forall hd hd' tl tl',
      R hd hd' ->
      InTSList R tl tl' ->
      InTSList R (cons hd tl) (cons hd' tl').

  Inductive SatisfiesList : list M -> Valuation -> list Pattern -> Prop :=
  | SatisfiesListNil :
    forall rho,
      SatisfiesList nil rho nil
  | SatisfiesListCons :
    forall gamma gammaTail phi phiTail rho,
      Satisfies gamma rho phi ->
      SatisfiesList gammaTail rho phiTail ->
      SatisfiesList (cons gamma gammaTail) rho (cons phi phiTail).

  Inductive TS (csys : CondSystem) : relation M :=
    | TSIntro : forall (rule : CondRule),
      csys rule ->
      forall (rho : Valuation),
        forall gamma gamma',
          Satisfies gamma rho (fst (fst rule)) ->
          Satisfies gamma' rho (snd (fst rule)) ->
          (forall gammaList : list M,
            SatisfiesList gammaList rho (mapFst (snd rule)) ->
            exists gammaList' : list M,
              InTSList (clos _ (TS csys) false) gammaList gammaList') ->
          TS csys gamma gamma'.

  Definition TSStar (csys : CondSystem) : relation M :=
    clos _ (TS csys) false.

  Parameter S : CondSystem.

  Definition RhoStronglyValid (rho : Valuation) (srule : SimpleRule) :=
    forall gamma,
      Satisfies gamma rho (fst srule) ->
      exists gamma',
        TSStar S gamma gamma' /\
        Satisfies gamma' rho (snd srule).

  Definition StronglyValid (srule : SimpleRule) :=
    forall rho,
      RhoStronglyValid rho srule.

  Inductive IsSuffix { A : Type } : list A -> list A -> Prop :=
  | here :
    forall slist,
      IsSuffix slist slist
  | next :
    forall slist blist bhead,
      IsSuffix slist blist ->
      IsSuffix slist (bhead :: blist).

  Notation "x 'SuffixOf' y" := (IsSuffix x y) (at level 60, right associativity).

  Definition TerminationDependence (csys : CondSystem) (gamma gamma' : M) : Prop :=
    (TSStar csys) gamma gamma' \/
    exists rule,
      (csys rule /\
        exists rho,
          exists rest : SimpleSystem,
            exists irule : SimpleRule,
              (cons irule rest) SuffixOf (snd rule) /\
              Satisfies gamma rho (fst (fst rule)) /\
              List.Forall (fun rule => RhoStronglyValid rho rule) rest /\
              Satisfies gamma' rho (fst irule)).

  Definition Prec csys gamma gamma' :=
    (clos _ (TerminationDependence csys) true) gamma' gamma.

  Definition Terminates (csys : CondSystem) gamma :=
    Acc (Prec csys) gamma.

  Definition TerminatesPattern (csys : CondSystem) phi :=
    forall gamma rho,
      Satisfies gamma rho phi ->
      Terminates csys gamma.

  Definition WeaklyWD (rule : CondRule) :=
    WeaklyWDPattern (snd (fst rule)) /\
    List.Forall WeaklyWDPattern (mapFst (snd rule)).

  Definition WD (rule : CondRule) :=
    WDPattern (snd (fst rule)) /\
    List.Forall WDPattern (mapFst (snd rule)).

  Definition WeaklyWDSystem (csys : CondSystem) :=
    forall rule,
      csys rule ->
      WeaklyWD rule.

  Definition WDSystem (csys : CondSystem) :=
    forall rule,
      csys rule ->
      WD rule.

  Definition GR : relation M := TSStar S.

  Notation "x 'Succ' y" := (TerminationDependence S x y) (at level 60).

  Notation "y 'SuccEq' yy" := (clos M (TerminationDependence S) false y yy) (at level 60).

  Definition GoesTo strict gamma gamma' := clos M (TS S) strict gamma gamma'.

  Definition RhoValid (srule : SimpleRule) (rho : Valuation) :=
    forall gamma,
      Satisfies gamma rho (fst srule) ->
      Terminates S gamma ->
      exists gamma',
        (Satisfies gamma' rho (snd srule) /\
          GoesTo false gamma gamma').

  Definition Valid (srule : SimpleRule) :=
    forall rho,
      RhoValid srule rho.

End Reachability.


Require Export ARS.
Require Export Prelims Base. (* Shared.Extra.Bijection. *)

Hint Constructors ARS.star : cbv.

(** * L: Weak Call-by-value Lambda Calculus*)

Inductive term : Type :=
| var (n : nat) : term
| app (s : term) (t : term) : term
| lam (s : term).

Coercion app : term >-> Funclass.
Coercion var : nat >-> term.

Notation "(λ  s )" := (lam s) (right associativity, at level 0).

Instance term_eq_dec : eq_dec term.
Proof.
  intros s t; unfold dec; repeat decide equality.
Defined.

Fixpoint size s :nat :=
  match s with
    var x => 1 + x
  | lam s => 1 +(size s)
  | app s t => 1 + size s + size t
  end.

(** ** Substitution and Closedness *)

Fixpoint subst (s : term) (k : nat) (u : term) :=
  match s with
      | var n => if Dec (n = k) then u else (var n)
      | app s t => app (subst s k u) (subst t k u)
      | lam s => lam (subst s (S k) u)
  end.

Inductive bound : nat -> term -> Prop :=
  | bndvar k n : k > n -> bound k (var n)
  | bndApp k s t : bound k s -> bound k t -> bound k (s t)
  | bndlam k s : bound (S k) s -> bound k (lam s).
Hint Constructors bound.

Definition closed s:= bound 0 s.

Hint Unfold closed.

Definition lambda s := exists t, s = lam t.

Lemma lambda_lam s : lambda (lam s).
Proof.
  exists s; reflexivity.
Qed.

Hint Resolve lambda_lam.

Lemma bound_closed_k s k u : bound k s -> subst s k u = s.
Proof with eauto.
  intros H; revert u; induction H; intros u; simpl.
  - decide (n = k)... omega.
  - rewrite IHbound1, IHbound2...
  - f_equal...
Qed.

Lemma bound_ge k s : bound k s -> forall m, m >= k -> bound m s.
Proof.
  induction 1; intros m Hmk; econstructor; eauto.
  - omega.
  - eapply IHbound. omega.
Qed.

Lemma bound_closed s k u : bound 0 s -> subst s k u = s.
Proof.
  intros H. destruct k.
  - eapply bound_closed_k. eassumption.
  - eapply bound_ge in H. eapply bound_closed_k. eassumption. omega.
Qed.

Lemma closed_k_bound k s : (forall n u, n >= k -> subst s n u = s) -> bound k s.
Proof.
  revert k. induction s; intros k H.
  - econstructor. specialize (H n (S n)). simpl in H.
    decide (n >= k) as [Heq | Heq].
    + decide (n = n) ; [injection (H Heq)|]; omega.
    + omega.
  - econstructor; [eapply IHs1 | eapply IHs2]; intros n u Hnk;
    injection (H n u Hnk); congruence.
  - econstructor. eapply IHs. intros n u Hnk.
    destruct n. omega.
    injection (H n u). tauto. omega.
Qed.

Lemma closed_subst_iff s : closed s <-> (forall k t, subst s k t = s).
Proof.
  split.
  -intros. rewrite bound_closed_k. tauto. eapply bound_ge;eauto. omega.
  -intros. apply closed_k_bound. eauto.
Qed.

Lemma closed_subst s : closed s -> (forall k t, subst s k t = s).
Proof.
  apply closed_subst_iff.
Qed.

Lemma closed_app (s t : term) : closed (s t) -> closed s /\ closed t.
Proof.
  intros H. now inv H.
Qed.

Lemma app_closed (s t : term) : closed s -> closed t -> closed (s t).
Proof.
  unfold closed. eauto using bound.
Qed.

Lemma bound_subst s t k : bound (S k) s -> bound k t -> bound k (subst s k t).
Proof.
  revert k t; induction s; intros k t cls_s cls_t; simpl; inv cls_s; eauto 6 using bound, bound_ge.
  decide (n = k); eauto. econstructor.  omega.
Qed.

Lemma closed_subst2 s t: closed (lam s) -> closed t -> closed (subst s 0 t).
Proof.
  intros cs ct. inv cs.  now apply bound_subst.
Qed.

(** ** Deterministic Reduction *)

Reserved Notation "s '>>' t" (at level 50).

Inductive step : term -> term -> Prop :=
| stepApp  s t     : app (lam s) (lam t) >> subst s 0 (lam t)
| stepAppR s t  t' : t >> t' -> app (lam s) t >> app (lam s) t'
| stepAppL s s' t  : s >> s' -> app s t >> app s' t
where "s '>>' t" := (step s t).
Hint Constructors step.

Ltac inv_step :=
  match goal with
    | H : step (lam _) _ |- _ => inv H
    | H : step (var _) _ |- _ => inv H
    | H : star step (lam _) _ |- _ => inv H
    | H : star step (var _) _ |- _ => inv H
  end.

(** ** Resource Measures *)

(** *** Small-Step Time Measure *)

Instance pow_step_congL k:
  Proper ((pow step k) ==> eq ==> (pow step k)) app.
Proof.
  intros s t R u ? <-. revert s t R u.
  induction k;cbn in *;intros ? ? R ?. congruence. destruct R as [s' [R1 R2]].
  exists (s' u). firstorder.
Defined.

Instance pow_step_congR k:
  Proper (eq ==>(pow step k) ==> (pow step k)) (fun s t => app (lam s) t).
Proof.
  intros s ? <- t u R. revert s t u R.
  induction k;cbn in *;intros ? ? ? R. congruence. destruct R as [t' [R1 R2]].
  exists (lam s t'). firstorder.
Defined.

(** *** Small-Step Space Measure *)

Definition redWithMaxSizeL := redWithMaxSize size step.

Lemma redWithMaxSizeL_congL m m' s s' t:
  redWithMaxSizeL m s s' -> m' = (1 + m + size t) -> redWithMaxSizeL m' (s t) (s' t).
Proof.
  intros R Heq.
  induction R as [s | s s'] in m',t,Heq |-* .
  -apply redWithMaxSizeR. cbn. omega.
  -eapply redWithMaxSizeC. now eauto using step. apply IHR. reflexivity.
   subst m m'. cbn. repeat eapply Nat.max_case_strong;omega.
Qed.

Lemma redWithMaxSizeL_congR m m' s t t':
  redWithMaxSizeL m t t' -> m' = (1 + m + size (lam s)) -> redWithMaxSizeL m' (lam s t) (lam s t').
Proof.
  intros R Heq.
  induction R as [t | t t'] in m',s,Heq |-* .
  -apply redWithMaxSizeR. cbn in *. omega.
  -eapply redWithMaxSizeC. now eauto using step. apply IHR. reflexivity.
   subst m m'. cbn -[plus]. repeat eapply Nat.max_case_strong;omega.
Defined.

(** *** Big-Step Time Measure*)

Inductive timeBS : nat -> term -> term -> Prop :=
  timeBSVal s : timeBS 0 (λ s) (λ s)
| timeBSBeta (s s' t t' u : term) i j k l: timeBS i s (λ s') -> timeBS j t (λ t') -> timeBS k (subst s' 0 (λ t')) u -> l = (i+j+1+k) -> timeBS l (s t) u.

Lemma step_evaluatesIn s s' t k: s >> s' -> timeBS k s' t -> timeBS (S k) s t.
Proof.
  intros H; induction H in t,k|-*; intros;  try inv H0; eauto 20 using timeBS.
  eapply IHstep in H4. econstructor; eauto. omega.
Qed.

Lemma timeBS_correct s t k:
  (timeBS k s t <-> pow step k s t /\ lambda t).
Proof.
  split.
  -intros R.
   induction R.
   +unfold pow;cbn. eauto.
   +destruct IHR1 as [R1'].
    destruct IHR2 as [R2'].
    destruct IHR3 as [R3'].
    split;[|assumption].
    subst l.
    replace (i+j+1+k) with (i+(j+(1+k))) by omega.
    eapply pow_add.
    eexists;split.
    eapply pow_step_congL;now eauto.
    eapply pow_add.
    eexists;split.
    eapply pow_step_congR;now eauto.
    eapply pow_add.
    eexists;split.
    eapply (rcomp_1 step). eauto.
    eauto.
  -intros [R lt].
   induction k in s,t,R,lt,R|-*.
   +inv R. inv lt. constructor.
   +change (S k) with (1+k) in R. eapply pow_add in R as (?&R'&?).
    eapply (rcomp_1 step) in R'. eapply step_evaluatesIn;eauto 10.
Qed.

(** *** Big-Step Space Measure*)

Inductive spaceBS : nat -> term -> term -> Prop :=
  spaceBSVal s : spaceBS (size (λ s)) (λ s) (λ s)
| spaceBSBeta (s s' t t' u : term) m1 m2 m3 m:
    spaceBS m1 s (λ s') ->
    spaceBS m2 t (λ t') ->
    spaceBS m3 (subst s' 0 (λ t')) u ->
    m = max (m1 + 1 + size t)
            (max (size (lam s') + 1 + m2) m3) ->
    spaceBS m (s t) u.

Lemma spaceBS_ge s t m: spaceBS m s t -> size s<= m /\ size t <= m.
Proof.
  induction 1. omega.
  subst m. cbn -[plus]. all:(repeat apply Nat.max_case_strong;try omega).
Qed.

Lemma step_evaluatesSpace s s' t m: s >> s' -> spaceBS m s' t -> spaceBS (max (size s) m) s t.
Proof.
  intros H; induction H in t,m|-*; intros H'.
  -inv H'.
   +destruct s;inv H1. decide _;inv H0.
    all:repeat econstructor.
    all:cbn -[plus].
    all:(repeat apply Nat.max_case_strong;try omega).
   +destruct s;inv H. now decide _.
    econstructor. 1,2:econstructor. cbn.
    econstructor.
    1-4:now eauto.
    cbn -[plus].
    (repeat apply Nat.max_case_strong;intros;try omega).
  -inv H'. inv H2.
   specialize (spaceBS_ge H3) as [? ?].
   apply IHstep in H3.
   specialize (spaceBS_ge H5) as [? ?].
   specialize (spaceBS_ge H3) as [_ H7].
   econstructor;[now eauto using spaceBS..|].
   revert H7. cbn -[plus] in *.
   (repeat apply Nat.max_case_strong;intros;try omega).
  -inv H'.
   specialize (spaceBS_ge H2) as [? ?].
   eapply IHstep in H2.
   specialize (spaceBS_ge H2) as [_ H7].
   specialize (spaceBS_ge H3) as [? ?].

   econstructor.
   1-3:eassumption.
   revert H7.
   cbn -[plus] in *.
   (repeat apply Nat.max_case_strong;intros;try omega).
Qed.

Lemma spaceBS_correct s t m:
  spaceBS m s t <->
   redWithMaxSizeL m s t /\ lambda t.
Proof.
  split.
  -intros R. induction R.
   +repeat econstructor.
   +destruct IHR1 as (R1'&?).
    destruct IHR2 as (R2'&?).
    destruct IHR3 as (R3'&?).
    split;[|firstorder].
    subst m.
    eapply redWithMaxSize_trans with (t:=(lam s') t).
    { eapply redWithMaxSizeL_congL. eassumption. reflexivity. }

    eapply redWithMaxSize_trans with (t:=(lam s') (lam t')).
    { eapply redWithMaxSizeL_congR. eassumption. reflexivity. }
    econstructor. constructor. eassumption. reflexivity. reflexivity.
    specialize (spaceBS_ge R2) as [_ H3];cbn in H3.
    cbn - [plus]. repeat eapply Nat.max_case_strong;omega.
  -intros (R&H).
   induction R as [t | s s' t m].
   +inv H;rename x into t. eauto using spaceBS.
   +inv H;rename m' into m. eapply step_evaluatesSpace. eassumption. eauto.
Qed.