Require Import rt.util.all.
Require Import rt.model.arrival.basic.task rt.model.arrival.basic.job rt.model.priority rt.model.arrival.basic.task_arrival.
Require Import rt.model.schedule.global.basic.schedule rt.model.schedule.global.basic.interference
               rt.model.schedule.global.basic.platform.
From mathcomp Require Import ssreflect ssrbool ssrfun eqtype ssrnat seq fintype bigop.

Module ConstrainedDeadlines.

  Import Job SporadicTaskset ScheduleOfSporadicTask SporadicTaskset
         TaskArrival Interference Priority Platform.

  Section Lemmas.

    Context {sporadic_task: eqType}.
    Variable task_cost: sporadic_task -> time.
    Variable task_period: sporadic_task -> time.
    Variable task_deadline: sporadic_task -> time.

    Context {Job: eqType}.
    Variable job_arrival: Job -> time.
    Variable job_cost: Job -> time.
    Variable job_deadline: Job -> time.
    Variable job_task: Job -> sporadic_task.

    (* Assume any job arrival sequence ... *)
    Variable arr_seq: arrival_sequence Job.

    (* ... and any schedule of this arrival sequence. *)
    Context {num_cpus: nat}.
    Variable sched: schedule Job num_cpus.
    Hypothesis H_jobs_come_from_arrival_sequence:
      jobs_come_from_arrival_sequence sched arr_seq.

    (* Assume all jobs have valid parameters, ...*)
    Hypothesis H_valid_job_parameters:
      forall j,
        arrives_in arr_seq j ->
        valid_sporadic_job task_cost task_deadline job_cost job_deadline job_task j.

    (* In this section we prove the absence of multiple jobs of the same
       task when constrained deadlines are assumed.  *)
    Section NoMultipleJobs.

      (* Assume any work-conserving priority-based scheduler. *)
      Variable higher_eq_priority: JLDP_policy Job.
      Hypothesis H_work_conserving: work_conserving job_arrival job_cost arr_seq sched.
      Hypothesis H_respects_JLDP_policy:
        respects_JLDP_policy job_arrival job_cost arr_seq sched higher_eq_priority.

      (* Consider task set ts. *)
      Variable ts: taskset_of sporadic_task.

      (* Assume that all jobs come from the taskset. *)
      Hypothesis H_all_jobs_from_taskset:
        forall j,
          arrives_in arr_seq j -> job_task j \in ts.

      (* Suppose that jobs are sequential, ...*)
      Hypothesis H_sequential_jobs: sequential_jobs sched.
      (* ... jobs must arrive to execute, ... *)
      Hypothesis H_completed_jobs_dont_execute:
        completed_jobs_dont_execute job_cost sched.
      (* ... and jobs do not execute after completion. *)
      Hypothesis H_jobs_must_arrive_to_execute:
        jobs_must_arrive_to_execute job_arrival sched.

      (* Assume that the schedule satisfies the sporadic task model ...*)
      Hypothesis H_sporadic_tasks:
        sporadic_task_model task_period job_arrival job_task arr_seq.

      (* Consider a valid task tsk, ...*)
      Variable tsk: sporadic_task.
      Hypothesis H_valid_task: is_valid_sporadic_task task_cost task_period task_deadline tsk.

      (*... whose job j ... *)
      Variable j: Job.
      Hypothesis H_j_arrives: arrives_in arr_seq j.
      Variable H_job_of_tsk: job_task j = tsk.

      (*... is backlogged at time t. *)
      Variable t: time.
      Hypothesis H_j_backlogged: backlogged job_arrival job_cost sched j t.

      (* Assume that any previous jobs of tsk have completed by the period. *)
      Hypothesis H_all_previous_jobs_completed :
        forall j_other tsk_other,
          arrives_in arr_seq j_other ->
          job_task j_other = tsk_other ->
          job_arrival j_other + task_period tsk_other <= t ->
          completed job_cost sched j_other (job_arrival j_other + task_period (job_task j_other)).

      Let scheduled_task_other_than (tsk tsk_other: sporadic_task) :=
        task_is_scheduled job_task sched tsk_other t && (tsk_other != tsk).

      (* Then, there can be at most one pending job of each task at time t. *)
      Lemma platform_at_most_one_pending_job_of_each_task :
        forall j1 j2,
          arrives_in arr_seq j1 ->
          arrives_in arr_seq j2 ->
          pending job_arrival job_cost sched j1 t ->
          pending job_arrival job_cost sched j2 t ->
          job_task j1 = job_task j2 ->
          j1 = j2.
      Proof.
        rename H_sporadic_tasks into SPO, H_all_previous_jobs_completed into PREV.
        intros j1 j2 ARR1 ARR2 PENDING1 PENDING2 SAMEtsk.
        apply/eqP; rewrite -[_ == _]negbK; apply/negP; red; move => /eqP DIFF.
        move: PENDING1 PENDING2 => /andP [ARRIVED1 /negP NOTCOMP1] /andP [ARRIVED2 /negP NOTCOMP2].
        destruct (leqP (job_arrival j1) (job_arrival j2)) as [BEFORE1 | BEFORE2].
        {
          specialize (SPO j1 j2 DIFF ARR1 ARR2 SAMEtsk BEFORE1).
          exploit (PREV j1 (job_task j1) ARR1);
            [by done | by apply leq_trans with (n := job_arrival j2) | intros COMP1].
          apply NOTCOMP1.
          apply completion_monotonic with (t0 := job_arrival j1 + task_period (job_task j1));
            try (by done).
          by apply leq_trans with (n := job_arrival j2).
        }
        {
          apply ltnW in BEFORE2.
          exploit (SPO j2 j1); try (by done); [by red; ins; subst | intro SPO'].
          exploit (PREV j2 (job_task j2) ARR2);
            [by done | by apply leq_trans with (n := job_arrival j1) | intros COMP2].
          apply NOTCOMP2.
          apply completion_monotonic with (t0 := job_arrival j2 + task_period (job_task j2));
            try (by done).
          by apply leq_trans with (n := job_arrival j1).
        }
      Qed.

      (* Therefore, all processors are busy with tasks other than tsk. *)
      Lemma platform_cpus_busy_with_interfering_tasks :
        count (scheduled_task_other_than tsk) ts = num_cpus.
      Proof.
        have UNIQ := platform_at_most_one_pending_job_of_each_task.
        rename H_all_jobs_from_taskset into FROMTS,
               H_sequential_jobs into SEQUENTIAL, H_work_conserving into WORK,
               H_respects_JLDP_policy into PRIO,
               H_j_backlogged into BACK, H_jobs_come_from_arrival_sequence into FROMarr,
               H_job_of_tsk into JOBtsk, H_valid_job_parameters into JOBPARAMS,
               H_valid_task into TASKPARAMS,
               H_all_previous_jobs_completed into PREV, H_completed_jobs_dont_execute into COMP,
               H_jobs_must_arrive_to_execute into ARRIVE.
        apply work_conserving_eq_work_conserving_count in WORK.
        unfold valid_sporadic_job, valid_realtime_job,
               respects_JLDP_policy, completed_jobs_dont_execute,
               sporadic_task_model, is_valid_sporadic_task,
               jobs_of_same_task_dont_execute_in_parallel,
               sequential_jobs in *.
        apply/eqP; rewrite eqn_leq; apply/andP; split.
        {
          apply leq_trans with (n := count (fun x => task_is_scheduled job_task sched x t) ts);
            first by apply sub_count; first by red; move => x /andP [SCHED _].
          unfold task_is_scheduled.
          apply count_exists; first by destruct ts.
          {
            intros cpu x1 x2 SCHED1 SCHED2.
            unfold task_scheduled_on in *.
            destruct (sched cpu t); last by done.
            move: SCHED1 SCHED2 => /eqP SCHED1 /eqP SCHED2.
            by rewrite -SCHED1 -SCHED2.
          }
        }
        {
          rewrite -(WORK j t) // -count_predT.
          apply leq_trans with (n := count (fun j => scheduled_task_other_than tsk (job_task j)) (jobs_scheduled_at sched t));
            last first.
          {
            rewrite -count_map.
            apply count_sub_uniqr; last first.
            {
              move => tsk' /mapP [j' INj' JOBtsk']; subst; apply FROMTS.
              by rewrite mem_scheduled_jobs_eq_scheduled in INj'; eauto 2.
            }
            rewrite map_inj_in_uniq; first by apply scheduled_jobs_uniq.
            red; intros j1 j2 SCHED1 SCHED2 SAMEtsk.
            rewrite 2!mem_scheduled_jobs_eq_scheduled in SCHED1 SCHED2.
            have ARRin1: arrives_in arr_seq j1 by apply (FROMarr j1 t).
              have ARRin2: arrives_in arr_seq j2 by apply (FROMarr j2 t).
              by apply UNIQ; try (by done); apply scheduled_implies_pending.
          }
          {
            apply sub_in_count; intros j' SCHED' _.
            rewrite mem_scheduled_jobs_eq_scheduled in SCHED'.
            unfold scheduled_task_other_than; apply/andP; split.
            {
              move: SCHED' => /existsP [cpu /eqP SCHED'].
              by apply/existsP; exists cpu; rewrite /task_scheduled_on SCHED' eq_refl.
            }
            {
              apply/eqP; red; intro SAMEtsk; symmetry in SAMEtsk.
              move: BACK => /andP [PENDING NOTSCHED].
              have ARRin': arrives_in arr_seq j' by apply (FROMarr j' t).
              exploit (UNIQ j j'); try (by done);
                [by apply scheduled_implies_pending | by rewrite -SAMEtsk |].
              by intro EQjob; subst; rewrite SCHED' in NOTSCHED.
            }
          }
        }
      Qed.

    End NoMultipleJobs.

    (* In this section we also prove the absence of multiple jobs of the same
       task when constrained deadlines are assumed, but in the specific case
       of fixed-priority scheduling.  *)
    Section NoMultipleJobsFP.

      (* Assume any work-conserving priority-based scheduler. *)
      Variable higher_eq_priority: FP_policy sporadic_task.
      Hypothesis H_work_conserving: work_conserving job_arrival job_cost arr_seq sched.
      Hypothesis H_respects_JLDP_policy:
        respects_FP_policy job_arrival job_cost job_task arr_seq sched higher_eq_priority.

      (* Consider any task set ts. *)
      Variable ts: taskset_of sporadic_task.

      (* Assume that all jobs come from the taskset. *)
      Hypothesis H_all_jobs_from_taskset:
        forall j,
          arrives_in arr_seq j -> job_task j \in ts.

      (* Suppose that jobs are sequential, ...*)
      Hypothesis H_sequential_jobs: sequential_jobs sched.
      (* ... jobs must arrive to execute, ... *)
      Hypothesis H_completed_jobs_dont_execute:
        completed_jobs_dont_execute job_cost sched.
      (* ... and jobs do not execute after completion. *)
      Hypothesis H_jobs_must_arrive_to_execute:
        jobs_must_arrive_to_execute job_arrival sched.

      (* Assume that jobs arrive sporadically. *)
      Hypothesis H_sporadic_tasks:
        sporadic_task_model task_period job_arrival job_task arr_seq.

      (* Consider a valid task tsk, ...*)
      Variable tsk: sporadic_task.
      Hypothesis H_valid_task: is_valid_sporadic_task task_cost task_period task_deadline tsk.

      (*... whose job j ... *)
      Variable j: Job.
      Hypothesis H_j_arrives: arrives_in arr_seq j.
      Variable H_job_of_tsk: job_task j = tsk.

      (*... is backlogged at time t <= job_arrival j + task_period tsk. *)
      Variable t: time.
      Hypothesis H_j_backlogged: backlogged job_arrival job_cost sched j t.
      Hypothesis H_t_before_period: t < job_arrival j + task_period tsk.

      (* Recall the definition of a higher-priority task (with respect to tsk). *)
      Let is_hp_task := higher_priority_task higher_eq_priority tsk.

      (* Assume that any jobs of higher-priority tasks complete by their period. *)
      Hypothesis H_all_previous_jobs_completed :
        forall j_other tsk_other,
          arrives_in arr_seq j_other ->
          job_task j_other = tsk_other ->
          is_hp_task tsk_other ->
          completed job_cost sched j_other (job_arrival j_other + task_period tsk_other).

      (* Assume that any jobs of tsk prior to j complete by their period. *)
      Hypothesis H_all_previous_jobs_of_tsk_completed :
        forall j0,
          arrives_in arr_seq j0 ->
          job_task j0 = tsk ->
          job_arrival j0 < job_arrival j ->
          completed job_cost sched j0 (job_arrival j0 + task_period tsk).

      Definition scheduled_task_with_higher_eq_priority (tsk tsk_other: sporadic_task) :=
        task_is_scheduled job_task sched tsk_other t &&
        is_hp_task tsk_other.

      (* Then, there can be at most one pending job of higher-priority tasks at time t. *)
      Lemma platform_fp_no_multiple_jobs_of_interfering_tasks :
        forall j1 j2,
          arrives_in arr_seq j1 ->
          arrives_in arr_seq j2 ->
          pending job_arrival job_cost sched j1 t ->
          pending job_arrival job_cost sched j2 t ->
          job_task j1 = job_task j2 ->
          is_hp_task (job_task j1) ->
          j1 = j2.
      Proof.
        unfold sporadic_task_model in *.
        rename H_sporadic_tasks into SPO, H_all_previous_jobs_of_tsk_completed into PREVtsk,
               H_all_previous_jobs_completed into PREV.
        intros j1 j2 ARR1 ARR2 PENDING1 PENDING2 SAMEtsk INTERF.
        apply/eqP; rewrite -[_ == _]negbK; apply/negP; red; move => /eqP DIFF.
        move: PENDING1 PENDING2 => /andP [ARRIVED1 /negP NOTCOMP1] /andP [ARRIVED2 /negP NOTCOMP2].
        destruct (leqP (job_arrival j1) (job_arrival j2)) as [BEFORE1 | BEFORE2].
        {
          specialize (SPO j1 j2 DIFF ARR1 ARR2 SAMEtsk BEFORE1).
          exploit (PREV j1 (job_task j1) ARR1); [by done | by apply INTERF | intros COMP1].
          apply NOTCOMP1.
          apply completion_monotonic with (t0 := job_arrival j1 + task_period (job_task j1));
            try (by done).
          by apply leq_trans with (n := job_arrival j2).
        }
        {
          apply ltnW in BEFORE2.
          exploit (SPO j2 j1); try (by done); [by red; ins; subst j2 | intro SPO'].
          exploit (PREV j2 (job_task j2) ARR2);
            [by done | by rewrite -SAMEtsk | intro COMP2 ].
          apply NOTCOMP2.
          apply completion_monotonic with (t0 := job_arrival j2 + task_period (job_task j2));
            try (by done).
          by apply leq_trans with (n := job_arrival j1).
        }
      Qed.

      (* Also, there can be at most one pending job of tsk at time t. *)
      Lemma platform_fp_no_multiple_jobs_of_tsk :
        forall j',
          arrives_in arr_seq j' ->
          pending job_arrival job_cost sched j' t ->
          job_task j' = tsk ->
          j' = j.
      Proof.
        unfold sporadic_task_model in *.
        rename H_sporadic_tasks into SPO,
               H_valid_task into PARAMS,
               H_all_previous_jobs_of_tsk_completed into PREVtsk,
               H_all_previous_jobs_completed into PREV,
               H_j_backlogged into BACK, H_job_of_tsk into JOBtsk.
        intros j' ARR' PENDING' SAMEtsk.
        apply/eqP; rewrite -[_ == _]negbK; apply/negP; red; move => /eqP DIFF.
        move: BACK PENDING' => /andP [/andP [ARRIVED /negP NOTCOMP] NOTSCHED]
                               /andP [ARRIVED' /negP NOTCOMP'].
        destruct (leqP (job_arrival j') (job_arrival j)) as [BEFORE | BEFORE'].
        {
          exploit (SPO j' j DIFF ARR' H_j_arrives); [by rewrite JOBtsk | by done | intro SPO'].
          apply NOTCOMP'.
          apply completion_monotonic with (t0 := job_arrival j' + task_period tsk); try (by done);
            first by apply leq_trans with (n := job_arrival j); [by rewrite -SAMEtsk | by done].
          {
            apply PREVtsk; try (by done).
            apply leq_trans with (n := job_arrival j' + task_period tsk); last by rewrite -SAMEtsk.
            rewrite -addn1; apply leq_add; first by done.
            by unfold is_valid_sporadic_task in *; des.
          }
        }
        {
          unfold has_arrived in *.
          rewrite leqNgt in ARRIVED'; move: ARRIVED' => /negP BUG; apply BUG.
          apply leq_trans with (n := job_arrival j + task_period tsk); first by done.
          by rewrite -JOBtsk; apply SPO; try (by done);
            [by red; ins; subst j' | by rewrite SAMEtsk | by apply ltnW].
        }
      Qed.

      (* Therefore, all processors are busy with tasks other than tsk. *)
      Lemma platform_fp_cpus_busy_with_interfering_tasks :
        count (scheduled_task_with_higher_eq_priority tsk) ts = num_cpus.
      Proof.
        have UNIQ := platform_fp_no_multiple_jobs_of_interfering_tasks.
        have UNIQ' := platform_fp_no_multiple_jobs_of_tsk.
        rename H_all_jobs_from_taskset into FROMTS, H_sequential_jobs into SEQUENTIAL,
               H_work_conserving into WORK, H_respects_JLDP_policy into PRIO,
               H_j_backlogged into BACK, H_job_of_tsk into JOBtsk,
               H_sporadic_tasks into SPO, H_valid_job_parameters into JOBPARAMS,
               H_valid_task into TASKPARAMS, H_all_previous_jobs_completed into PREV,
               H_completed_jobs_dont_execute into COMP, H_jobs_come_from_arrival_sequence into FROMarr,
               H_all_previous_jobs_of_tsk_completed into PREVtsk,
               H_jobs_must_arrive_to_execute into ARRIVE.
        apply work_conserving_eq_work_conserving_count in WORK.
        unfold valid_sporadic_job, valid_realtime_job,
               respects_FP_policy, respects_JLDP_policy, FP_to_JLDP,
               completed_jobs_dont_execute, sequential_jobs,
               sporadic_task_model, is_valid_sporadic_task,
               jobs_of_same_task_dont_execute_in_parallel,
               is_hp_task in *.
        apply/eqP; rewrite eqn_leq; apply/andP; split.
        {
          apply leq_trans with (n := count (fun x => task_is_scheduled job_task sched x t) ts);
            first by apply sub_count; red; move => x /andP [SCHED _].
          unfold task_is_scheduled.
          apply count_exists; first by destruct ts.
          {
            intros cpu x1 x2 SCHED1 SCHED2.
            unfold task_scheduled_on in *.
            destruct (sched cpu t); last by done.
            move: SCHED1 SCHED2 => /eqP SCHED1 /eqP SCHED2.
            by rewrite -SCHED1 -SCHED2.
          }
        }
        {
          rewrite -(WORK j t) // -count_predT.
          apply leq_trans with (n := count (fun j =>
            scheduled_task_with_higher_eq_priority tsk (job_task j)) (jobs_scheduled_at sched t));
            last first.
          {
            rewrite -count_map.
            apply leq_trans with (n := count predT [seq x <- (map (fun j => job_task j) (jobs_scheduled_at sched t)) | scheduled_task_with_higher_eq_priority tsk x]);
              first  by rewrite count_filter; apply sub_count; red; ins.
            apply leq_trans with (n := count predT [seq x <- ts | scheduled_task_with_higher_eq_priority tsk x]);
              last by rewrite count_predT size_filter.
            apply count_sub_uniqr; last first.
            {
              red; intros tsk' IN'.
              rewrite mem_filter in IN'; move: IN' => /andP [SCHED IN'].
              rewrite mem_filter; apply/andP; split; first by done.
              move: IN' => /mapP [j' IN'] ->; apply FROMTS.
              by rewrite mem_scheduled_jobs_eq_scheduled in IN'; eauto 2.
            }
            {
              rewrite filter_map.
              rewrite map_inj_in_uniq; first by apply filter_uniq, scheduled_jobs_uniq.
              red; intros j1 j2 SCHED1 SCHED2 SAMEtsk.
              rewrite 2!mem_filter in SCHED1 SCHED2.
              move: SCHED1 SCHED2 => /andP [/andP [_ HP1] SCHED1] /andP [/andP [_ HP2] SCHED2].
              rewrite 2!mem_scheduled_jobs_eq_scheduled in SCHED1 SCHED2.
              have ARRin1: arrives_in arr_seq j1 by apply (FROMarr j1 t).
              have ARRIn2: arrives_in arr_seq j2 by apply (FROMarr j2 t).
              by apply UNIQ; try (by done); apply scheduled_implies_pending.
            }
          }
          {
            apply sub_in_count; intros j' SCHED' _.
            rewrite mem_scheduled_jobs_eq_scheduled in SCHED'.
            apply/andP; split.
            {
              move: SCHED' => /existsP [cpu /eqP SCHED'].
              by apply/existsP; exists cpu; rewrite /task_scheduled_on SCHED' eq_refl.
            }
            apply/andP; split; first by rewrite -JOBtsk; apply PRIO with (t := t).
            {
              apply/eqP; red; intro SAMEtsk.
              generalize SCHED'; intro PENDING'.
              have ARRin': arrives_in arr_seq j' by apply (FROMarr j' t).
              apply scheduled_implies_pending with (job_arrival0 := job_arrival)
                (job_cost0 := job_cost) in PENDING'; try (by done).
              specialize (UNIQ' j' ARRin' PENDING' SAMEtsk); subst j'.
              by move: BACK => /andP [_ NOTSCHED]; rewrite SCHED' in NOTSCHED.
            }
          }
        }
      Qed.

    End NoMultipleJobsFP.

  End Lemmas.

End ConstrainedDeadlines.


Theorem no_intermediate: ~ exists p : Prop, ~ (p <-> True) /\ ~ (p <-> False).
Proof.
  unfold "~". intros H.
  inversion H.
  inversion H0.
  apply H2. unfold "<->".
  split.
  (* left *)
  intros. apply H1.
  split. (* left *)
         intros. exact I.
         (* right *)
         intros. exact H3.
  (* right *)
  intros.
  inversion H3.
Qed.


Require Import Coq.NArith.NArith.
Require Import Coq.Arith.Arith.
Require Import Coq.ZArith.ZArith.
Require Import Coq.Lists.List.
Require Import Coq.Sorting.Permutation.
Require Import Coq.Setoids.Setoid.
Require Import Coq.Classes.Morphisms.
Require Import Crypto.Util.Tactics.BreakMatch.
Import ListNotations.

Lemma Permutation_partition {A} (l : list A) f :
    Permutation.Permutation l (fst (partition f l) ++ snd (partition f l)).
Proof using Type.
  induction l; cbn; break_match; cbn in *; eauto.
  etransitivity. econstructor. eassumption.
  eapply Permutation.Permutation_middle.
Qed.

Global Instance fold_right_Proper_commutative_associative_Permutation {A op}
       (Hcomm : forall x y, op x y = op y x)
       (Hassoc : forall x y z, op x (op y z) = op (op x y) z)
  : Proper (eq ==> @Permutation A ==> eq) (fold_right op) | 1000.
Proof using Type.
  intros init init' <- xs ys; induction 1; cbn;
    rewrite ?Hassoc;
    repeat match goal with
           | [ |- context[op ?x ?y] ]
             => lazymatch goal with
                | [ |- context[op y x] ]
                  => rewrite (Hcomm x y)
                end
           end;
    try reflexivity;
    try now repeat match goal with H : _ |- _ => rewrite H; clear H end.
Qed.

Module Nat.
  Global Instance fold_right_Proper_Permutation_add
  : Proper (eq ==> @Permutation _ ==> eq) (fold_right Nat.add) | 100.
  Proof using Type. apply fold_right_Proper_commutative_associative_Permutation; [ hnf; apply Nat.add_comm | hnf; apply Nat.add_assoc ]. Qed.

  Global Instance fold_right_Proper_Permutation_mul
  : Proper (eq ==> @Permutation _ ==> eq) (fold_right Nat.mul) | 100.
  Proof using Type. apply fold_right_Proper_commutative_associative_Permutation; [ hnf; apply Nat.mul_comm | hnf; apply Nat.mul_assoc ]. Qed.

  Global Instance fold_right_Proper_Permutation_land
  : Proper (eq ==> @Permutation _ ==> eq) (fold_right Nat.land) | 100.
  Proof using Type. apply fold_right_Proper_commutative_associative_Permutation; [ hnf; apply Nat.land_comm | hnf; apply Nat.land_assoc ]. Qed.

  Global Instance fold_right_Proper_Permutation_lor
  : Proper (eq ==> @Permutation _ ==> eq) (fold_right Nat.lor) | 100.
  Proof using Type. apply fold_right_Proper_commutative_associative_Permutation; [ hnf; apply Nat.lor_comm | hnf; apply Nat.lor_assoc ]. Qed.

  Global Instance fold_right_Proper_Permutation_lxor
  : Proper (eq ==> @Permutation _ ==> eq) (fold_right Nat.lxor) | 100.
  Proof using Type. apply fold_right_Proper_commutative_associative_Permutation; [ hnf; apply Nat.lxor_comm | hnf; intros; symmetry; apply Nat.lxor_assoc ]. Qed.
End Nat.

Module N.
  Global Instance fold_right_Proper_Permutation_add
  : Proper (eq ==> @Permutation _ ==> eq) (fold_right N.add) | 100.
  Proof using Type. apply fold_right_Proper_commutative_associative_Permutation; [ hnf; apply N.add_comm | hnf; apply N.add_assoc ]. Qed.

  Global Instance fold_right_Proper_Permutation_mul
  : Proper (eq ==> @Permutation _ ==> eq) (fold_right N.mul) | 100.
  Proof using Type. apply fold_right_Proper_commutative_associative_Permutation; [ hnf; apply N.mul_comm | hnf; apply N.mul_assoc ]. Qed.

  Global Instance fold_right_Proper_Permutation_land
  : Proper (eq ==> @Permutation _ ==> eq) (fold_right N.land) | 100.
  Proof using Type. apply fold_right_Proper_commutative_associative_Permutation; [ hnf; apply N.land_comm | hnf; apply N.land_assoc ]. Qed.

  Global Instance fold_right_Proper_Permutation_lor
  : Proper (eq ==> @Permutation _ ==> eq) (fold_right N.lor) | 100.
  Proof using Type. apply fold_right_Proper_commutative_associative_Permutation; [ hnf; apply N.lor_comm | hnf; apply N.lor_assoc ]. Qed.

  Global Instance fold_right_Proper_Permutation_lxor
  : Proper (eq ==> @Permutation _ ==> eq) (fold_right N.lxor) | 100.
  Proof using Type. apply fold_right_Proper_commutative_associative_Permutation; [ hnf; apply N.lxor_comm | hnf; intros; symmetry; apply N.lxor_assoc ]. Qed.
End N.

Module Z.
  Global Instance fold_right_Proper_Permutation_add
  : Proper (eq ==> @Permutation _ ==> eq) (fold_right Z.add) | 100.
  Proof using Type. apply fold_right_Proper_commutative_associative_Permutation; [ hnf; apply Z.add_comm | hnf; apply Z.add_assoc ]. Qed.

  Global Instance fold_right_Proper_Permutation_mul
  : Proper (eq ==> @Permutation _ ==> eq) (fold_right Z.mul) | 100.
  Proof using Type. apply fold_right_Proper_commutative_associative_Permutation; [ hnf; apply Z.mul_comm | hnf; apply Z.mul_assoc ]. Qed.

  Global Instance fold_right_Proper_Permutation_land
  : Proper (eq ==> @Permutation _ ==> eq) (fold_right Z.land) | 100.
  Proof using Type. apply fold_right_Proper_commutative_associative_Permutation; [ hnf; apply Z.land_comm | hnf; apply Z.land_assoc ]. Qed.

  Global Instance fold_right_Proper_Permutation_lor
  : Proper (eq ==> @Permutation _ ==> eq) (fold_right Z.lor) | 100.
  Proof using Type. apply fold_right_Proper_commutative_associative_Permutation; [ hnf; apply Z.lor_comm | hnf; apply Z.lor_assoc ]. Qed.

  Global Instance fold_right_Proper_Permutation_lxor
  : Proper (eq ==> @Permutation _ ==> eq) (fold_right Z.lxor) | 100.
  Proof using Type. apply fold_right_Proper_commutative_associative_Permutation; [ hnf; apply Z.lxor_comm | hnf; intros; symmetry; apply Z.lxor_assoc ]. Qed.
End Z.


(** Parametric Transition Matrices **)
Require Import ListSet.
Require Import Reals.
Require Import List.
Require Import SetoidList.

Require Import State.
Require Import TransitionMatrix.Definitions.
Require Import TransitionMatrix.Real.
Require Import RatExpr.
Require Import Probabilities.

Open Scope R_scope.

Definition Var : Type := RatExpr.Var.
Definition Var_eq_dec := RatExpr.Var_eq_dec.

(** ** Variables in the Matrix **)

Fixpoint var_set_list' (l: list (State * RatExpr)): set Var :=
match l with
  | nil       => empty_set Var
  | e :: rest => set_union Var_eq_dec (vars (snd e)) (var_set_list' rest)
end.

Definition var_set' (r: StateMaps.t RatExpr): set Var :=
    var_set_list' (StateMaps.elements r).

Fixpoint var_set_list (l: list (State * StateMaps.t RatExpr)): set Var :=
match l with
  | nil       => empty_set Var
  | e :: rest => set_union Var_eq_dec (var_set' (snd e)) (var_set_list rest)
end.

Definition var_set (m: TransitionMatrix RatExpr): set Var :=
    var_set_list (StateMaps.elements m).

Lemma vars_injection: forall (m: TransitionMatrix RatExpr) (r: StateMaps.t RatExpr) (s: State) (x: Var),
    StateMaps.MapsTo s r m -> In x (var_set' r) -> In x (var_set m).
Proof.
    intros m r s x H_mapsto H_x_in_r.
    unfold var_set.
    apply StateMaps.elements_1 in H_mapsto.
    elim H_mapsto.
    - (* r is the first element *)
      intros.
      red in H. red in H. destruct H as [H_fst H_snd].
      unfold var_set_list.
      rewrite <- H_snd. simpl.
      apply set_union_intro1. assumption.
    - (* r is in the tail of elements list *)
      intros.
      unfold var_set_list.
      apply set_union_intro2. assumption.
Qed.

(** ** Evaluation **)

Definition eval_row (r: StateMaps.t RatExpr) (u: Evaluation) : StateMaps.t R :=
    StateMaps.map (fun val => eval_expr val u) r.

Definition eval_matrix (m: TransitionMatrix RatExpr) (u: Evaluation): TransitionMatrix R :=
    StateMaps.map (fun r => eval_row r u) m.

Definition expr_sum_in_map (m: StateMaps.t RatExpr) : RatExpr :=
    sum_f_in_map m expr_sum (Const 0).

Lemma commutative_eval_elements:
  forall (m: StateMaps.t RatExpr) (u: Evaluation),
      let eval_snd := (fun p => (fst p, eval_expr (snd p) u)) in
      StateMaps.elements (eval_row m u) = map eval_snd (StateMaps.elements m).
Proof.
    intros.
    unfold eval_row. unfold StateMaps.elements. unfold StateMaps.Raw.elements.
    unfold StateMaps.map. simpl.
    induction StateMaps.this.
    - reflexivity.
    - simpl. rewrite IHt.
      destruct a as [k v]. reflexivity.
Qed.

Lemma eval_sum_in_map:
    forall (m: StateMaps.t RatExpr) (u: Evaluation),
        eval_expr (expr_sum_in_map m) u = sum_in_map (eval_row m u).
Proof.
    intros.
    unfold expr_sum_in_map.
    unfold sum_in_map. unfold sum_f_in_map.
    unfold sum_f.
    rewrite commutative_eval_elements.
    induction (StateMaps.elements (elt:=RatExpr) m).
    - (* nil *)
      reflexivity.
    - (* a:: l *)
      rewrite map_cons.
      unfold fold_right. simpl.
      rewrite distributive_eval_sum.
      unfold fold_right in IHl.
      rewrite IHl. reflexivity.
Qed.

Lemma commutative_eval_find:
    forall (m: TransitionMatrix RatExpr) (s: State) (u: Evaluation),
        StateMaps.In s m ->
        StateMaps.find s (eval_matrix m u) = option_map (fun row => eval_row row u) (StateMaps.find s m).
Proof.
    intros m s u H_s_in_m.
    unfold eval_matrix.
    inversion H_s_in_m as [e H_mapsto].
    apply StateMaps.find_1 in H_mapsto. rewrite -> H_mapsto.
    apply StateMaps.find_2 in H_mapsto.
    apply StateMaps.map_1 with (elt':=StateMaps.t R) (f:=fun r : StateMaps.t RatExpr => eval_row r u) in H_mapsto.
    apply StateMaps.find_1 in H_mapsto. rewrite -> H_mapsto.
    unfold option_map. trivial.
Qed.

Lemma commutative_eval_mapsto:
    forall (r: StateMaps.t RatExpr) (s: State) (u: Evaluation) (v: R),
        StateMaps.MapsTo s v (eval_row r u) ->
            ( exists e: RatExpr, StateMaps.MapsTo s e r
              /\ v = eval_expr e u ).
Proof.
    intros.
    assert (H':=H). (* backing H up *)
  (* we first prove that s in (eval_row r u) implies s in r *)
    apply mapsto_in in H.
    apply StateMaps.map_2 in H. rewrite in_mapsto in H.
    destruct H as [e H].
  (* now we have a witness that s in r *)
    exists e.
    split. apply H.
  (* to prove the equality, we rely on results about 'map' *)
    unfold eval_row in H'.
    apply StateMaps.map_1 with (f:=(fun val : RatExpr => eval_expr val u)) in H.
    symmetry. apply (StateMapsFacts.MapsTo_fun H H').
Qed.

Lemma eval_matrix_in:
    forall (m: TransitionMatrix RatExpr) (s: State) (u: Evaluation),
        StateMaps.In s m <-> StateMaps.In s (eval_matrix m u).
Proof.
    intros. unfold eval_matrix.
    split.
    - (* -> *)
      intros.
      rewrite in_mapsto in H. elim H. intros.
      rewrite in_mapsto. exists (eval_row x u).
      apply StateMaps.map_1 with (f:=(fun r : StateMaps.t RatExpr => eval_row r u)).
      apply H0.
    - (* <- *)
      intros. apply (StateMaps.map_2 H).
Qed.

Lemma eval_matrix_mapsto:
    forall (m: TransitionMatrix RatExpr) (s: State) (r: StateMaps.t RatExpr) (u: Evaluation),
        StateMaps.MapsTo s r m -> StateMaps.MapsTo s (eval_row r u) (eval_matrix m u).
Proof.
    intros m s r u. unfold eval_matrix. apply StateMaps.map_1.
Qed.

Lemma eval_matrix_evaluated_expr:
    forall (m: TransitionMatrix RatExpr) (s: State) (r: StateMaps.t R) (u: Evaluation),
        StateMaps.MapsTo s r (eval_matrix m u) ->
            ( exists e: StateMaps.t RatExpr, StateMaps.MapsTo s e m
              /\ r = eval_row e u ).
Proof.
    intros.
    assert (H':=H). (* backing H up *)
  (* we first prove that s in (eval_row r u) implies s in r *)
    apply mapsto_in in H.
    apply StateMaps.map_2 in H. rewrite in_mapsto in H.
    destruct H as [e H].
  (* now we have a witness that s in r *)
    exists e.
    split. apply H.
  (* to prove the equality, we rely on results about 'map' *)
    unfold eval_matrix in H'.
    apply StateMaps.map_1 with (f:=(fun r : StateMaps.t RatExpr => eval_row r u)) in H.
    symmetry. apply (StateMapsFacts.MapsTo_fun H H').
Qed.

Lemma sum_accumulator_commutative (T: Type):
    forall (g: T -> RatExpr)
           (x: RatExpr)
           (l: list T),
        let f := fun (v: T) (e: RatExpr) => expr_sum e (g v) in
        fold_right f x l = expr_sum x (fold_right f (Const 0) l).
Proof.
    intros.
    induction l.
    - simpl. unfold expr_sum.
      destruct x;
          try rewrite Rplus_0_r;
          try rewrite expr_0_identity;
          reflexivity.
    - simpl. rewrite IHl.
      unfold f.
      rewrite expr_sum_associative'.
      trivial.
Qed.

Lemma sum_in_map: forall (r: StateMaps.t RatExpr) (s: State) (e: RatExpr),
    StateMaps.MapsTo s e r ->
    expr_sum_in_map r = expr_sum e (expr_sum_in_map (StateMaps.remove s r)).
Proof.
    intros.
  (* expose elements of r *)
    unfold expr_sum_in_map. unfold sum_f_in_map.
    unfold sum_f.
  (* now we manipulate the list of elements to expose (s, e) *)
    assert (H':=H). (* backup *)
    apply StateMaps.elements_1 in H.
    apply SetoidList.InA_split in H.
    destruct H as [l1 [a [l2 [H_eq_a H_elem]]]].
    rewrite H_elem.
    rewrite remove_elements' with (e:=e) (l1:=l1) (l2:=l2) (a:=a); try assumption.
  (* finally, let us leverage sums and folds commutative and
     associative properties *)
    rewrite fold_right_app.
    unfold fold_right at 2.
    rewrite expr_sum_commutative'.
    rewrite sum_accumulator_commutative.
    rewrite <- expr_sum_associative'.
    replace
        ((fix fold_right (l : list (StateMaps.key * RatExpr)) : RatExpr :=
           match l with
           | Datatypes.nil => Const 0
           | b :: t => expr_sum (fold_right t) (snd b)
           end) l2)
      with
        (fold_right
          (fun (v : StateMaps.key * RatExpr) (e0 : RatExpr) =>
           expr_sum e0 (snd v)) (Const 0) l2).
      Focus 2. unfold fold_right. trivial.
    rewrite <- sum_accumulator_commutative.
    rewrite <- fold_right_app.
  (* e = snd a *)
    unfold StateMaps.eq_key_elt, StateMaps.Raw.PX.eqke in H_eq_a.
    simpl in H_eq_a.
    destruct H_eq_a as [_ H_e].
    rewrite H_e.
    trivial.
Qed.

Lemma sum_zero_in_map:
    forall (r: StateMaps.t RatExpr),
        (forall s: State, StateMaps.In s r -> StateMaps.MapsTo s (Const 0) r) ->
            expr_sum_in_map r = Const 0.
Proof.
    intros.
    unfold expr_sum_in_map. unfold sum_f_in_map.
    assert (H':
        (forall s: State, StateMaps.In s r -> StateMaps.MapsTo s (Const 0) r ) ->
          (forall p: State * RatExpr, (In p (StateMaps.elements r)) -> snd p = Const 0)).
    {
        intros.
        specialize (H0 (fst p)).
        apply SetoidList.In_InA with (eqA:=StateMaps.eq_key_elt (elt:=RatExpr)) in H1.
        - destruct p as [s e].
          apply StateMaps.elements_2 in H1. assert (H':= H1).
          simpl in H0.
          apply mapsto_in in H1. apply H0 in H1.
          apply StateMapsFacts.MapsTo_fun with (e':=Const 0) (e:=e) in H1.
          + auto.
          + assumption.
        - exact (StateMapsProperties.eqke_equiv RatExpr).
    }
    induction (StateMaps.elements (elt:=RatExpr) r).
    - (* empty list *)
      reflexivity.
    - (* a::l *)
      eapply H' in H.
      + (* main goal *)
        unfold sum_f. unfold fold_right.
        unfold sum_f in IHl. unfold fold_right in IHl.
        rewrite IHl.
        * rewrite H. simpl. rewrite Rplus_0_r. reflexivity.
        * intros.
          { apply H' with (p:=p) in H0.
            - assumption.
            - apply in_cons. assumption.
          }
      + (* proof obligation *)
        apply in_eq.
Qed.

(** ** Evaluation and stochasticity **)
Definition expr_is_stochastic_row (r: StateMaps.t RatExpr) : Prop :=
    (forall (s: State) (e: RatExpr), StateMaps.MapsTo s e r -> expr_is_valid_prob e)
    /\ expr_sum_in_map r = Const 1.

Lemma eval_stochastic_row:
    forall (r: StateMaps.t RatExpr) (u: Evaluation),
        expr_is_stochastic_row r -> is_stochastic_row (eval_row r u).
Proof.
    intros.
    unfold expr_is_stochastic_row in H.
    destruct H as [H_valid_prob H_sum_1].
    unfold is_stochastic_row. split.
    - (* all values are valid probabilities *)
      intros.
      apply commutative_eval_mapsto in H.
      inversion H as [e [H_s_mapsto_e H_v_eval_e]].
      apply H_valid_prob in H_s_mapsto_e.
      rewrite H_v_eval_e.
      apply valid_prob_preservation. assumption.
    - (* sum is 1 *)
      rewrite <- eval_sum_in_map.
      rewrite H_sum_1. reflexivity.
Qed.


(*
  Author(s):
    Andrej Dudenhefner (1)
  Affiliation(s):
    (1) Saarland University, Saarbrücken, Germany
*)

(*
  Reduction from:
    Left-uniform Two-Inequality Semi-unification (LU2SemiU)
  to:
    System F Typability (SysF_TYP)

  Related Work:
  [1] Wells, Joe B. "Typability and Type-Checking in the Second-Order lambda-Calculus are Equivalent and Undecidable."
      Proceedings Ninth Annual IEEE Symposium on Logic In Computer Science. IEEE, 1994.
*)

Require Import List Lia.
Import ListNotations.
Require Import Undecidability.SystemF.SysF Undecidability.SystemF.Autosubst.syntax.
From Undecidability.SystemF.Util Require Import Facts poly_type_facts pure_term_facts term_facts typing_facts pure_typing_facts pure_typable_prenex.

Require Undecidability.SemiUnification.SemiU.
Module SU := SemiU.

Require Import ssreflect ssrbool ssrfun.

Set Default Goal Selector "!".

Module Argument.
Section LU2SemiU_to_SysF_TYP.

Variables SUs SUt0 SUt1 : SU.term. (* semi-unification instance *)

(* construct polymorphic type from semi-unification term *)
Fixpoint SU_term_to_poly_type (s: SU.term) :=
  match s with
  | SU.atom n => poly_var n
  | SU.arr s t => poly_arr (SU_term_to_poly_type s) (SU_term_to_poly_type t)
  end.

Definition SUs' := SU_term_to_poly_type SUs.
Definition SUt0' := SU_term_to_poly_type SUt0.
Definition SUt1' := SU_term_to_poly_type SUt1.

Definition N_W := pure_abs (many_pure_app (pure_var 2) [pure_var 0; pure_var 0]).
(* x (\z.y z z) as given by Wells *)
Definition M_W : pure_term := many_pure_app (pure_var 0) [N_W].

Definition t_0_W := (poly_arr (poly_arr (poly_var 0) (poly_var 0)) (poly_var 1)).
Definition t_1_W := poly_arr (poly_arr SUt0' SUt0') (poly_arr (poly_arr SUt1' SUt1') (poly_arr SUs' SUs')).
Definition n_W : nat := locked (poly_var_bound t_1_W).
(* similar to Wells' construction
 ∀((α -> α) -> β) ;
 ∀ .. ∀ (t1 -> t1) -> (t2 -> t2) -> (s -> s) *)
Definition Gamma_W : list poly_type := [poly_abs (poly_abs t_0_W); many_poly_abs n_W t_1_W].

Lemma is_simple_SU_term_to_poly_type {s} : is_simple (SU_term_to_poly_type s) <-> True.
Proof. elim: s; [done | by move=> /= ? + ? + => -> ->]. Qed.

Lemma is_simple_t_1_W : is_simple t_1_W.
Proof.
  by rewrite /t_1_W /= /SUt0' /SUt1' /SUs' ?is_simple_SU_term_to_poly_type.
Qed.

Lemma closed_t_1_W : allfv_poly_type (gt n_W) t_1_W.
Proof. rewrite /n_W -lock. by apply: poly_var_boundP. Qed.

Lemma is_simple_t_0_W : is_simple t_0_W.
Proof. done. Qed.

Lemma closed_t_0_W : allfv_poly_type (gt 2) t_0_W.
Proof. move=> /=. by lia. Qed.

(* introduce assumption poly_abs (poly_abs t_0_W) *)
Definition HM_W1 := pure_typable_intro_prenex M_W t_0_W 2 is_simple_t_0_W closed_t_0_W.

(* introduce assumption many_poly_abs n_W t_1_W *)
Definition HM_W2 := pure_typable_intro_prenex (sval HM_W1) t_1_W n_W is_simple_t_1_W closed_t_1_W.

Lemma SysF_TYPI Gamma {M} : pure_typable Gamma M -> SysF_TYP M.
Proof. move=> [t] /pure_typing_iff_type_assignment ?. by exists Gamma, t. Qed.

Lemma prenex_Gamma_W :
  Forall (fun t => exists n s, t = many_poly_abs n s /\ is_simple s /\ allfv_poly_type (gt n) s) Gamma_W.
Proof.
  apply /Forall_cons_iff. constructor.
  { exists 2, t_0_W. constructor; [done | constructor; [done | by move=> /=; lia ]]. }
  apply /Forall_cons_iff. constructor; last done.
  do 2 eexists. constructor; first by reflexivity.
  by constructor; [apply: is_simple_t_1_W | apply: closed_t_1_W ].
Qed.

(* if Gamma_W ⊢ M_W : poly_var 0 typechecks, then the semi-unification instance (SUs, SUt0, SUt1) is solvable *)
Section InverseTransport.
Local Arguments many_pure_app : simpl never.
Local Arguments map : simpl never.

Lemma map_ren_poly_type_Gamma_W {ξ}: map (ren_poly_type ξ) Gamma_W = Gamma_W.
Proof.
  congr cons. congr cons. apply: ren_poly_type_closed_id.
  rewrite allfv_poly_type_many_poly_abs.
  apply: allfv_poly_type_impl closed_t_1_W.
  move=> ? ?. rewrite iter_scons_lt; by [|lia].
Qed.

(* decompose first layer of Wells' type checking term *)
Lemma decompose_M_W : pure_typable Gamma_W M_W ->
  exists s, pure_typing Gamma_W N_W (poly_arr s s).
Proof.
  move=> [?] /pure_typingE /= [?] [?] [?] [?] [+] [+] [_ ?].
  move=> /pure_typingE' [?] [[?]]. subst.
  move=> /(contains_poly_arrE (n := 2)) [?] [?] [? ?]. subst.
  rewrite map_ren_poly_type_Gamma_W => ?. eexists. by eassumption.
Qed.

(* decompose second layer of Wells' type checking term *)
Lemma decompose_N_W {s} : pure_typing Gamma_W N_W (poly_arr s s) ->
  exists n σ,
    s = many_poly_abs n (subst_poly_type σ (poly_arr SUs' SUs')) /\
    pure_typing (map (ren_poly_type (Nat.add n)) (s :: Gamma_W)) (pure_var 0) (subst_poly_type σ (poly_arr SUt0' SUt0')) /\
    pure_typing (map (ren_poly_type (Nat.add n)) (s :: Gamma_W)) (pure_var 0) (subst_poly_type σ (poly_arr SUt1' SUt1')).
Proof.
  move=> /pure_typingE' /=.
  move=> /pure_typingE /= [n1] [?] [?] [?] [+] [+] [H1C ?].
  move=> /pure_typingE' [?] [?] [+] [+ H2C].
  move=> /pure_typingE' [?] [[?]] H3C. subst.
  move=> H1 H2.
  move: H3C. rewrite ren_poly_type_many_poly_abs /t_1_W /=.
  move=> /contains_poly_arrE [?] [?] [? ?]. subst.
  move: H2C => /containsE [? ?]. subst.
  move: H1C => /containsE ?. subst.
  do 2 eexists.
  constructor; first by rewrite ?poly_type_norm.
  constructor; [move: H1 | move: H2]; by rewrite ?poly_type_norm.
Qed.

(* decompose third layer of Wells' type checking term *)
Lemma decompose_x_W {n σ Gamma s' t'} :
  let s := many_poly_abs n (subst_poly_type σ (poly_arr s' s')) in
  pure_typing ((ren_poly_type (Nat.add n) s) :: Gamma) (pure_var 0) (subst_poly_type σ (poly_arr t' t')) ->
  exists τ, subst_poly_type (funcomp (subst_poly_type τ) σ) s' = subst_poly_type σ t'.
Proof.
  move=> /= /pure_typingE' /= [?] [[?]]. subst.
  rewrite ren_poly_type_many_poly_abs /=.
  move=> /contains_poly_arrE [?] [?] [-> _].
  rewrite ?poly_type_norm. by eexists.
Qed.

(* construct semi-unification term from a polymorphic type pruning quantification *)
Fixpoint prune (s: poly_type) : SU.term :=
  match s with
  | poly_var x => SU.atom (1 + x)
  | poly_arr s t => SU.arr (prune s) (prune t)
  | poly_abs _ => SU.atom 0
  end.

Fact substitute_prune_ex τ t : exists ψ, SU.substitute ψ (prune t) = prune (subst_poly_type τ t).
Proof.
  exists (fun n=> if n is S n then prune (τ n) else SU.atom 0).
  elim: t; [done | by move=> ? + ? /= => -> -> | done ].
Qed.

Fact substitute_prune {σ t} :
  SU.substitute (funcomp prune σ) t = prune (subst_poly_type σ (SU_term_to_poly_type t)).
Proof. elim: t; [done | by move=> ? + ? => /= -> -> ]. Qed.

(* construct semi-unification valuations from polymorphic substitutions *)
Lemma introduce_simple_substitutions {s t σ τ} :
  subst_poly_type (funcomp (subst_poly_type τ) σ) (SU_term_to_poly_type s) = subst_poly_type σ (SU_term_to_poly_type t) ->
  exists ψ, SU.substitute ψ (SU.substitute (funcomp prune σ) s) = SU.substitute (funcomp prune σ) t.
Proof.
  move=> H. have [ψ Hψ] := substitute_prune_ex τ (subst_poly_type σ (SU_term_to_poly_type s)).
  exists ψ. by rewrite ?substitute_prune Hψ -H poly_type_norm.
Qed.

End InverseTransport.

(* typability to semi-unification solution *)
Lemma inverse_transport : SysF_TYP (sval HM_W2) -> SU.LU2SemiU (SUs, SUt0, SUt1).
Proof.
  move=> [Gamma] [t] /pure_typing_iff_type_assignment.
  move=> /pure_typing_tidyI /pure_typableI /(svalP HM_W2) /(svalP HM_W1).
  move=> [?] /(pure_typing_ren_pure_term_allfv_pure_term id (Delta := map tidy Gamma_W)).
  apply: unnest; first done. rewrite ren_pure_term_id => /pure_typableI.
  move=> /(pure_typable_tidy_iff (Gamma1 := [])). apply: unnest; first by apply: prenex_Gamma_W.
  move=> /decompose_M_W [?] /decompose_N_W [?] [σ] [->].
  move=> [/decompose_x_W [?] + /decompose_x_W [?] +].
  move=> /introduce_simple_substitutions [ψ0 ?] /introduce_simple_substitutions [ψ1 ?].
  by exists (funcomp prune σ), ψ0, ψ1.
Qed.

(* if the semi-unification instance (SUs, SUt0, SUt1) is solvable, then Gamma_W ⊢ M_W : poly_var 0 typechecks *)
Section Transport.

(* semi-unification solution *)
Variables φ ψ0 ψ1 : nat -> SU.term.
Variable Hφψ0 : SU.substitute ψ0 (SU.substitute φ SUs) = SU.substitute φ SUt0.
Variable Hφψ1 : SU.substitute ψ1 (SU.substitute φ SUs) = SU.substitute φ SUt1.

Definition t_x_W' := SU_term_to_poly_type (SU.substitute φ SUs).
Definition n_W' := locked (poly_var_bound (poly_arr t_x_W' t_x_W')).
Definition t_x_W : poly_type := many_poly_abs n_W' (poly_arr t_x_W' t_x_W').

Definition t_x_0_W : poly_type := SU_term_to_poly_type (SU.substitute φ SUt0).
Definition t_x_1_W : poly_type := SU_term_to_poly_type (SU.substitute φ SUt1).

Definition ts_0_W : list poly_type := map (fun n => SU_term_to_poly_type (φ n)) (rev (seq 0 n_W)).
Definition ts_1_W : list poly_type := map (fun n => SU_term_to_poly_type (ψ0 n)) (rev (seq 0 n_W')).
Definition ts_2_W : list poly_type := map (fun n => SU_term_to_poly_type (ψ1 n)) (rev (seq 0 n_W')).

Definition x_1_W : term := many_ty_app (var 0) ts_1_W.
Definition x_2_W : term := many_ty_app (var 0) ts_2_W.
Definition Q_W : term := abs t_x_W (many_ty_abs n_W' (app (app (many_ty_app (var 2) ts_0_W) x_1_W) x_2_W)).
(* correctly annotated M_W, similar to Wells' construction *)
Definition P_W : term := app (ty_app (ty_app (var 0) (poly_var 0)) t_x_W) Q_W.

Fact poly_var_bound_SUt0'_n_W : poly_var_bound SUt0' < n_W.
Proof. rewrite /n_W /t_1_W /= -lock. by lia. Qed.

Fact poly_var_bound_SUt1'_n_W : poly_var_bound SUt1' < n_W.
Proof. rewrite /n_W /t_1_W /= -lock. by lia. Qed.

Fact poly_var_bound_SUs'_n_W : poly_var_bound SUs' < n_W.
Proof. rewrite /n_W /t_1_W /= -lock. by lia. Qed.

Lemma SU_term_to_poly_type_SU_substitute {ζ t} :
  subst_poly_type (fun x : nat => SU_term_to_poly_type (ζ x)) (SU_term_to_poly_type t) =
    SU_term_to_poly_type (SU.substitute ζ t).
Proof. elim: t; [done | by move=> /= ? -> ? ->]. Qed.

Lemma subst_poly_type_SU_term_to_poly_type {ζ n t} :
  poly_var_bound (SU_term_to_poly_type t) < n ->
  subst_poly_type
    (fold_right scons poly_var (map (fun x => SU_term_to_poly_type (ζ x)) (seq 0 n)))
    (SU_term_to_poly_type t) =
  SU_term_to_poly_type (SU.substitute ζ t).
Proof.
  move=> ?. rewrite subst_poly_type_map_seq; first by lia.
  by apply: SU_term_to_poly_type_SU_substitute.
Qed.

Fact poly_arr_eq_eq {s t} : s = t -> poly_arr s s = poly_arr t t.
Proof. by move=> ->. Qed.

Lemma typing_x_1_W :
  (forall Gamma, typing (many_poly_abs n_W' (poly_arr t_x_W' t_x_W') :: Gamma) x_1_W (poly_arr t_x_0_W t_x_0_W)).
Proof using Hφψ0.
  move=> Gamma. rewrite /x_1_W. evar (t: poly_type).
  have := @typing_var (many_poly_abs n_W' (poly_arr t_x_W' t_x_W') :: Gamma) 0 t.
  subst t. move=> /(_ ltac:(reflexivity)) /typing_many_ty_appI => /(_ ts_1_W).
  apply: unnest. { by rewrite /ts_1_W map_length rev_length seq_length. }
  congr typing => /=. apply: poly_arr_eq_eq. rewrite /ts_1_W /t_x_W' /t_x_0_W -Hφψ0.
  rewrite -map_rev rev_involutive. apply: subst_poly_type_SU_term_to_poly_type.
  rewrite /n_W' /t_x_W' -lock /=. by lia.
Qed.

Lemma typing_x_2_W :
  (forall Gamma, typing (many_poly_abs n_W' (poly_arr t_x_W' t_x_W') :: Gamma) x_2_W (poly_arr t_x_1_W t_x_1_W)).
Proof using Hφψ1.
  move=> Gamma. rewrite /x_2_W. evar (t: poly_type).
  have := @typing_var (many_poly_abs n_W' (poly_arr t_x_W' t_x_W') :: Gamma) 0 t.
  subst t. move=> /(_ ltac:(reflexivity)) => /typing_many_ty_appI => /(_ ts_2_W).
  apply: unnest. { by rewrite /ts_2_W map_length rev_length seq_length. }
  congr typing => /=. apply: poly_arr_eq_eq. rewrite /ts_2_W /t_x_W' /t_x_1_W -Hφψ1.
  rewrite -map_rev rev_involutive. apply: subst_poly_type_SU_term_to_poly_type.
  rewrite /n_W' /t_x_W' -lock /=. by lia.
Qed.

Lemma typing_Q_W :
  (forall Gamma, typing (many_poly_abs n_W' (poly_arr t_x_W' t_x_W') :: Gamma) x_1_W (poly_arr t_x_0_W t_x_0_W)) ->
  (forall Gamma, typing (many_poly_abs n_W' (poly_arr t_x_W' t_x_W') :: Gamma) x_2_W (poly_arr t_x_1_W t_x_1_W)) ->
  typing Gamma_W Q_W (poly_arr t_x_W t_x_W).
Proof.
  move=> H1 H2.
  apply /typing_abs /typing_many_ty_absI.
  apply: (typing_app (s := (poly_arr t_x_1_W t_x_1_W))); [apply: (typing_app (s := (poly_arr t_x_0_W t_x_0_W))) |].
  - evar (t : poly_type).
    have := @typing_var (map (ren_poly_type (Nat.add n_W')) (t_x_W :: Gamma_W)) 2 t.
    subst t. move=> /(_ ltac:(reflexivity)).
    rewrite ren_poly_type_many_poly_abs.
    move /typing_many_ty_appI => /(_ ts_0_W). apply: unnest.
    { by rewrite /ts_0_W map_length rev_length seq_length. }
    congr typing. rewrite /t_1_W /= -map_rev rev_involutive.
    congr poly_arr; [| congr poly_arr]; apply: poly_arr_eq_eq.
    + have := poly_var_bound_SUt0'_n_W. rewrite /t_x_0_W /SUt0' /ts_0_W.
      move=> ?. rewrite ren_poly_type_poly_var_bound; first by lia.
      by apply: subst_poly_type_SU_term_to_poly_type.
    + have := poly_var_bound_SUt1'_n_W. rewrite /t_x_1_W /SUt1' /ts_0_W.
      move=> ?. rewrite ren_poly_type_poly_var_bound; first by lia.
      by apply: subst_poly_type_SU_term_to_poly_type.
    + have := poly_var_bound_SUs'_n_W. rewrite /t_x_W' /SUs' /ts_0_W.
      move=> ?. rewrite ren_poly_type_poly_var_bound; first by lia.
      by apply: subst_poly_type_SU_term_to_poly_type.
  - rewrite /(map _ (_ :: _)) -/(map _ _). move: (map _ Gamma_W) => Gamma.
    rewrite /t_x_W ren_poly_type_many_poly_abs ren_poly_type_poly_var_bound;
      first by (rewrite /n_W' -lock /=; lia).
    by apply: H1.
  - rewrite /(map _ (_ :: _)) -/(map _ _). move: (map _ Gamma_W) => Gamma.
    rewrite /t_x_W ren_poly_type_many_poly_abs ren_poly_type_poly_var_bound;
      first by (rewrite /n_W' -lock /=; lia).
    by apply: H2.
Qed.

Lemma erase_P_W : erase P_W = M_W.
Proof.
  by rewrite /P_W /= erase_many_ty_abs /= /x_1_W /x_2_W ?erase_many_ty_app.
Qed.

Lemma typing_P_W :
  typing Gamma_W Q_W (poly_arr t_x_W t_x_W) ->
  typing Gamma_W P_W (poly_var 0).
Proof.
  rewrite /P_W. apply: typing_app.
  apply: typing_ty_appI; [apply: typing_ty_appI |]; [by apply: typing_var | done | done].
Qed.

Lemma pure_typable_M_W :
  typing Gamma_W P_W (poly_var 0) ->
  pure_typable Gamma_W M_W.
Proof.
  move=> /typing_to_pure_typing /pure_typableI. by rewrite erase_P_W.
Qed.

End Transport.

(* semi-unification solution to typability *)
Lemma transport : SU.LU2SemiU (SUs, SUt0, SUt1) -> SysF_TYP (sval HM_W2).
Proof.
  move=> [φ] [ψ0] [ψ1] [Hφψ0 HHφψ1]. apply: (SysF_TYPI (map tidy [])).
  apply /(svalP HM_W2) /(svalP HM_W1).
  have := iffRL (pure_typable_tidy_iff _ (Gamma1 := [])).
  move=> /(_ [many_poly_abs 2 t_0_W; many_poly_abs n_W t_1_W]).
  apply; first by apply: prenex_Gamma_W.
  apply: (pure_typable_M_W φ ψ0 ψ1).
  apply: typing_P_W.
  apply: typing_Q_W; [apply: typing_x_1_W | apply: typing_x_2_W]; by eassumption.
Qed.

End LU2SemiU_to_SysF_TYP.
End Argument.

Require Import Undecidability.Synthetic.Definitions.
Import SemiU.

(* Left-uniform Two-Inequality Semi-unification many-one reduces to System F Typability *)
Theorem reduction : LU2SemiU ⪯ SysF_TYP.
Proof.
  exists (fun '(s, t0, t1) => (sval (Argument.HM_W2 s t0 t1))).
  move=> [[s t0] t1]. constructor.
  - exact: Argument.transport.
  - exact: Argument.inverse_transport.
Qed.


(* This file is generated by Why3's Coq-realize driver *)
(* Beware! Only edit allowed sections below    *)
Require Import BuiltIn.
Require Reals.R_sqrt.
Require Reals.Rbasic_fun.
Require Reals.Rtrigo_def.
Require Reals.Rtrigo1.
Require Reals.Ratan.
Require BuiltIn.
Require real.Real.
Require real.Abs.
Require real.Square.

Require Import Reals.

(* Why3 comment *)
(* cos is replaced with (Reals.Rtrigo_def.cos x) by the coq driver *)

(* Why3 comment *)
(* sin is replaced with (Reals.Rtrigo_def.sin x) by the coq driver *)

(* Why3 goal *)
Lemma Pythagorean_identity : forall (x:R),
  (((Reals.RIneq.Rsqr (Reals.Rtrigo_def.cos x)) + (Reals.RIneq.Rsqr (Reals.Rtrigo_def.sin x)))%R = 1%R).
Proof.
intros x.
rewrite Rplus_comm.
apply sin2_cos2.
Qed.

(* Why3 goal *)
Lemma Cos_le_one : forall (x:R),
  ((Reals.Rbasic_fun.Rabs (Reals.Rtrigo_def.cos x)) <= 1%R)%R.
Proof.
intros x.
apply Abs.Abs_le.
apply COS_bound.
Qed.

(* Why3 goal *)
Lemma Sin_le_one : forall (x:R),
  ((Reals.Rbasic_fun.Rabs (Reals.Rtrigo_def.sin x)) <= 1%R)%R.
Proof.
intros x.
apply Abs.Abs_le.
apply SIN_bound.
Qed.

(* Why3 goal *)
Lemma Cos_0 : ((Reals.Rtrigo_def.cos 0%R) = 1%R).
Proof.
apply cos_0.
Qed.

(* Why3 goal *)
Lemma Sin_0 : ((Reals.Rtrigo_def.sin 0%R) = 0%R).
Proof.
apply sin_0.
Qed.

(* Why3 comment *)
(* pi is replaced with Reals.Rtrigo1.PI by the coq driver *)

(* Why3 goal *)
Lemma Pi_interval : ((314159265358979323846264338327950288419716939937510582097494459230781640628620899862803482534211706798214808651328230664709384460955058223172535940812848111745028410270193852110555964462294895493038196 / 100000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000)%R < Reals.Rtrigo1.PI)%R /\
  (Reals.Rtrigo1.PI < (314159265358979323846264338327950288419716939937510582097494459230781640628620899862803482534211706798214808651328230664709384460955058223172535940812848111745028410270193852110555964462294895493038197 / 100000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000)%R)%R.
Proof.
replace PI with (4 * (PI / 4))%R by field.
rewrite <- atan_1.
admit. (* to avoid a dependency on CoqInterval *)
(* split ; interval with (i_prec 680). *)
Qed.

(* Why3 goal *)
Lemma Cos_pi : ((Reals.Rtrigo_def.cos Reals.Rtrigo1.PI) = (-1%R)%R).
Proof.
apply cos_PI.
Qed.

(* Why3 goal *)
Lemma Sin_pi : ((Reals.Rtrigo_def.sin Reals.Rtrigo1.PI) = 0%R).
Proof.
apply sin_PI.
Qed.

(* Why3 goal *)
Lemma Cos_pi2 : ((Reals.Rtrigo_def.cos ((05 / 10)%R * Reals.Rtrigo1.PI)%R) = 0%R).
Proof.
replace (5 / 10 * PI)%R with (PI / 2)%R by field.
apply cos_PI2.
Qed.

(* Why3 goal *)
Lemma Sin_pi2 : ((Reals.Rtrigo_def.sin ((05 / 10)%R * Reals.Rtrigo1.PI)%R) = 1%R).
Proof.
replace (5 / 10 * PI)%R with (PI / 2)%R by field.
apply sin_PI2.
Qed.

(* Why3 goal *)
Lemma Cos_plus_pi : forall (x:R),
  ((Reals.Rtrigo_def.cos (x + Reals.Rtrigo1.PI)%R) = (-(Reals.Rtrigo_def.cos x))%R).
Proof.
intros x.
apply neg_cos.
Qed.

(* Why3 goal *)
Lemma Sin_plus_pi : forall (x:R),
  ((Reals.Rtrigo_def.sin (x + Reals.Rtrigo1.PI)%R) = (-(Reals.Rtrigo_def.sin x))%R).
Proof.
intros x.
apply neg_sin.
Qed.

(* Why3 goal *)
Lemma Cos_plus_pi2 : forall (x:R),
  ((Reals.Rtrigo_def.cos (x + ((05 / 10)%R * Reals.Rtrigo1.PI)%R)%R) = (-(Reals.Rtrigo_def.sin x))%R).
Proof.
intros x.
rewrite cos_sin.
replace (PI / 2 + (x + 5 / 10 * PI))%R with (x + PI)%R by field.
apply neg_sin.
Qed.

(* Why3 goal *)
Lemma Sin_plus_pi2 : forall (x:R),
  ((Reals.Rtrigo_def.sin (x + ((05 / 10)%R * Reals.Rtrigo1.PI)%R)%R) = (Reals.Rtrigo_def.cos x)).
Proof.
intros x.
rewrite cos_sin.
apply f_equal.
field.
Qed.

(* Why3 goal *)
Lemma Cos_neg : forall (x:R),
  ((Reals.Rtrigo_def.cos (-x)%R) = (Reals.Rtrigo_def.cos x)).
Proof.
intros x.
apply cos_neg.
Qed.

(* Why3 goal *)
Lemma Sin_neg : forall (x:R),
  ((Reals.Rtrigo_def.sin (-x)%R) = (-(Reals.Rtrigo_def.sin x))%R).
Proof.
intros x.
apply sin_neg.
Qed.

(* Why3 goal *)
Lemma Cos_sum : forall (x:R) (y:R),
  ((Reals.Rtrigo_def.cos (x + y)%R) = (((Reals.Rtrigo_def.cos x) * (Reals.Rtrigo_def.cos y))%R - ((Reals.Rtrigo_def.sin x) * (Reals.Rtrigo_def.sin y))%R)%R).
Proof.
intros x y.
apply cos_plus.
Qed.

(* Why3 goal *)
Lemma Sin_sum : forall (x:R) (y:R),
  ((Reals.Rtrigo_def.sin (x + y)%R) = (((Reals.Rtrigo_def.sin x) * (Reals.Rtrigo_def.cos y))%R + ((Reals.Rtrigo_def.cos x) * (Reals.Rtrigo_def.sin y))%R)%R).
Proof.
intros x y.
apply sin_plus.
Qed.

(* Why3 goal *)
Lemma tan_def : forall (x:R),
  ((Reals.Rtrigo1.tan x) = ((Reals.Rtrigo_def.sin x) / (Reals.Rtrigo_def.cos x))%R).
Proof.
intros x.
apply eq_refl.
Qed.

(* Why3 comment *)
(* atan is replaced with (Reals.Ratan.atan x) by the coq driver *)

(* Why3 goal *)
Lemma Tan_atan : forall (x:R),
  ((Reals.Rtrigo1.tan (Reals.Ratan.atan x)) = x).
Proof.
intros x.
apply atan_right_inv.
Qed.


Require Import ssreflect.
From coq_sortscan Require Import helper.
From stdpp Require Import sorting list.

Context {A} (R : relation A) `{∀ x y, Decision (R x y)}
  `{!Total R, !Transitive R, !AntiSymm (=) R, !Reflexive R}.
Implicit Type (mem : list A).
Implicit Type (L : list (list A)).

Section sorted.
  Inductive is_sorted_lists : list (list A) → Prop :=
  | isl_nil : is_sorted_lists ([] : list (list A))
  | isl_S (l : list A) (L' : list (list A)) :
      StronglySorted R l → is_sorted_lists L' →
      is_sorted_lists (l :: L').

  Fixpoint join_lists L :=
    match L with
    | [] => []
    | l::L' => l ++ join_lists L'
    end.

  Definition is_min_of_join L v :=
    ∃ l, merge_sort R (join_lists L) = v :: l.

  Lemma join_lists_pop_Permutation L i v l :
    L !! i = Some (v::l) →
    v :: join_lists (<[i:=l]> L) ≡ₚ join_lists L.
  Proof.
    intros Hi.
    assert (i < length L) as Hlt. by eapply lookup_lt_Some.
    destruct L. { inversion Hlt. }
    revert i l0 L v l Hi Hlt. induction i; intros.
    - simpl in *. by inversion Hi.
    - simpl in *. destruct L. { inversion Hlt. lia. }
      simpl in *. apply IHi in Hi; [|lia].
      etrans. 1: apply Permutation_middle.
      by apply Permutation_app.
  Qed.

  Lemma merge_sort_Permutation_eq l1 l2 :
    merge_sort R l1 = l2 → l1 ≡ₚ l2.
  Proof.
    intros. rewrite -H0. symmetry. apply merge_sort_Permutation.
  Qed.

  Lemma merge_sort_equivalent `{!Total R, !Transitive R} l1 l2 :
    l1 ≡ₚ l2 → merge_sort R l1 = merge_sort R l2.
  Proof.
    intros HP. apply (StronglySorted_unique R);
      try by apply StronglySorted_merge_sort.
    trans l1. 1: apply merge_sort_Permutation.
    trans l2; auto. symmetry. apply merge_sort_Permutation.
  Qed.

  Lemma merge_sort_identity `{!Total R, !Transitive R} l :
    StronglySorted R l → merge_sort R l = l.
  Proof.
    intros HP. apply (StronglySorted_unique R); auto.
    - apply StronglySorted_merge_sort; auto.
    - apply merge_sort_Permutation.
  Qed.

  Lemma merge_sort_criterion l1 l2 :
    l1 ≡ₚ l2 → StronglySorted R l2 →
    merge_sort R l1 = l2.
  Proof.
    intros HP SS.
    replace l2 with (merge_sort R l2).
    - by apply merge_sort_equivalent.
    - by apply merge_sort_identity.
  Qed.

  Lemma is_min_of_join_cons `{!Total R, !Transitive R, !Reflexive R}
  v' l L v :
    is_sorted_lists ((v'::l) :: L) →
    is_min_of_join L v → R v' v →
    is_min_of_join ((v'::l) :: L) v'.
  Proof.
    intros SS [rest Hmin] HR.
    exists (merge_sort R (l ++ join_lists L)). simpl.
    apply merge_sort_criterion.
    - constructor. symmetry. apply merge_sort_Permutation.
    - constructor. { by apply StronglySorted_merge_sort. }
      eapply (Forall_Permutation (R v') _ _ (l ++ join_lists L)).
      { symmetry. apply merge_sort_Permutation. }
      inversion SS; subst. inversion H2; subst.
      assert (StronglySorted R (v :: rest)) as SSvr.
      { rewrite -Hmin. by apply StronglySorted_merge_sort. }
      apply Forall_app; split; auto.
      apply (Forall_impl (R v)).
      + inversion SSvr; subst.
        eapply (Forall_Permutation (R v) _ _ (v :: rest)).
        { rewrite -Hmin. apply merge_sort_Permutation. }
        by constructor.
      + intros. by trans v.
    Unshelve. all: unfold pointwise_relation; auto.
  Qed.

  Lemma is_min_of_join_cons2 `{!Total R, !Transitive R, !Reflexive R}
  v' l L v :
    is_sorted_lists ((v'::l) :: L) →
    is_min_of_join L v → R v v' →
    is_min_of_join ((v'::l) :: L) v.
  Proof.
    intros SS [rest Hmin] HR.
    exists (merge_sort R (v' :: l ++ rest)). simpl.
    apply merge_sort_criterion.
    - trans (v :: v' :: l ++ rest).
      2: { constructor. symmetry. apply merge_sort_Permutation. }
      apply merge_sort_Permutation_eq in Hmin.
      trans (v' :: l ++ v :: rest).
      1: { constructor. apply Permutation_app; auto. }
      rewrite Permutation_swap. constructor.
      by rewrite Permutation_middle.
    - constructor. { by apply StronglySorted_merge_sort. }
      eapply (Forall_Permutation (R v) _ _ (v' :: l ++ rest)).
      { symmetry. apply merge_sort_Permutation. }
      constructor; auto.
      inversion SS; subst. inversion H2; subst.
      assert (StronglySorted R (v :: rest)) as SSvr.
      { rewrite -Hmin. by apply StronglySorted_merge_sort. }
      apply Forall_app; split.
      + apply (Forall_impl (R v')); auto. intros; by trans v'.
      + inversion SSvr; subst. auto.
    Unshelve. all: unfold pointwise_relation; auto.
  Qed.
End sorted.

Section Phase1.
  Fixpoint phase1_loop mem M (z : nat) :=
    match z with
    | O => ([] : list (list A))
    | S z' => (merge_sort R (take M mem)) ::
        phase1_loop (drop M mem) M z'
    end.

  Definition phase1 mem M :=
    phase1_loop mem M (S (length mem)).

  Lemma phase1_loop_correct `{!Transitive R, !Total R} mem M z :
    0 < M → length mem ≤ M * z →
    let L := phase1_loop mem M z in
    is_sorted_lists L ∧ join_lists L ≡ₚ mem.
  Proof.
    revert mem. simpl.
    induction z; intros mem HM Hz; simpl; split.
    - constructor.
    - destruct mem; auto. simpl in Hz. lia.
    - constructor.
      + by apply StronglySorted_merge_sort.
      + apply IHz; auto. rewrite drop_length. lia.
    - assert (mem = take M mem ++ drop M mem) as HTD
        by by rewrite take_drop.
      rewrite {3} HTD. apply Permutation_app.
      + apply merge_sort_Permutation.
      + apply IHz; auto. rewrite drop_length. lia.
  Qed.

  Lemma phase1_correct mem M :
    0 < M →
    let L := phase1 mem M in
    is_sorted_lists L ∧ join_lists L ≡ₚ mem.
  Proof.
    simpl. unfold phase1. intros HM.
    apply phase1_loop_correct; auto.
    destruct M; lia.
  Qed.
End Phase1.

Section Phase2.
  Fixpoint num_elements L :=
    match L with
    | [] => 0
    | l::L' => length l + num_elements L'
    end.

  Definition look_head L i :=
    match L !! i with
    | None | Some [] => None
    | Some (x::l) => Some x
    end.

  Fixpoint find_min L :=
    match L with
    | [] => 0
    | l::L' =>
        let j := S (find_min L') in
        match look_head L j with
        | None => 0
        | Some xj =>
            match l with
            | [] => j
            | x0::l' => if (bool_decide (R x0 xj)) then 0 else j
            end
        end
    end.

  Fixpoint phase2_loop {V} L (f : A → V → V) (v : V) (z : nat) :=
    match z with
    | O => v
    | S z' =>
        let i := find_min L in
        match L !! i with
        | None | Some [] => v
        | Some (x::l) => f x (phase2_loop (<[i:=l]> L) f v z')
        end
    end.

  Definition phase2 {V} L (f : A → V → V) (start : V) :=
    phase2_loop L f start (num_elements L).

  (* proof *)

  Lemma join_lists_empty L :
    num_elements L = 0 → join_lists L = [].
  Proof.
    induction L; auto; intros.
    simpl in *. assert (length a = 0) by lia.
    destruct a; auto. simpl in H1. lia.
  Qed.

  Lemma look_head_empty L i :
    num_elements L = 0 → look_head L i = None.
  Proof.
    revert i. induction L; auto; intros.
    simpl in *. assert (length a = 0) by lia.
    destruct a.
    - rewrite H1 /= in H0. destruct i; auto.
      by apply (IHL i) in H0.
    - simpl in H1. lia.
  Qed.

  Lemma look_head_cons l L i :
    look_head (l :: L) (S i) = look_head L i.
  Proof. auto. Qed.

  Lemma look_head_Some L i v l :
    L !! i = Some (v :: l) → look_head L i = Some v.
  Proof. intros H. unfold look_head. by rewrite H. Qed.

  Lemma phase2_first_correct `{!Transitive R, !Total R} L:
    num_elements L > 0 → is_sorted_lists L →
    let i := find_min L in
    ∃ v l,
      L !! i = Some (v :: l) ∧
      merge_sort R (join_lists L) =
        v :: merge_sort R (join_lists (<[i:=l]> L)).
  Proof.
    intros Hn Hsort. simpl.
    induction L as [|l L']. { inversion Hn. }
    destruct l.
    - apply IHL' in Hn as IHL; clear IHL'; last by inversion Hsort.
      destruct IHL as [v [l [Hmin Hrest]]].
      simpl. rewrite look_head_cons. erewrite look_head_Some; eauto.
    - rename a into v0.
      destruct (le_gt_dec (num_elements L') 0) as [g|g].
      + simpl. rewrite look_head_cons look_head_empty; try lia.
        exists v0, l. split; auto. simpl.
        rewrite join_lists_empty; [lia|]. rewrite app_nil_r.
        assert (StronglySorted R (v0 :: l)) as SS by by inversion Hsort.
        repeat rewrite merge_sort_identity; auto.
        by inversion SS.
      + apply IHL' in g; clear IHL'; last by inversion Hsort.
        destruct g as [vj [l' [Hmin' Hrest']]].
        simpl. rewrite look_head_cons. erewrite look_head_Some; eauto.
        assert (is_min_of_join L' vj) as Hminjoin'. { eexists. eauto. }
        destruct (bool_decide (R v0 vj)) eqn: ER.
        * apply bool_decide_eq_true_1 in ER.
          exists v0, l. split; auto. simpl.
          assert (is_min_of_join ((v0::l) :: L') v0) as Hminjoin.
          { eapply is_min_of_join_cons; eauto. }
          destruct Hminjoin as [rest' Hrest'2].
          simpl in Hrest'2. rewrite Hrest'2.
          assert (rest' = merge_sort R (l ++ join_lists L')) as Heq.
            2: by rewrite Heq.
          assert (StronglySorted R rest') as SSrest'.
          { assert (StronglySorted R (v0 :: rest')).
            - rewrite -Hrest'2. by apply StronglySorted_merge_sort.
            - by inversion H0. }
          apply merge_sort_Permutation_eq,
            Permutation_cons_inv, merge_sort_equivalent in Hrest'2.
          rewrite (merge_sort_identity rest') in Hrest'2; auto.
        * apply bool_decide_eq_false_1, total_not in ER.
          exists vj, l'. split; auto. simpl.
          assert (is_min_of_join ((v0::l) :: L') vj) as Hminjoin.
          { eapply is_min_of_join_cons2; eauto. }
          destruct Hminjoin as [rest' Hrest'2]. rewrite Hrest'2.
          assert (rest' = merge_sort R (v0::l ++
            join_lists (<[find_min L':=l']> L'))
          ) as Heq.
            2: by rewrite Heq.
          assert (StronglySorted R rest') as SSrest'.
          { assert (StronglySorted R (vj :: rest')).
            - rewrite -Hrest'2. by apply StronglySorted_merge_sort.
            - by inversion H0. }
          apply (StronglySorted_unique R); auto.
            1: by apply StronglySorted_merge_sort.
          eapply Permutation_cons_inv. instantiate (1:=vj).
          rewrite -Hrest'2.
          etrans. 1: apply merge_sort_Permutation.
          etrans. 2: apply Permutation_cons. 2: eauto.
          2: symmetry; apply merge_sort_Permutation.
          etrans. 1: {
            symmetry; apply (join_lists_pop_Permutation
              _ (S (find_min L'))
            ). by rewrite lookup_cons.
          }
          auto.
    Qed.

  Lemma sorted_lists_pop L n i v l :
    is_sorted_lists L → S n = num_elements L →
    L !! i = Some (v :: l) →
    is_sorted_lists (<[i:=l]> L) ∧ n = num_elements (<[i:=l]> L).
  Proof.
    revert L n v l. induction i; intros L n v l SS Hn Hi.
    - destruct L; inversion Hi. subst; simpl. clear Hi.
      inversion SS; subst. inversion H2; subst.
      split; auto. by constructor.
    - destruct L; inversion Hi. clear Hi.
      inversion SS; subst.
      remember (num_elements L) as m.
      destruct m.
      { exfalso. clear IHi. revert i H1. induction L; intros.
        - inversion H1.
        - simpl in Heqm. destruct a; simpl in Heqm; try lia.
          inversion H4; subst. destruct i; inversion H1.
          eapply IHL; eauto. by constructor.
      }
      eapply IHi in H4 as [SSm Hm]; eauto.
      split.
      + constructor; auto.
      + simpl in *. rewrite -Heqm in Hn. rewrite -Hm. lia.
  Qed.

  Lemma phase2_correct {V} L (f : A → V → V) (start : V) :
    is_sorted_lists L →
    phase2 L f start = foldr f start (merge_sort R (join_lists L)).
  Proof.
    intros Hsort. unfold phase2.
    remember (num_elements L) as n. revert L start Hsort Heqn.
    induction n; intros.
    { by rewrite join_lists_empty. }
    simpl. destruct (phase2_first_correct L) as [v [l Hv]].
    all: try lia; auto.
    destruct Hv as [Hmin Hv1].
    eapply sorted_lists_pop in Hsort; eauto. destruct Hsort.
    rewrite Hmin Hv1. simpl. rewrite IHn; auto.
  Qed.
End Phase2.

Section combine.
  Definition TPMMS {V} mem M (f : A → V → V) (start : V) :=
    phase2 (phase1 mem M) f start.

  Lemma TPMMS_correct {V} mem M (f : A → V → V) (start : V) :
    0 < M →
    TPMMS mem M f start = foldr f start (merge_sort R mem).
  Proof.
    intros HM. unfold TPMMS.
    apply (phase1_correct mem) in HM as [ISS HP].
    apply (phase2_correct _ f start) in ISS.
    rewrite ISS.
    replace (merge_sort R (join_lists (phase1 mem M)))
      with (merge_sort R mem); auto.
    apply merge_sort_equivalent. by symmetry.
  Qed.
End combine.


Require Export ZF.

(** Stable functions *)

Definition stable_class (P:set->Prop) (F:set->set) :=
  forall X, (forall x, x∈X -> P x) -> inter (replf X F) ⊆ F (inter X).

Lemma cst_stable_class A K : stable_class K (fun _ => A).
red; red; intros.
apply inter_elim with (1:=H0) (y:=A).
destruct inter_wit with (2:=H0).
 red; red; reflexivity.
rewrite replf_ax; auto.
exists x; auto with *.
Qed.

Lemma id_stable_class K : stable_class K (fun x => x).
red; red; intros.
destruct inter_wit with (2:=H0).
 red; red; auto.
apply inter_intro; eauto.
intros.
apply inter_elim with (1:=H0).
rewrite replf_ax.
2:red; red; auto.
exists y;auto with *.
Qed.

Lemma compose_stable_class K1 K2 F G :
  Proper (eq_set ==> iff) K1 ->
  Proper (incl_set ==> incl_set) F ->
  morph1 G ->
  stable_class K1 F ->
  stable_class K2 G ->
  (forall x, K2 x -> K1 (G x)) ->
  stable_class K2 (fun o => F (G o)).
intros K1m Fm Gm Fs Gs Gty.
red; intros.
transitivity (F (inter (replf X G))).
 red; intros.
 apply Fs.
  intros.
  rewrite replf_ax in H1; auto with *.
  destruct H1.
  rewrite H2; auto.

  rewrite compose_replf; trivial.
   red; red; intros; apply Gm; trivial.

   apply Fmono_morph in Fm.
   red; red; intros; apply Fm; trivial.

 apply Fm.
 apply Gs; trivial.
Qed.

Lemma power_stable K : stable_class K power.
red; red; intros.
apply power_intro; intros.
destruct inter_non_empty with (1:=H0).
rewrite replf_ax in H2.
2:red;red;intros;apply power_morph; trivial.
destruct H2.
apply inter_intro; eauto.
clear H3 H4 H2 x0 x.
intros.
assert (z ∈ power y).
 apply inter_elim with (1:=H0).
 rewrite replf_ax.
 2:red;red;intros;apply power_morph; trivial.
 exists y; auto with *.
rewrite power_ax in H3; auto.
Qed.

Lemma union2_stable_disjoint K F G :
  morph1 F ->
  morph1 G ->
  stable_class K F ->
  stable_class K G ->
  (forall X Y z, K X -> K Y -> z ∈ F X -> z ∈ G Y -> False) ->
  stable_class K (fun X => F X ∪ G X).
intros Fm Gm Fs Gs disj.
intros X KX z zty.
destruct inter_wit with (2:=zty) as (w,winX).
 do 2 red; intros.
 rewrite H; reflexivity.
assert (forall x, x ∈ X -> z ∈ F x ∪ G x).
 intros.
 apply inter_elim with (1:=zty).
 rewrite replf_ax.
  exists x; auto with *.

  red; red; intros.
  rewrite H1; reflexivity.
clear zty.
assert (z ∈ F w ∪ G w) by auto.
apply union2_elim in H0; destruct H0.
 apply union2_intro1.
 apply Fs; auto.
 apply inter_intro.
  intros.
  rewrite replf_ax in H1.
  2:red;red;intros;apply Fm; trivial.
  destruct H1.
  rewrite H2; clear H2 y.
  assert (z ∈ F x ∪ G x) by auto.
  apply union2_elim in H2; destruct H2; trivial.
  elim disj with (3:=H0) (4:=H2); auto.

  exists (F w).
  rewrite replf_ax.
  2:red;red;intros;apply Fm;trivial.
  exists w; auto with *.

 apply union2_intro2.
 apply Gs; auto.
 apply inter_intro.
  intros.
  rewrite replf_ax in H1.
  2:red;red;intros;apply Gm; trivial.
  destruct H1.
  rewrite H2; clear H2 y.
  assert (z ∈ F x ∪ G x) by auto.
  apply union2_elim in H2; destruct H2; trivial.
  elim disj with (3:=H2) (4:=H0); auto.

  exists (G w).
  rewrite replf_ax.
  2:red;red;intros;apply Gm;trivial.
  exists w; auto with *.
Qed.

Definition stable := stable_class (fun _ => True).


Fixpoint eqb (n m : nat) : bool :=
    match n with
    | O => match m with
        | O => true
        | S m' => false
        end
    | S n' => match m with
        | O => false
        | S m' => eqb n' m'
        end
    end.

Notation "x =? y" := (eqb x y) (at level 70) : nat_scope.

Theorem zero_nbeq_plus_1 : forall n : nat,
    0 =? (n + 1) = false.
Proof.
    intros [| n'].
    - reflexivity.
    - reflexivity. Qed.